Pytorch 模型保存/加载

Pytorch是常用的深度学习框架之一，在训练神经网络时，我们经常需要保存模型或者加载模型。模型保存有两种方式：保存模型和保存模型的参数。保存模型即把整个模型的结构和参数保存下来，后者只保存参数（这种方式在加载模型时，必须知道模型的结构，才可以加载）

保存模型

• 保存

torch.save(model, path)

model是模型，path是模型保存的路径，例如：'./*.pkl' 或者 './*.tar' './*.pth', 保存参数的文件一定要有后缀扩展名。

• 加载

model = torch.load(path)

保存模型

• 保存

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

或者保存更多的信息

state = {'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch}
torch.save(state, path)

• 加载

针对第一种情况

model.load_state_dict(torch.load(path))

针对第二种情况

checkpoint = torch.load(path)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
epoch = checkpoint(['epoch'])

上述模型model必须实例化，即model=Net()

下面给出一个具体的例子程序，该程序只保存最新的参数：

#-*- coding:utf-8 -*-

import torch as torch
import torchvision as tv
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision.transforms as transforms
from torchvision.transforms import ToPILImage
import torch.backends.cudnn as cudnn
import datetime
import argparse

# 参数
log_dir = '/home/pxt/pytorch/logs/cifar_model.pth'  # 模型保存路径


# 构造模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
        ......
    
    
    def forward(self, x):
        ......
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)


# 模型训练
def train(model, train_loader, epoch):
    model.train()
    train_loss = 0
    .......


# 模型测试
def test(model, test_loader):
    model.eval()
    ......

if __name__ == '__main__':
    # 如果test_flag=True,则加载已保存的模型
    if test_flag:
        # 加载保存的模型直接进行测试机验证，不进行此模块以后的步骤
        checkpoint = torch.load(log_dir)
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        epochs = checkpoint['epoch']
        test(model, test_load)
        return

    for epoch in range(0, epochs):
        train(model, train_load, epoch)
        test(model, test_load)
        # 保存模型
        state = {'model':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}
        torch.save(state, log_dir)

• 程序中断后继续训练

在训练模型的时候可能会因为一些问题导致程序中断，或者常常需要观察训练情况的变化来更改学习率等参数，这时候就需要加载中断前保存的模型

def main():

    # 如果test_flag=True,则加载已保存的模型
    if test_flag:
        # 加载保存的模型直接进行测试机验证，不进行此模块以后的步骤
        checkpoint = torch.load(log_dir)
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        start_epoch = checkpoint['epoch']
        test(model, test_load)
        return

    # 如果有保存的模型，则加载模型，并在其基础上继续训练
    if os.path.exists(log_dir):
        checkpoint = torch.load(log_dir)
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        start_epoch = checkpoint['epoch']
        print('加载 epoch {} 成功！'.format(start_epoch))
    else:
        start_epoch = 0
        print('无保存模型，将从头开始训练！')

    for epoch in range(start_epoch+1, epochs):
        train(model, train_load, epoch)
        test(model, test_load)
        # 保存模型
        state = {'model':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}
        torch.save(state, log_dir)