torch.jit.load#
- torch.jit.load(f, map_location=None, _extra_files=None, _restore_shapes=False)[source]#
載入之前使用
torch.jit.save儲存的ScriptModule或ScriptFunction。所有之前儲存的模組,無論其裝置如何,都會首先載入到 CPU,然後移動到它們儲存時的裝置。如果失敗(例如,因為執行時系統沒有某些裝置),則會引發異常。
- 引數
f – 一個檔案物件(必須實現 read, readline, tell, 和 seek),或一個包含檔名字串
map_location (string 或 torch.device) – torch.jit.save 中
map_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'])