Rate this Page

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 usingtorch::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 ofScriptModule as well.

Danger

All modules, no matter their device, are always loaded onto the CPUduring loading. This is different fromtorch.load()’s semanticsand may change in the future.

Parameters
  • m – AScriptModule to 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)