torch.jit.load

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

載入先前使用 torch.jit.save 儲存的 ScriptModule 或 ScriptFunction。

所有先前儲存的模組，無論其裝置為何，都會先載入到 CPU 上，然後再移動到儲存它們的裝置。 如果此操作失敗（例如，因為運行時系統沒有某些裝置），則會引發例外。

參數

• f – 一個類檔案物件（必須實現 read、readline、tell 和 seek），或一個包含檔案名稱的字串

• map_location (字串 或 torch.device) – torch.jit.save 中 map_location 的簡化版本，用於將儲存動態重新映射到另一組裝置。

• _extra_files (filename 到 content 的字典) – 在映射中給出的額外檔案名稱將被載入，並且它們的內容將儲存在提供的映射中。

• _restore_shapes (bool) – 是否在使用儲存的輸入載入時重新追蹤模組

回傳

一個 ScriptModule 物件。

範例： .. testcode

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'])