Table of Contents
Shortcuts

# Source code for torch.jit._serialization

    """Serialization

    This module contains functionality for serializing TorchScript modules, notably:
        * torch.jit.save
        * torch.jit.load

    This is not intended to be imported directly; please use the exposed
    functionalities in `torch.jit`.
    """
    import os
    import pathlib

    import torch
    from torch.jit._recursive import wrap_cpp_module
    from torch.serialization import validate_cuda_device

    def save(m, f, _extra_files=None):
        r"""
        Save an offline version of this module for use in a separate process. The
        saved module serializes all of the methods, submodules, parameters, and
        attributes of this module. It can be loaded into the C++ API using
        ``torch::jit::load(filename)`` or into the Python API with
        :func:`torch.jit.load <torch.jit.load>`.

        To be able to save a module, it must not make any calls to native Python
        functions.  This means that all submodules must be subclasses of
        :class:`ScriptModule` as well.

        .. DANGER::
            All modules, no matter their device, are always loaded onto the CPU
            during loading.  This is different from :func:`torch.load`'s semantics
            and may change in the future.

        Args:
            m: A :class:`ScriptModule` to save.
            f: A file-like object (has to implement write and flush) or a string
               containing a file name.
            _extra_files: Map from filename to contents which will be stored as part of `f`.

        .. note::
            torch.jit.save attempts to preserve the behavior of some operators
            across versions. For example, dividing two integer tensors in
            PyTorch 1.5 performed floor division, and if the module
            containing that code is saved in PyTorch 1.5 and loaded in PyTorch 1.6
            its division behavior will be preserved. The same module saved in
            PyTorch 1.6 will fail to load in PyTorch 1.5, however, since the
            behavior of division changed in 1.6, and 1.5 does not know how to
            replicate the 1.6 behavior.

        Example:

        .. testcode::

            import torch
            import io

            class MyModule(torch.nn.Module):
                def forward(self, x):
                    return x + 10

            m = torch.jit.script(MyModule())

            # Save to file
            torch.jit.save(m, 'scriptmodule.pt')
            # This line is equivalent to the previous
            m.save("scriptmodule.pt")

            # Save to io.BytesIO buffer
            buffer = io.BytesIO()
            torch.jit.save(m, buffer)

            # Save with extra files
            extra_files = {'foo.txt': b'bar'}
            torch.jit.save(m, 'scriptmodule.pt', _extra_files=extra_files)
        """
        if _extra_files is None:
            _extra_files = {}
        if isinstance(f, (str, pathlib.Path)):
            m.save(f, _extra_files=_extra_files)
        else:
            ret = m.save_to_buffer(_extra_files=_extra_files)
            f.write(ret)

    def load(f, map_location=None, _extra_files=None, _restore_shapes=False):
        r"""
        Load a :class:`ScriptModule` or :class:`ScriptFunction` previously
        saved with :func:`torch.jit.save <torch.jit.save>`

        All previously saved modules, no matter their device, are first loaded onto CPU,
        and then are moved to the devices they were saved from. If this fails (e.g.
        because the run time system doesn't have certain devices), an exception is
        raised.

        Args:
            f: a file-like object (has to implement read, readline, tell, and seek),
                or a string containing a file name
            map_location (string or torch.device): A simplified version of
                ``map_location`` in `torch.jit.save` used to dynamically remap
                storages to an alternative set of devices.
            _extra_files (dictionary of filename to content): The extra
                filenames given in the map would be loaded and their content
                would be stored in the provided map.
            _restore_shapes (bool): Whether or not to retrace the module on load using stored inputs

        Returns:
            A :class:`ScriptModule` object.

        Example:

        .. testcode::

            import torch
            import io

            torch.jit.load('scriptmodule.pt')

            # Load ScriptModule from io.BytesIO object
            with open('scriptmodule.pt', 'rb') as f:
                buffer = io.BytesIO(f.read())

            # Load all tensors to the original device
            torch.jit.load(buffer)

            # Load all tensors onto CPU, using a device
            buffer.seek(0)
            torch.jit.load(buffer, map_location=torch.device('cpu'))

            # Load all tensors onto CPU, using a string
            buffer.seek(0)
            torch.jit.load(buffer, map_location='cpu')

            # Load with extra files.
            extra_files = {'foo.txt': ''}  # values will be replaced with data
            torch.jit.load('scriptmodule.pt', _extra_files=extra_files)
            print(extra_files['foo.txt'])

        .. testoutput::
            :hide:

            ...

        .. testcleanup::

            import os
            os.remove("scriptmodule.pt")
        """

        if isinstance(f, str):
            if not os.path.exists(f):  # type: ignore[type-var]
                raise ValueError(f"The provided filename {f} does not exist")  # type: ignore[str-bytes-safe]
            if os.path.isdir(f):
                raise ValueError(f"The provided filename {f} is a directory")  # type: ignore[str-bytes-safe]

        map_location = validate_map_location(map_location)
        if _extra_files is None:
            _extra_files = {}

        cu = torch._C.CompilationUnit()
        if isinstance(f, (str, pathlib.Path)):
            cpp_module = torch._C.import_ir_module(cu, str(f), map_location, _extra_files, _restore_shapes)  # type: ignore[call-arg]
        else:
            cpp_module = torch._C.import_ir_module_from_buffer(
                cu, f.read(), map_location, _extra_files, _restore_shapes
            )  # type: ignore[call-arg]

        # TODO: Pretty sure this approach loses ConstSequential status and such
        return wrap_cpp_module(cpp_module)

    def validate_map_location(map_location=None):
        if isinstance(map_location, str):
            map_location = torch.device(map_location)
        elif not (map_location is None or isinstance(map_location, torch.device)):
            raise ValueError(
                "map_location should be either None, string or torch.device, "
                "but got type: " + str(type(map_location))
            )

        if str(map_location).startswith("cuda"):
            validate_cuda_device(map_location)

        return map_location

    def jit_module_from_flatbuffer(f):
        if isinstance(f, (str, pathlib.Path)):
            f = str(f)
            return wrap_cpp_module(torch._C._load_jit_module_from_file(f))
        else:
            return wrap_cpp_module(torch._C._load_jit_module_from_bytes(f.read()))

    def save_jit_module_to_flatbuffer(m, f, _extra_files=None):
        r"""
        Save an offline version of this module for use in a separate process. The
        saved module serializes all of the methods, submodules, parameters, and
        attributes of this module. It can be loaded into the C++ API using
        ``torch::jit::load_jit_module_from_file(filename)`` or into the Python API with
        :func:`torch.jit.jit_module_from_flatbuffer<torch.jit.jit_module_from_flatbuffer>`.

        To be able to save a module, it must not make any calls to native Python
        functions.  This means that all submodules must be subclasses of
        :class:`ScriptModule` as well.

        .. DANGER::
            All modules, no matter their device, are always loaded onto the CPU
            during loading.  This is different from :func:`torch.load`'s semantics
            and may change in the future.

        Args:
            m: A :class:`ScriptModule` to save.
            f: A string for file path

        Example:

        .. testcode::

            import torch
            import io

            class MyModule(torch.nn.Module):
                def forward(self, x):
                    return x + 10

            m = torch.jit.script(MyModule())

            # Save to file
            torch.jit.save_jit_module_to_flatbuffer(m, 'scriptmodule.ff')
        """

        extra_files = _extra_files
        if extra_files is None:
            extra_files = {}

        if isinstance(f, (str, pathlib.Path)):
            f = str(f)
            torch._C._save_jit_module(m._c, f, extra_files)
        else:
            s = torch._C._save_jit_module_to_bytes(m._c, extra_files)
            f.write(s)

    def get_flatbuffer_module_info(path_or_file):
        r"""Get some information regarding a model file in flatbuffer format.

        Args:
            path_or_file: Either str, Path or file like object (BytesIO OK).
                If it's str or Path, we will read the file referenced by that
                path as Bytes.

        Returns:
            A dict with metadata on what that file contains, currently looks like
            this:
            {
                'bytecode_version': 4,  # int
                'operator_version': 4,  # int
                'function_names': {
                    '__torch__.___torch_mangle_0.Foo.forward'}, # set
                'type_names': set(),  # set
                'opname_to_num_args': {'aten::linear': 3} # Dict[str, int]
            }
        """
        if isinstance(path_or_file, (str, pathlib.Path)):
            with open(path_or_file, "rb") as f:
                all_bytes = f.read()
        else:
            all_bytes = path_or_file.read()
        return torch._C._get_module_info_from_flatbuffer(all_bytes)

## Docs

Access comprehensive developer documentation for PyTorch

View Docs

## Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

## Resources

Find development resources and get your questions answered

View Resources

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.