torch.jit.save

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

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

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

Внимание

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

Параметры

• m – ScriptModule для сохранения.
• f – Объект типа «файл» (должен реализовывать write и flush) или строка, содержащая имя файла.
• _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)