Python PyTorch load用法及代碼示例

本文簡要介紹python語言中 torch.jit.load 的用法。

用法:

torch.jit.load(f, map_location=None, _extra_files=None)

參數：

• f-file-like 對象(必須實現 read、readline、tell 和 seek)，或包含文件名的字符串

• map_location(string或者torch.device) -torch.jit.save 中map_location 的簡化版本，用於將存儲動態重新映射到一組備用設備。

• _extra_files(文件名到內容的字典) -Map中給出的額外文件名將被加載，它們的內容將存儲在提供的Map中。

返回：

ScriptModule 對象。

加載之前使用 torch.jit.save 保存的 ScriptModule 或 ScriptFunction

所有以前保存的模塊，無論它們的設備如何，都首先加載到 CPU 上，然後移動到保存它們的設備。如果失敗(例如，因為運行時係統沒有某些設備)，則會引發異常。

例子：

import torch
import io

torch.jit.load('scriptmodule.pt')

# Load ScriptModule from io.BytesIO object
with open('scriptmodule.pt', 'rb') as f:
    buffer = io.BytesIO(f.read())

# Load all tensors to the original device
torch.jit.load(buffer)

# Load all tensors onto CPU, using a device
buffer.seek(0)
torch.jit.load(buffer, map_location=torch.device('cpu'))

# Load all tensors onto CPU, using a string
buffer.seek(0)
torch.jit.load(buffer, map_location='cpu')

# Load with extra files.
extra_files = {'foo.txt': ''}  # values will be replaced with data
torch.jit.load('scriptmodule.pt', _extra_files=extra_files)
print(extra_files['foo.txt'])