本人学习pytorch主要参考官方文档和 莫烦Python中的pytorch视频教程。
后文主要是对pytorch官网的文档的总结。
代码来自pytorch官网
-
torch.save：将对象序列化到硬盘上，该对象可以是 Models, tensors和 dictionaries 等。实际上是使用了python的pickle方法。 -
torch.load：将硬盘上序列化的对象加载设备中。实际是使用了pickle的解包方法。 -
torch.nn.Module.load_state_dict：通过反序列化state_dict加载模型参数字典。
模型参数
在pytorch中torch.nn.Module模型的参数存放在模型的parameters中(model.parameters()),而state_dict是参数tensor的字典。仅当该层有可学习的参数才会有属性state_dict。torch.optim也有state_dict属性，其中包含的是优化器的参数，即网络所使用的超参数。
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
class TheModelClass(nn.Module):
def __init__(self):
super(TheModelClass, 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 = TheModelClass()
# 初始化优化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# 打印模型参数
print("Model's state_dict:")
for param_tensor in model.state_dict():
print(param_tensor, "\t", model.state_dict()[param_tensor].size())
# 打印优化器参数
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
print(var_name, "\t", optimizer.state_dict()[var_name])
保存和加载模型
保存的文件结尾一般使用.pt或.pth。最后的model.eval()表示将drop和batch nromalization层设置为测试模式。可通过mode.train()转化为训练模式。
模型参数
#保存模型参数
torch.save(model.state_dict(), PATH)
#加载模型参数
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()
模型
#保存模型
torch.save(model, PATH)
#加载模型
model = torch.load(PATH)
model.eval()
保存多个模型
torch.save({
'modelA_state_dict': modelA.state_dict(),
'modelB_state_dict': modelB.state_dict(),
'optimizerA_state_dict': optimizerA.state_dict(),
'optimizerB_state_dict': optimizerB.state_dict(),
...
}, PATH)
加载多个模型
modelA = TheModelAClass(*args, **kwargs)
modelB = TheModelBClass(*args, **kwargs)
optimizerA = TheOptimizerAClass(*args, **kwargs)
optimizerB = TheOptimizerBClass(*args, **kwargs)
checkpoint = torch.load(PATH)
modelA.load_state_dict(checkpoint['modelA_state_dict'])
modelB.load_state_dict(checkpoint['modelB_state_dict'])
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])
modelA.eval()
modelB.eval()
# - or -
modelA.train()
modelB.train()
加载不同的模型
加载他人训练的模型，可能需要忽略部分层。则将load_state_dict方法的strict参数设置为False。则模型金价在modelB中名称对应的层的参数。
torch.save(modelA.state_dict(), PATH)
modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)
加载不同设备的模型
将由GPU保存的模型加载到CPU上。
将torch.load()函数中的map_location参数设置为torch.device('cpu')
device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))
将由GPU保存的模型加载到GPU上。确保对输入的tensors调用input = input.to(device)方法。
device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)
将由CPU保存的模型加载到GPU上。确保对输入的tensors调用input = input.to(device)方法。map_location是将模型加载到GPU上，model.to(torch.device('cuda'))是将模型参数加载为CUDA的tensor。最后保证使用.to(torch.device('cuda'))方法将需要使用的参数放入CUDA。
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)