pytorch保存模型的方式有两种
①将整个网络都都保存下来
保存整个神经网络的的结构信息和模型参数信息,save的对象是网络net
后缀一般命名为.pkl
②仅保存和加载模型参数(推荐使用这样的方法)
当需要为预测保存一个模型的时候,只需要保存训练模型的可学习参数即可。只保存神经网络的训练模型参数,save的对象是net.state_dict()
后缀一般命名为 .pt 或者 .pth
但是记住在进行预测之前,必须调用 model.eval() 方法来将 dropout 和 batch normalization 层设置为验证模型。否则,只会生成前后不一致的预测结果。
加载方式
①加载模型时通过torch.load('.pth')直接初始化新的神经网络对象
②需要首先导入对应的网络,再通过net.load_state_dict(torch.load('.pth'))完成模型参数的加载
net = Net() # 保存和加载整个模型 torch.save(net, 'model.pkl') model = torch.load('model.pkl')
# 仅保存和加载模型参数(推荐使用) torch.save(net.state_dict(), 'params.pkl') net.load_state_dict(torch.load('params.pkl'))
(后缀是pth还是pkl还是ckpt.mdl都无所谓的,你命名是啥就是啥)
记住这两种方式都是需要为预测保存模型的时候,即你已经把模型train完了,需要保存下来以备预测使用
但是如果在train的过程中保存,即你训练还没完不得已先结束或者断电断网了之类的,那就需要加载和保存一个通用的检查点(Checkpoint),也就是断点续传
直接读取,而不加载到模型中
import torch checkpoint = torch.load('output/trained_model.pth') print(checkpoint)
然后如果删掉几个key值元素再保存的话
del checkpoint['conv'] torch.save(checkpoint, 'work_dirs/self_annotated_IN_backbone_last_2_stages/epoch_12_without_char-bbox-head-fc-cls.pth')
甚至从其他模型中取一部分,写入到当前模型中再让程序加载也是ok的
import torch checkpoint = torch.load('work_dirs/self_annotated_IN_backbone_last_2_stages/epoch_12_without_char-bbox-head-fc-cls.pth') checkpoint_landmark = torch.load('work_dirs/self_annotated_IN_backbone_last_2_stages_landmark_3-layer_conv_10w+3w-data_divloss*50_6-epoch_feed-testdata-3k/epoch_6_only_with_landmark-conv.pth') checkpoint['state_dict'].update(checkpoint_landmark['state_dict']) torch.save(checkpoint, 'work_dirs/self_annotated_IN_backbone_last_2_stages/epoch_12_without_char-bbox-head-fc-cls_with_landmark.pth')
注意加载权重是要key名称匹配的
即使是相同结构,名称不同也是没法load的
删除部分模型
其实删除模型层就是替换模型部分层,我们只需要写一个空层替换局可以了,
如下:#删除最后一个层,或者说删除任意一个层 # 写一个空的层 classifier = nn.Sequential() #然后替换一下就可以了 CNN_shuffmodel.fc = classifier
删除层就是这么简单
模型的训练过程中加载和保存Checkpoint(断点续传)
当保存一个通用的检查点(checkpoint)时,无论是用于继续训练还是预测,都需要保存更多的信息,不仅仅是 state_dict ,比如说优化器的 state_dict 也是非常重要的,它包含了用于模型训练时需要更新的参数和缓存信息,还可以保存的信息包括 epoch,即中断训练的批次,最后一次的训练 loss,额外的 torch.nn.Embedding 层等等。
上述保存代码就是介绍了如何保存这么多种信息,通过用一个字典来进行组织,然后继续调用 torch.save 方法
保存
checkpoint = { "net": model.state_dict(), 'optimizer':optimizer.state_dict(), "epoch": epoch } torch.save(checkpoint, 'checkpoints/ckpt_%s.pth' %(str(epoch)))
加载
checkpoint = torch.load('checkpoints/ckpt_100.pth') # 加载断点 model.load_state_dict(checkpoint['net']) # 加载模型可学习参数 optimizer.load_state_dict(checkpoint['optimizer']) # 加载优化器参数 start_epoch = checkpoint['epoch'] # 设置开始的epoch #加载之后再训练就直接从start_epoch+1开始了
更详细的见
https://blog.csdn.net/hxxjxw/article/details/129257330
加载完后,根据后续步骤,调用 model.eval() 用于预测,model.train() 用于恢复训练。
在GPU上保存模型,在 CPU 上加载模型
保存模型的示例代码:
torch.save(model.state_dict(), PATH)
加载模型的示例代码:
device = torch.device('cpu') model = TheModelClass(*args, **kwargs) model.load_state_dict(torch.load(PATH, map_location=device))
在 CPU 上加载在 GPU 上训练的模型,必须在调用
torch.load()
的时候,设置参数map_location
,指定采用的设备是torch.device('cpu')
,这个做法会将张量都重新映射到 CPU 上。在GPU上保存模型,在 GPU 上加载模型
保存模型的示例代码:
torch.save(model.state_dict(), PATH)
加载模型的示例代码:
device = torch.device('cuda') model = TheModelClass(*args, **kwargs) model.load_state_dict(torch.load(PATH) model.to(device) # Make sure to call input = input.to(device) on any input tensors that you feed to the model
在 GPU 上训练和加载模型,调用
torch.load()
加载模型后,还需要采用model.to(torch.device('cuda'))
,将模型调用到 GPU 上,并且后续输入的张量都需要确保是在 GPU 上使用的,即也需要采用my_tensor.to(device)
。在CPU上保存,在GPU上加载模型
保存模型的示例代码:
torch.save(model.state_dict(), PATH)
加载模型的示例代码:
device = torch.device("cuda") model = TheModelClass(*args, **kwargs) model.load_state_dict(torch.load(PATH, map_location="cuda:0")) # Choose whatever GPU device number you want model.to(device) # Make sure to call input = input.to(device) on any input tensors that you feed to the model
这次是 CPU 上训练模型,但在 GPU 上加载模型使用,那么就需要通过参数
map_location
指定设备。然后继续记得调用model.to(torch.device('cuda'))
。
序列化 & 反序列化
PyTorch中模型,参数的保存
序列化——从内存到硬盘
反序列化——加载,从硬盘到内存
实操
我们使用vgg19模型 https://blog.csdn.net/hxxjxw/article/details/108923779
vgg19.py
import os import torch from torch.utils.data import DataLoader from torchvision import datasets from torchvision import transforms from torch import nn,optim #from lenet5 import Lenet5 from vgg import VGG19 from vgg import VGG34 def main(): batchsz = 32 cifar_train = datasets.CIFAR10('dataset/', train=True, transform=transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]) ]), download=True) cifar_train = DataLoader(cifar_train, batch_size=batchsz, shuffle=True) cifar_test = datasets.CIFAR10('dataset/', train=False, transform=transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]) ]), download=True) cifar_test = DataLoader(cifar_test, batch_size=batchsz, shuffle=True) x, label = iter(cifar_train).next() print('x:', x.shape, 'label:', label.shape) device = torch.device('cuda') model = VGG19().to(device) criton = nn.CrossEntropyLoss().to(device) #包含了softmax optimizer = optim.Adam(model.parameters(),lr=1e-3) print(model) if os.path.exists('model.pkl'): model.load_state_dict(torch.load('model.pkl')) print('model loaded from model.pkl') for epoch in range(20): model.train() for batchidx, (x,label) in enumerate(cifar_train): #x:[b,3,32,32] #label: [b] x,label = x.to(device), label.to(device) logits=model(x) #logits:[b,10] #label:b[b] #loss: tensor scalar 长度为0的标量 loss = criton(logits,label) #反向传播 optimizer.zero_grad() #梯度清零 loss.backward() #得到新的梯度 optimizer.step() #走一步,把梯度更新到weight里面去了 print('epoch',epoch, ' loss: ',loss.item()) model.eval() #test total_correct = 0 total_num = 0 for x, label in cifar_test: # x:[b,3,32,32] # label: [b] x,label = x.to(device) ,label.to(device) #logits:[b,10] logits = model(x) pred = logits.argmax(dim=1) total_correct += torch.eq(pred,label).float().sum().item() total_num += x.size(0) #即batch_size acc = total_correct / total_num print('epoch',epoch, ' acc: ',acc) torch.save(model.state_dict(),'model.pkl') print('model in epoch',epoch,' saved in model.pkl') if __name__ == '__main__': main()
模型训练结果
现在我们再运行一次,看一下直接加载model后的训练结果
直接就在原来acc 67%的模型上继续训练,上来就70%多了
如果只是想验证一下这个模型,不用训练
import os import torch from torch.utils.data import DataLoader from torchvision import datasets from torchvision import transforms from torch import nn,optim #from lenet5 import Lenet5 from vgg import VGG19 from vgg import VGG34 def main(): batchsz = 32 cifar_train = datasets.CIFAR10('dataset/', train=True, transform=transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]) ]), download=True) cifar_train = DataLoader(cifar_train, batch_size=batchsz, shuffle=True) cifar_test = datasets.CIFAR10('dataset/', train=False, transform=transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]) ]), download=True) cifar_test = DataLoader(cifar_test, batch_size=batchsz, shuffle=True) x, label = iter(cifar_train).next() print('x:', x.shape, 'label:', label.shape) device = torch.device('cuda') model = VGG19().to(device) criton = nn.CrossEntropyLoss().to(device) #包含了softmax optimizer = optim.Adam(model.parameters(),lr=1e-3) print(model) if os.path.exists('model.pkl'): model.load_state_dict(torch.load('model.pkl')) print('model loaded from model.pkl') model.eval() #test total_correct = 0 total_num = 0 for x, label in cifar_test: x,label = x.to(device) ,label.to(device) logits = model(x) pred = logits.argmax(dim=1) total_correct += torch.eq(pred,label).float().sum().item() total_num += x.size(0) #即batch_size acc = total_correct / total_num print('acc: ',acc) if __name__ == '__main__': main()
参考: