torch.jit.save¶
- torch.jit.save(m, f, _extra_files=None)[源][源]¶
保存此模块的离线版本,以便在单独的进程中使用。
保存的模块序列化此模块的所有方法、子模块、参数和属性。可以使用 C++ API 中的
torch::jit::load(filename)
或 Python API 中的torch.jit.load
将其加载。为了能够保存模块,它不能调用任何原生的 Python 函数。这意味着所有子模块也必须是
ScriptModule
的子类。警告
所有模块,无论其设备如何,在加载时总是加载到 CPU 上。这与
torch.load()
的语义不同,将来可能会改变。- 参数
m – 要保存的
ScriptModule
。f – 文件类对象(必须实现 write 和 flush 方法)或包含文件名的字符串。
_extra_files – 文件名到内容的映射,这些内容将作为 f 的一部分存储。
注意
torch.jit.save 尝试跨版本保留某些操作符的行为。例如,在 PyTorch 1.5 中,对两个整数张量进行除法运算会执行向下取整除法。如果包含该代码的模块在 PyTorch 1.5 中保存并在 PyTorch 1.6 中加载,其除法行为将得到保留。然而,在 PyTorch 1.6 中保存的同一模块将无法在 PyTorch 1.5 中加载,因为除法行为在 1.6 中发生了变化,而 1.5 不知道如何复现 1.6 的行为。
示例:.. 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)