wehub-resource-sync
36 lines
1.9 KiB
Python
36 lines
1.9 KiB
Python
"""
|
|
This module imports contents from CloudPickle in a way that is compatible with the
|
|
``pickle_module`` parameter of PyTorch's model persistence function: ``torch.save``
|
|
(see https://github.com/pytorch/pytorch/blob/692898fe379c9092f5e380797c32305145cd06e1/torch/
|
|
serialization.py#L192). It is included as a distinct module from :mod:`mlflow.pytorch` to avoid
|
|
polluting the namespace with wildcard imports.
|
|
|
|
Calling ``torch.save(..., pickle_module=mlflow.pytorch.pickle_module)`` will persist PyTorch model
|
|
definitions using CloudPickle, leveraging improved pickling functionality such as the ability
|
|
to capture class definitions in the "__main__" scope.
|
|
|
|
TODO: Remove this module or make it an alias of CloudPickle when CloudPickle and PyTorch have
|
|
compatible pickling APIs.
|
|
"""
|
|
|
|
# Import all contents of the CloudPickle module in an attempt to include all functions required
|
|
# by ``torch.save``.
|
|
|
|
# CloudPickle does not include `Unpickler` in its namespace, which is required by PyTorch for
|
|
# deserialization. Noting that CloudPickle's `load()` and `loads()` routines are aliases for
|
|
# `pickle.load()` and `pickle.loads()`, we therefore import Unpickler from the native
|
|
# Python pickle library.
|
|
from pickle import Unpickler # noqa: F401
|
|
|
|
from cloudpickle import * # noqa: F403
|
|
|
|
# PyTorch uses the ``Pickler`` class of the specified ``pickle_module``
|
|
# (https://github.com/pytorch/pytorch/blob/692898fe379c9092f5e380797c32305145cd06e1/torch/
|
|
# serialization.py#L290). Unfortunately, ``cloudpickle.Pickler`` is an alias for Python's native
|
|
# pickling class: ``pickle.Pickler``, instead of ``cloudpickle.CloudPickler``.
|
|
# https://github.com/cloudpipe/cloudpickle/pull/235 has been filed to correct the issue,
|
|
# but this import renaming is necessary until either the requested change has been incorporated
|
|
# into a CloudPickle release or the ``torch.save`` API has been updated to be compatible with
|
|
# the existing CloudPickle API.
|
|
from cloudpickle import CloudPickler as Pickler # noqa: F401
|