torch.jit.save#
- torch.jit.save(m,f,_extra_files=None)[source]#
Save an offline version of this module for use in a separate process.
The saved module serializes all of the methods, submodules, parameters, andattributes of this module. It can be loaded into the C++ API using
torch::jit::load(filename)or into the Python API withtorch.jit.load.To be able to save a module, it must not make any calls to native Pythonfunctions. This means that all submodules must be subclasses of
ScriptModuleas well.Danger
All modules, no matter their device, are always loaded onto the CPUduring loading. This is different from
torch.load()’s semanticsand may change in the future.- Parameters
m – A
ScriptModuleto save.f – A file-like object (has to implement write and flush) or a stringcontaining a file name.
_extra_files – Map from filename to contents which will be stored as part off.
Note
torch.jit.save attempts to preserve the behavior of some operatorsacross versions. For example, dividing two integer tensors inPyTorch 1.5 performed floor division, and if the modulecontaining that code is saved in PyTorch 1.5 and loaded in PyTorch 1.6its division behavior will be preserved. The same module saved inPyTorch 1.6 will fail to load in PyTorch 1.5, however, since thebehavior of division changed in 1.6, and 1.5 does not know how toreplicate the 1.6 behavior.
Example:.. testcode:
importtorchimportioclassMyModule(torch.nn.Module):defforward(self,x):returnx+10m=torch.jit.script(MyModule())# Save to filetorch.jit.save(m,'scriptmodule.pt')# This line is equivalent to the previousm.save("scriptmodule.pt")# Save to io.BytesIO bufferbuffer=io.BytesIO()torch.jit.save(m,buffer)# Save with extra filesextra_files={'foo.txt':b'bar'}torch.jit.save(m,'scriptmodule.pt',_extra_files=extra_files)