Guardado y carga de modelos existentes (versión PyTorch)

Tomamos la red VGG16 como ejemplo para ilustrar las operaciones de guardar y cargar de los modelos existentes.

Hay dos formas de guardar y cargar, y luego aprenderemos estas dos formas por separado. Nota: guardar y cargar no están en el mismo archivo py, configuramos la operación de guardar en el archivo save.py y la operación de carga en el archivo load.py.
Hay dos formas de guardar el modelo como se muestra en el código siguiente: la primera guarda tanto la estructura del modelo como los parámetros del modelo; la segunda guarda solo los parámetros del modelo y los guarda en forma de diccionario.

import torch
import torchvision

vgg16 = torchvision.models.vgg16(pretrained=False)
# 保存方式1，模型结构+模型参数
torch.save(vgg16, "vgg16_method1.pth") # 保存路径：vgg16_method1.pth

# 保存方式2，模型参数
torch.save(vgg16.state_dict(), "vgg16_method2.pth") # 保存路径：vgg16_method2.pth

Las dos formas de cargar el modelo se muestran en el siguiente código.

import torch

# 方式1 --》保存方式1，加载模型
model1 = torch.load("vgg16_method1.pth")
print(model1)

# 方式2 --》保存方式2，加载模型
model2 = torch.load("vgg16_method2.pth")
print(model2)

El resultado impreso es:

 El resultado del modelo 1 es:

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

Los resultados anteriores son la estructura del modelo de red VGG16 y los parámetros del modelo de red. 

El resultado del modelo 2 es: (debido a que hay demasiados resultados, solo se dan una parte de los resultados)

OrderedDict([('features.0.weight', tensor([[[[-0.1638, -0.0292,  0.0316],
          [-0.0149,  0.0681,  0.0458],
          [ 0.0633, -0.0374, -0.0047]],

         [[-0.0123, -0.0461,  0.0343],
          [ 0.0207, -0.0128,  0.0107],
          [-0.0181,  0.0154,  0.0320]],

         [[-0.0759, -0.1384, -0.0318],
          [ 0.0244, -0.0424,  0.0332],
          [-0.0244,  0.0524,  0.1292]]],
……………………………………

Los resultados anteriores son los parámetros del modelo de red VGG16. 

Entonces, si queremos usar los parámetros del modelo guardados al guardar el método 2, ¿cómo usarlos? Consulte el código a continuación.
Primero construimos la estructura del modelo de red y luego cargamos los parámetros del modelo de red guardados en la estructura del modelo de red.

vgg16 = torchvision.models.vgg16(pretrained=False) # 网络模型结构
vgg16.load_state_dict(torch.load("vgg16_method2.pth")) # 加载保存的网络模型参数
print(vgg16)

El método de ahorro 1 tiene un pequeño inconveniente.

Ilustramos este problema construyendo nosotros mismos una red.
Construimos nuestra estructura de red en el archivo save.py y la guardamos.

# 陷阱1
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=5)
    def forward(self, x):
        x = self.conv1(x)
        return x


tudui = Tudui()
torch.save(tudui, "tudui_method1.pth")

A continuación, cargamos este modelo en el archivo load.py de acuerdo con el método de carga 1.

# 陷阱1
model = torch.load("tudui_method1.pth")
print(model)

El resultado impreso es:

AttributeError: Can't get attribute 'Tudui' on <module '__main__' from 'D:/graduate0/pytorch_practice/model_load.py'>

Descubrimos que se informó un error y el motivo del error fue que no se pudo obtener la propiedad de Tudui.

Agregamos la estructura de red en el archivo load.py. Nota: No es necesario crear un modelo de red en este momento, es decir, no es necesario ejecutar el código tudui=Tudui().

from torch import nn
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=5)
    def forward(self, x):
        x = self.conv1(x)
        return x


# tudui = Tudui()
model = torch.load("tudui_method1.pth")
print(model)

El resultado impreso es: 

Tudui(
  (conv1): Conv2d(3, 64, kernel_size=(5, 5), stride=(1, 1))
)