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)