torch.jit.load

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

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

所有之前保存的模块，无论它们的设备是什么，首先加载到 CPU 上，然后移动到它们保存的设备。如果这失败（例如，因为运行时系统没有某些设备），则会引发异常。

参数:

• f – 一个类似文件的对象（必须实现 read、readline、tell 和 seek），或者包含文件名的字符串

• map_location（字符串或 torch.device）- torch.jit.save 中 map_location 的简化版本，用于动态重映射存储到一组替代设备。

• _extra_files（文件名到内容的字典）- 在 map 中给出的额外文件名将被加载，其内容将存储在提供的 map 中。

• _restore_shapes（布尔值）- 是否在加载时使用存储的输入重新跟踪模块

返回值:

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