温故而知新,可以为师矣!
一、参考资料
SAVING AND LOADING MODELS
pytorch模型保存、加载与续训练
PyTorch保存和加载模型详解(一)
二、模型保存与加载
1. 搭建网络模型
import torch
import torch.nn as nn
import torch.nn.functional as F
# 模型定义
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
#模型初始化
model = Net()
2. 方式一(推荐)
2.1 保存模型
# 保存模型
torch.save(model.state_dict(), './model/model_state_dict.pth')
该方法后面的参数'./model/model_state_dict.pth'
为模型的保存路径,模型后缀名官方推荐使用.pth
和.pt
,当然,取别的后缀名也是可以的。
2.2 加载模型
使用load_state_dict()
加载模型时,先使用 load()
方法将保存的模型参数反序列化,load()
后的结果是一个字典,这时就可以通过 load_state_dict()
方法来加载。
model_test1 = Net() # 加载模型时应先实例化模型
# 加载模型
model_test1.load_state_dict(torch.load('./model/model_state_dict.pth'))
model_test1.eval() # 模型推理时设置
load_state_dict()
函数接收一个字典,所以不能直接将 './model/model_state_dict.pth'
传入,而是先使用 load()
函数将保存的模型参数反序列化。
序列化就是把内存中的数据保存到磁盘中,使用
torch.save()
方法保存模型就是序列化;而反序列化则是将硬盘中的数据加载到内存当中,显然我们加载模型的过程就是反序列化过程。
3. 方式二(推荐)
如果因为某种原因导致训练异常中止,采用checkpoint方式可以很方便的接着上次继续训练。正因为这样,非常推荐大家使用这种方式进行模型的保存与加载。
3.1 保存模型
# 保存checkpoint
torch.save({
'epoch':epoch,
'model_state_dict':model.state_dict(),
'optimizer_state_dict':optimizer.state_dict(),
'loss':loss
}, './model/model_checkpoint.tar' #这里的后缀名官方推荐使用.tar
)
3.2 加载模型
# 加载checkpoint
model = Net()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
checkpoint = torch.load('./model/model_checkpoint.tar') # 先反序列化模型
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']
3.3 代码示例
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torchvision
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#1、准备数据集
train_dataset = torchvision.datasets.CIFAR10("./data", train=True, transform=torchvision.transforms.ToTensor(), download= True)
test_dataset = torchvision.datasets.CIFAR10("./data", train=False, transform=torchvision.transforms.ToTensor(), download= True)
#2、加载数据集
train_dataset_loader = DataLoader(dataset=train_dataset, batch_size=100)
test_dataset_loader = DataLoader(dataset=test_dataset, batch_size=100)
#3、搭建神经网络
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.model1 = nn.Sequential(
nn.Conv2d(3, 32, 5, padding=2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, padding=2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, padding=2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(1024, 64),
nn.Linear(64, 10)
)
def forward(self, input):
input = self.model1(input)
return input
#4、创建网络模型
model = Net()
model.to(device)
#5、设置损失函数
loss_fun = nn.CrossEntropyLoss() #交叉熵函数
# 设置优化器
learning_rate = 1e-2
optimizer = torch.optim.SGD(model.parameters(), learning_rate) #SGD:梯度下降算法
#6、设置网络训练中的一些参数
Max_epoch = 10 #设置训练轮数
total_train_step = 0 #记录总训练次数
total_test_step = 0 #记录总测试次数
#7、开始进行训练
for epoch in range(Max_epoch):
print("---第{}轮训练开始---".format(epoch))
model.train() #开始训练,不是必须的,在网络中有BN,dropout时需要
#由于训练集数据较多,这里我没用训练集训练,而是采用测试集(test_dataset_loader)当训练集,但思想是一致的
for data in test_dataset_loader: # 遍历所有batch
imgs, targets = data
imgs, targets = imgs.to(device), targets.to(device)
#反向传播,更新参数
optimizer.zero_grad() # 重置每个batch的梯度
outputs = model(imgs) # 前向传播计算预测值
loss = loss_fun(outputs, targets) # 计算当前损失
loss.backward() # 反向传播计算梯度
optimizer.step() # 更新所有的参数
total_train_step += 1
if total_train_step % 50 == 0:
print("---第{}次训练结束, Loss:{})".format(total_train_step, loss.item()))
if epoch > 5:
print("---意外中断---")
break
if (epoch+1) % 2 == 0:
# 保存checkpoint
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss
}, './model/model_checkpoint_epoch_{}.tar'.format(epoch) # 这里的后缀名官方推荐使用.tar
)
两个epoch保存一次。当epoch=6时,设置一个break模拟程序意外中断,中断后可以来看到终端的输出信息,如下图所示:
从上图可以看到,在进行第6轮循环时,程序中断,此时最新的保存的模型是第五次训练结果,如下图所示:
同时注意到第5次训练结束的loss在2.0左右,如果我们下次接着训练,损失应该是在2.0附近。
此时,接着上次训练的结果继续训练,代码如下所示:
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torchvision
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#1、准备数据集
train_dataset = torchvision.datasets.CIFAR10("./data", train=True, transform=torchvision.transforms.ToTensor(), download= True)
test_dataset = torchvision.datasets.CIFAR10("./data", train=False, transform=torchvision.transforms.ToTensor(), download= True)
#2、加载数据集
train_dataset_loader = DataLoader(dataset=train_dataset, batch_size=100)
test_dataset_loader = DataLoader(dataset=test_dataset, batch_size=100)
#3、搭建神经网络
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.model1 = nn.Sequential(
nn.Conv2d(3, 32, 5, padding=2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, padding=2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, padding=2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(1024, 64),
nn.Linear(64, 10)
)
def forward(self, input):
input = self.model1(input)
return input
#4、创建网络模型
model = Net()
model.to(device)
#5、设置损失函数
loss_fun = nn.CrossEntropyLoss() #交叉熵损失
# 设置优化器
learning_rate = 1e-2
optimizer = torch.optim.SGD(model.parameters(), learning_rate) #SGD:梯度下降算法
#6、设置网络训练中的一些参数
Max_epoch = 10 #设置训练轮数
total_train_step = 0 #记录总训练次数
total_test_step = 0 #记录总测试次数
##########################################################################################
# 加载checkpoint
checkpoint = torch.load('./model/model_checkpoint_epoch_5.tar') # 先反序列化模型
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch']
loss = checkpoint['loss']
##########################################################################################
#7、开始进行训练
for epoch in range(start_epoch+1, Max_epoch):
print("---第{}轮训练开始---".format(epoch))
model.train() #开始训练,不是必须的,在网络中有BN,dropout时需要
for data in test_dataset_loader: # 遍历所有batch
imgs, targets = data
imgs, targets = imgs.to(device), targets.to(device)
#反向传播,更新参数
optimizer.zero_grad() # 重置每个batch的梯度
outputs = model(imgs) # 前向传播计算预测值
loss = loss_fun(outputs, targets) # 计算当前损失
loss.backward() # 反向传播计算梯度
optimizer.step() # 更新所有的参数
total_train_step += 1
if total_train_step % 50 == 0:
print("---第{}次训练结束, Loss:{})".format(total_train_step, loss.item()))
if (epoch+1) % 2 == 0:
# 保存checkpoint
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss
}, './model/model_checkpoint_epoch_{}.tar'.format(epoch) # 这里的后缀名官方推荐使用.tar
)
这里的代码相较之前的多了一个加载checkpoint
的过程,我将其截取出来,如下图所示:
通过加载checkpoint
,我们就保存了之前训练的参数,进而实现断点继续训练,我们直接来看执行此代码的结果,如下图所示:
从上图可以看出,训练是从第6轮开始的,并且初始的loss为1.99,和2.0接近。这就说明我们已经实现了中断后恢复训练的操作。
4. 方法三
4.1 保存模型
# 保存模型
torch.save(model, './model/model.pt') #这里我们保存模型的后缀名取.pt
4.2 加载模型
# 加载模型
model_test2 = torch.load('./model/model.pt')
model_test2.eval() # 模型推理时设置
这种方式是不推荐使用的,因为使用这种方式保存模型,在加载时会遇到各种各样的错误。为了加深大家理解,举一个例子。文件的结构如下图所示:
models.py
文件中存储的是模型的定义,其位于文件夹models下。save_model.py
文件中写的是保存模型的代码,代码如下所示:
from models.models import Net
from torch import optim
import torch
#模型初始化
model = Net()
# 初始化优化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# ## 保存加载方式2——save/load
# # 保存模型
# torch.save(models, './models/models.pt')
执行此文件后,会生成models.pt
文件,我们再执行load_mode.py
文件即可实现加载,load_mode.py
内容如下:
from models.models import Net
import torch
## 加载方式2
# 加载模型
model_test2 = Net()
model_test2 = torch.load('./models/models.pt')
model_test2.eval() # 模型推理时设置
print(model_test2)
此时我们可以正常加载。但如果我们将models文件夹修改为model,如下所示:
此时我们在使用如下代码加载模型的话就会出现错误:
from models.models import Net
import torch
## 加载方式2
# 加载模型
model_test2 = Net()
model_test2 = torch.load('./model/models.pt') #这里需要修改一下文件路径
model_test2.eval() # 模型推理时设置
print(model_test2)
出现这种错误的原因是,使用本方式进行模型保存的时候,会把模型结构定义文件路径记录下来,加载的时候就会根据路径解析它然后装载参数;当把模型定义文件路径修改以后,使用torch.load(path)
就会报错。
其实,使用本方式进行模型的保存和加载还会存在各种问题,感兴趣的可以看看这篇 博文 。总之,在我们今后的使用中,尽量不要用本方式来加载模型。