Pytorch 模型保存与加载

模型保存

pytorch模型保存一般使用torch.save(),保存方法一般分为两种,分别为保存模型状态字典和模型本身两种

保存状态字典

若要保存状态字典,需要将save的两个参数设置为model.state_dict()和模型保存的路径和名称

import numpy as np
import torch
from torch import set_grad_enabled

device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')

class SimpleNet(torch.nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = torch.nn.Linear(10, 1)
    def forward(self, x):
        return self.fc1(x)

my_model = SimpleNet()
torch.save(my_model.state_dict(), 'my_model.pt')

保存模型

若要保存模型本身,需要将save的两个参数设置为model和模型保存的路径和名称

import numpy as np
import torch
from torch import set_grad_enabled

device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')

class SimpleNet(torch.nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = torch.nn.Linear(10, 1)
    def forward(self, x):
        return self.fc1(x)

my_model = SimpleNet()
torch.save(my_model, 'my_model1.pt')

关于保存模型状态字典和模型本身的区别将在模型加载处分析

模型加载

模型加载会使用函数torch.load()

这里就会体现出保存模型状态字典和模型本身的区别,保存模型本身实际会占用更多的字节数,也就是空间,原因是若只保存状态字典,则只会保存模型每层的权重和偏置,若是保存模型,则会在保存权重和偏置的基础上额外保存模型的架构,这样的优势在于,加载模型状态字典的时候会需要先实例化相同的模型架构,然后使用load_state_dict(torch.load(...))去获取模型参数,而加载整个模型则只需要torch.load(...)就可以

加载状态字典

import numpy as np
import torch
from torch import set_grad_enabled

device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')

class SimpleNet(torch.nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = torch.nn.Linear(10, 1)
    def forward(self, x):
        return self.fc1(x)

my_model = SimpleNet()
my_model.load_state_dict(torch.load('my_model.pt'))
print(my_model)
# SimpleNet(
#   (fc1): Linear(in_features=10, out_features=1, bias=True)
# )

加载完整模型

import numpy as np
import torch
from torch import set_grad_enabled

device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')

class SimpleNet(torch.nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = torch.nn.Linear(10, 1)
    def forward(self, x):
        return self.fc1(x)

my_model = torch.load('my_model1.pt')
print(my_model)
# SimpleNet(
#   (fc1): Linear(in_features=10, out_features=1, bias=True)
# )

注意,这两种加载方法分别对应了加载状态字典和加载模型,不可混用,如直接使用torch.load()加载模型字典,则只会读取到参数;若使用加载状态字典的方式加载完整模型则会报错