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)