PyTorch入门十 || 模型保存

PyTorch入门十 || 模型保存

方式一：

保存训练好的参数，而不是整个模型，再次使用时需要搭建一个一模一样的神经网络

#模型
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(1,10,(5,5))
        self.conv2 = torch.nn.Conv2d(10,20,(5,5))
        self.maxpooling = torch.nn.MaxPool2d(2)
        self.fc = torch.nn.Linear(320,10)

    def forward(self,x):
        # 将数据从 (n,1,28,28) 转为 (n,784)

        #统计 minibatch 的大小
        batch_size = x.size(0)

        x = F.relu(self.maxpooling(self.conv1(x)))
        x = F.relu(self.maxpooling(self.conv2(x)))

        #将批量输入的图片转为 张数 * N
        # 注意 批量数据矩阵 一行表示一个数据
        x = x.view(batch_size,-1)
        x = self.fc(x)
        return x
#实例化
net = Net()

#训练好后保存参数
# net是上面的实例模型，‘CNN01’是文件名
torch.save(net.state_dict(),'CNN01')

# 或者 下面这样，其中 ‘.../para’是路径，一定要先存在路径，它不会自己新建文件夹，否则报错 ‘CNN01’ 是文件名
torch.save(net.state_dict(),'../para/CNN01')

使用时

#需要搭建和上述一样的模型，然后实例化

#实例化
net = Net()

#调用
cnn_dict = torch.load('../para/CNN01')
net.load_state_dict(cnn_dict)
test()

#其中 test()函数
def test():
    correct = 0
    total = 0
    with torch.no_grad():
        for data in test_loader:
            x,y = data
            x,y = x.to(device),y.to(device)
            y_pred = model(x)
            _,predicted = torch.max(y_pred.data,dim=1)
            total += y.size(0)
            correct += (predicted==y).sum().item()
    print('accuracy on test set:%d%% [%d/%d]'%(100*correct/total,correct,total))
    accuracy_list.append(100*correct/total)

调用结果：

方式二

保存整个模型，再次使用时不需要搭建一样的模型

#保存整个模型,  .pt 后缀可要可不要，建议要
torch.save(net, '../para/CNN_all.pt')       

使用：

#测试函数
def test(model):
    correct = 0
    total = 0
    with torch.no_grad():
        for data in test_loader:
            x,y = data
            x,y = x.to(device),y.to(device)
            y_pred = model(x)
            _,predicted = torch.max(y_pred.data,dim=1)
            total += y.size(0)
            correct += (predicted==y).sum().item()
    print('accuracy on test set:%d%% [%d/%d]'%(100*correct/total,correct,total))
    accuracy_list.append(100*correct/total)

    
#加载模型
new_net = torch.load('../para/CNN_all.pt')
new_net.to(device)		#模型加载进GPU
test(new_net)			#测试

结果：

总结：

保存时

torch.save(net.state_dict(),'../para/CNN01.pt')	#保存局部参数
# or
torch.save(net, '../para/CNN_all.pt')      		#保存整个模型

使用时

cnn_dict = torch.load('../para/CNN01')			#参数加载进字典
net.load_state_dict(cnn_dict)					#字典加载进模型，模型需要提前实例化

# or
new_net = torch.load('../para/CNN_all.pt')		#加载模型，模型不需提前定义

tips

#训练过程中保留最好的模型 ，其中min_loss 为 0.001
if loss.item() <= min_loss:
    torch.save(net.state_dict(),'../para/best_para.pt')