pytorch模型保存格式

pytorch模型保存格式

对着手写数字识别实例讲讲pytorch模型保存的格式。首先讲讲保存模型或权重参数的后缀，pytorch保存数据的格式为.t7文件或者.pth文件，或者.pkl格式，t7文件是沿用torch7中读取模型权重的方式。而pth文件是python中存储文件的常用格式。而在keras中则是使用.h5文件
1、pytorch保存和加载模型以及权重参数（强烈推荐使用这种）
1.1 首先新建model.py模块
将模型单独新建一个模块

from torch import nn

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 使用序列工具快速构建
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(2))
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2))
        self.fc = nn.Linear(7 * 7 * 32, 10) 
    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = out.view(out.size(0), -1)  # reshape
        out = self.fc(out)
        return out

1.2 新建一个train.py模块
保存模型和权重参数的格式为：

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

具体的实例如下所示：

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import torchvision
from torchvision import datasets, transforms
from torch.autograd import Variable
import os
batch_size = 64
learning_rate = 0.001
# 将数据处理成Variable, 如果有GPU, 可以转成cuda形式
def get_variable(x):
    x = Variable(x)
    return x.cuda() if torch.cuda.is_available() else x
train_dataset = datasets.MNIST(
    root='./mnist/',
    train=True,
    transform=transforms.ToTensor(),
    download=True)
transforms