評價此頁

torch.jit.load#

torch.jit.load(f, map_location=None, _extra_files=None, _restore_shapes=False)[source]#

載入之前使用 torch.jit.save 儲存的 ScriptModuleScriptFunction

所有之前儲存的模組,無論其裝置如何,都會首先載入到 CPU,然後移動到它們儲存時的裝置。如果失敗(例如,因為執行時系統沒有某些裝置),則會引發異常。

引數
  • f – 一個檔案物件(必須實現 read, readline, tell, 和 seek),或一個包含檔名字串

  • map_location (stringtorch.device) – torch.jit.savemap_location 的一個簡化版本,用於動態地將儲存重新對映到一組備用裝置。

  • _extra_files (filename 到 content 的字典) – 在 map 中給出的額外檔名將被載入,並且其內容將儲存在提供的 map 中。

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

返回

一個 ScriptModule 物件。

警告

有可能構造惡意的 pickle 資料,這些資料會在 func:torch.jit.load 期間執行任意程式碼。切勿載入可能來自不受信任來源或可能被篡改過的資料。僅載入您信任的資料

示例: .. 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'])