Сохранение и загрузка модели в факел

программист

Ссылка на ссылку

Сохранение и загрузка моделей — официальное руководство по PyTorch, китайская версия (pytorch123.com)

state_dict

state_dict базовое содержимое

Доступ к параметрам модели torch.nn.Module в pytorch можно получить через функцию model.parameters(), которая возвращает генератор, а результатом просмотра model.parameters() через цикл for являются данные параметров каждого слоя .

for x in model.parameters():
    print(x)

#output:
Parameter containing:
tensor([[[[ 0.0908, -0.1033, -0.0175, -0.0567, -0.0543],
          [-0.0666,  0.0097,  0.0017, -0.1141,  0.1127],
          [ 0.0113, -0.1119,  0.0696, -0.0383, -0.0037],
          [-0.0199,  0.1131,  0.0044,  0.0083, -0.0113],
          [-0.0267, -0.0676,  0.0433, -0.0297,  0.0166]],
          ......

model.state_dict() возвращает формат словаря (ключ-значение), где ключ — это имя каждого слоя модели, а значение — данные параметра соответствующего слоя.

Форма, возвращаемая вызовом state_dict(), имеет форму словаря. Кроме модели в оптимизаторе есть еще и параметры, поэтому параметры получаются для сохранения через optimizer.state_dict(), а в словаре будут отображаться только слои с параметрами.

import torch.nn as nn
import torch.nn.functional as F
from torch import optim
import torch


class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

if __name__=="__main__":
        model = TheModelClass()
        optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
        
        print("model's state_dict: ")
        for param_tensor in model.state_dict():
            print(param_tensor, "\t", model.state_dict()[param_tensor].size())

        print(optimizer.state_dict())

#output:
model's state_dict: 
conv1.weight 	 torch.Size([6, 3, 5, 5])
conv1.bias 	 torch.Size([6])
conv2.weight 	 torch.Size([16, 6, 5, 5])
conv2.bias 	 torch.Size([16])
fc1.weight 	 torch.Size([120, 400])
fc1.bias 	 torch.Size([120])
fc2.weight 	 torch.Size([84, 120])
fc2.bias 	 torch.Size([84])
fc3.weight 	 torch.Size([10, 84])
fc3.bias 	 torch.Size([10])
{'state': {}, 'param_groups': [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]}]}

Обратите внимание на приведенный выше код:

  1. model.state_dict() возвращает словарную форму, ключ — это имя каждого слоя + вес или смещение, а значение — соответствующие данные параметра;
  2. optimizer.state_dict() возвращает словарную форму, первая — «state»: {}, вторая — «param_groups»:[], где [] — словарная форма. Поэтому нужно обратить внимание, когда нужно получить параметр, второй словарь параметра хранится в списке.

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

Существует три способа сохранения и загрузки моделей.

  • Первый — сохранить только параметры и сохранить их в виде словаря;
  • Второй — сохранить модель и параметры в целом;
  • Третий — использовать форму словаря для сохранения параметров модели, параметров оптимизатора и рабочей среды в формате checkpoint.tar.

Второй тип имеет меньше кода, но сохраняемый файл больше, и pickle не может сохранить модель, после использования или рефакторинга других проектов код может быть сломан различными способами. Поэтому изучайте только два других.

  • спасти:torch.save(state_dict, path)Сохраните сериализованный объект на диск. Он может сериализовать модель или словарь с помощью рассола и сохранить его. часто сохраняется как.pthтип файла.

  • Десериализовать:torch.load(path)Десериализовать файл в словарь, а затем использовать параметры в виде словаряload_state_dict(torch.load(path))импортируется в модель.

  • нагрузка:torch.nn.Module.load_state_dict: загрузить десериализованные данные в модель. существуетload_state_dict()один из параметровstrictпараметр, когдаstrict=Falseкогда это означает, что в параметре импортаkeyНе требует точного соответствия каждому слою в модели, может использоваться для отсутствующих определенных ключей.state_dictЕсть больше ключей для загрузки или из загруженной моделиstate_dict.

Сохраняйте только параметры модели

Если вы хотите исправить параметры слоя BN и отсева в тестовом примере, вам нужно использоватьmodel.eval().

  • спасти
torch.save(model.state_dict(), "dict.pth")
  • нагрузка
model1 = TheModelClass()
model1.load_state_dict(torch.load("dict.pth"))
model1.eval()

Сохраните модель, параметры оптимизатора и среду.

Во время обучения возможны перебои с питанием или другие непредвиденные ситуации в процессе обучения.В это время необходимо постоянно сохранять текущие результаты обучения.Результаты обучения включают в себя параметры модели, параметры оптимизатора и среду, в которой работает модель . . . Таким образом, вы можете сохранить его как словарь и обращаться к нему как к словарю после загрузки.

  • спасти
epoch = 60
loss = 0.12
torch.save({
    'epoch': 60,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss
}, "checkpoint1.tar")
  • нагрузка
model = TheModelClass()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
checkpoint = torch.load("checkpoint1.tar")

model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

оборудование для загрузки модели

Сохранить на GPU, загрузить на GPU

torch.save(model.state_dict(), PATH)

device = torch.device("cuda") 
model = TheModelClass(*args, **kwargs) 
model.load_state_dict(torch.load(PATH)) 
model.to(device) # 确保在你提供给模型的任何输入张量上调用input = input.to(device)

При обучении на графическом процессоре и сохранении модели на графическом процессоре просто используйтеmodel.to(torch.device('cuda')), который преобразует инициализированную модель в модель, оптимизированную для CUDA.

Кроме того, обязательно используйте на всех входных данных модели.to(torch.device('cuda'))Функция подготовки данных для модели.

Обратите внимание, что вызовmy_tensor.to(device)вернется на GPUmy_tensorкопия. Поэтому не забудьте переопределить тензоры вручную:my_tensor= my_tensor.to(torch.device('cuda')).

Категории