torch.jit.save

torch.jit.save(m, f, _extra_files=None) [source]

Сохраните автономную версию данного модуля для использования в отдельном процессе. Сохранённый модуль сериализует все методы, подмодули, параметры и атрибуты этого модуля. Его можно загрузить в API на C++ с помощью torch::jit::load(filename) или в Python API с помощью torch.jit.load.

Для сохранения модуля он не должен вызывать какие-либо родные функции Python. Это означает, что все подмодули должны быть подклассами ScriptModule.

Внимание

Все модули, независимо от устройства, всегда загружаются на ЦП во время загрузки. Это отличается от семантики torch.load() и может измениться в будущем.

Параметры:

• m – ScriptModule для сохранения.
• f – Объект, подобный файлу (должен реализовывать запись и сброс), или строка, содержащая имя файла.
• _extra_files – Отображение имени файла на содержимое, которое будет храниться как часть f.

Примечание

torch.jit.save пытается сохранить поведение некоторых операторов через версии. Например, деление двух целочисленных тензоров в PyTorch 1.5 выполняло целочисленное деление, и если модуль, содержащий этот код, сохранён в PyTorch 1.5 и загружен в PyTorch 1.6, его поведение при делении будет сохранено. Однако тот же модуль, сохранённый в PyTorch 1.6, не сможет загрузиться в PyTorch 1.5, так как поведение деления изменилось в 1.6, а 1.5 не знает, как воспроизвести поведение 1.6.

Пример:

import torch
import io

class MyModule(torch.nn.Module):
    def forward(self, x):
        return x + 10

m = torch.jit.script(MyModule())

# Save to file
torch.jit.save(m, 'scriptmodule.pt')
# This line is equivalent to the previous
m.save("scriptmodule.pt")

# Save to io.BytesIO buffer
buffer = io.BytesIO()
torch.jit.save(m, buffer)

# Save with extra files
extra_files = {'foo.txt': b'bar'}
torch.jit.save(m, 'scriptmodule.pt', _extra_files=extra_files)