Функции torch.save и torch.load в PyTorch

Pytorch torch 参考手册 Справочное руководство Pytorch torch

torch.saveиtorch.loadявляются функциями в PyTorch для сериализации (сохранения) и десериализации (загрузки) тензоров, моделей и других объектов Python.

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

Определение функции

torch.save(obj, f, pickle_module, pickle_protocol)
torch.load(f, map_location, pickle_module, weights_only)

Параметры torch.save:

obj: Объект для сохранения, который может быть тензором, моделью, словарём, списком и т. д.
f: Путь к файлу (строка или файловый объект).
pickle_module(Необязательно): Модуль, используемый для сериализации.
pickle_protocol(Необязательно): Версия протокола сериализации.

Параметры torch.load:

f: Путь к файлу (строка или файловый объект).
map_location(Необязательно): Указывает, как сопоставить хранилище с различными устройствами.
pickle_module(Необязательно): Модуль, используемый для десериализации.
weights_only(bool, optional): Загружать только веса, а не объекты Python.

Примеры использования

Пример 1: Сохранение и загрузка тензоров

Пример

import torch

# Создать несколько тензоров
x = torch.tensor([1, 2, 3, 4, 5])
y = torch.randn(3, 4)

# Сохранить тензоры в файл
torch.save({'x': x, 'y': y}, 'tensors.pth')

# Загрузить тензоры из файла
loaded = torch.load('tensors.pth')

print("Загруженные данные:", loaded)
print("x:", loaded['x'])
print("y:", loaded['y'])

Вывод:

加载的数据: {'x': tensor([1, 2, 3, 4, 5]), 'y': tensor([[-0.3042, -0.9077, -1.0826,  0.9333],
        [ 0.0551,  0.6728,  0.5942, -0.1522],
        [-0.3744,  0.9239, -0.2104, -0.5239]])}
x: tensor([1, 2, 3, 4, 5])
y: tensor([[-0.3042, -0.9077, -1.0826,  0.9333],
        [ 0.0551,  0.6728,  0.5942, -0.1522],
        [-0.3744,  0.9239, -0.2104, -0.5239]])

Пример 2: Сохранение и загрузка моделей

Пример

import torch
import torch.nn as nn

# Определить простую нейронную сеть
class SimpleNet(nn.Module):
def __init__(self):
super(SimpleNet, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.fc2 = nn.Linear(5, 2)

def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x

# Создать экземпляр модели
model = SimpleNet()

# Сохранить модель (сохранить всю модель)
torch.save(model, 'model.pth')

# Загрузить модель
loaded_model = torch.load('model.pth')

print("Модель сохранена и загружена")
print(loaded_model)

Пример 3: Сохранение только параметров модели (рекомендуемый метод)

Пример

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
def __init__(self):
super(SimpleNet, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.fc2 = nn.Linear(5, 2)

def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x

model = SimpleNet()

# Сохранить только параметры модели (рекомендуемый метод)
torch.save(model.state_dict(), 'model_weights.pth')

# Создать новую модель и загрузить параметры
new_model = SimpleNet()
new_model.load_state_dict(torch.load('model_weights.pth'))

print("Параметры модели сохранены и загружены")
print(new_model.state_dict().keys())

Вывод:

模型参数已保存并加载
odict_keys(['fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias'])


It is recommended to save only the model parameters (state_dict), rather than saving the entire model. This allows parameters to be reused between models with different architectures.

Example 4: Transferring Between CPU and GPU


Example

import torch



# Assume the model was saved on GPU

if torch.cuda.is_available():

    x = torch.randn(2, 3, device='cuda')

    torch.save(x, 'tensor_gpu.pth')



    # Load on CPU

    x_cpu = torch.load('tensor_gpu.pth', map_location='cpu')

    print("Loaded to CPU:", x_cpu.device)




Usingmap_locationthe parameter, you can load tensors onto different devices.


Notes


Save files usually use.pth или  .ptextension.
It is recommended to save only the model parameters (state_dict), rather than saving the entire model.
When loading, security and version compatibility should be considered.



Pytorch torch Reference Manual