Table of Contents
Shortcuts

torch.jit.load

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

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

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

Parameters

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

  • map_location (字符串 或 torch.device) – 一个简化的版本 map_location 在 torch.jit.save 中用于动态地将存储重新映射到一组替代设备。

  • _extra_files (字典 的 文件名到内容) – 在映射中给出的额外文件名将被加载，并且它们的内容将被存储在提供的映射中。

  • _restore_shapes (bool) – 是否在加载时使用存储的输入重新追踪模块

Returns

一个 ScriptModule 对象。

示例: .. testcode:

import torch
import io

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

# 从 io.BytesIO 对象加载 ScriptModule
with open('scriptmodule.pt', 'rb') as f:
    buffer = io.BytesIO(f.read())

# 将所有张量加载到原始设备
torch.jit.load(buffer)

# 将所有张量加载到 CPU，使用设备
buffer.seek(0)
torch.jit.load(buffer, map_location=torch.device('cpu'))

# 将所有张量加载到 CPU，使用字符串
buffer.seek(0)
torch.jit.load(buffer, map_location='cpu')

# 加载带有额外文件
extra_files = {'foo.txt': ''}  # 值将被数据替换
torch.jit.load('scriptmodule.pt', _extra_files=extra_files)
print(extra_files['foo.txt'])
优云智算