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')