File size: 9,687 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 |
# mypy: allow-untyped-defs
"""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 torch
from torch._jit_internal import _get_model_id
from torch._utils_internal import log_torchscript_usage
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)
"""
log_torchscript_usage("save", model_id=_get_model_id(m))
if _extra_files is None:
_extra_files = {}
if isinstance(f, (str, os.PathLike)):
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.
.. warning::
It is possible to construct malicious pickle data which will execute arbitrary code
during func:`torch.jit.load`. Never load data that could have come from an untrusted
source, or that could have been tampered with. **Only load data you trust**.
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, os.PathLike)):
if not os.path.exists(f):
raise ValueError(f"The provided filename {f} does not exist")
if os.path.isdir(f):
raise ValueError(f"The provided filename {f} is a directory")
map_location = validate_map_location(map_location)
if _extra_files is None:
_extra_files = {}
cu = torch._C.CompilationUnit()
if isinstance(f, (str, os.PathLike)):
cpp_module = torch._C.import_ir_module(
cu, os.fspath(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
ret = wrap_cpp_module(cpp_module)
log_torchscript_usage("load", model_id=_get_model_id(ret))
return ret
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, os.PathLike)):
f = os.fspath(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, os.PathLike)):
f = os.fspath(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, os.PathLike)):
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)
|