torch.jit.load

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

Загрузка ScriptModule или ScriptFunction, сохранённых ранее с помощью torch.jit.save

Все сохранённые ранее модули, независимо от их устройства, сначала загружаются в ЦП, а затем перемещаются на устройства, с которых они были сохранены. Если это не удаётся (например, потому что в среде выполнения отсутствуют определённые устройства), генерируется исключение.

Параметры

• f – объект типа «файл» (должен реализовывать методы read, readline, tell и seek) или строка, содержащая имя файла
• map_location (string или torch.device) – Упрощённая версия map_location в torch.jit.save, используемая для динамического переназначения хранилищ на другой набор устройств.
• _extra_files (словарь имени файла к содержимому) – Дополнительные имена файлов, указанные в map, будут загружены, а их содержимое будет сохранено в предоставленном map.
• _restore_shapes (bool) – Восстановить форму модуля при загрузке с помощью сохранённых входных данных

Возвращает

Объект ScriptModule.

Пример:

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