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开始了
更详细的见
加载完后,根据后续步骤,调用 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模型
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()