本人學(xué)習(xí)pytorch主要參考官方文檔和 莫煩Python中的pytorch視頻教程斩祭。
后文主要是對pytorch官網(wǎng)的文檔的總結(jié)汹碱。
代碼來自pytorch官網(wǎng)
-
torch.save：將對象序列化到硬盤上，該對象可以是 Models, tensors和 dictionaries 等项玛。實際上是使用了python的pickle方法貌笨。 -
torch.load：將硬盤上序列化的對象加載設(shè)備中。實際是使用了pickle的解包方法襟沮。 -
torch.nn.Module.load_state_dict：通過反序列化state_dict加載模型參數(shù)字典锥惋。
模型參數(shù)
在pytorch中torch.nn.Module模型的參數(shù)存放在模型的parameters中(model.parameters()),而state_dict是參數(shù)tensor的字典昌腰。僅當(dāng)該層有可學(xué)習(xí)的參數(shù)才會有屬性state_dict。torch.optim也有state_dict屬性膀跌，其中包含的是優(yōu)化器的參數(shù)遭商，即網(wǎng)絡(luò)所使用的超參數(shù)。
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()
# 初始化優(yōu)化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# 打印模型參數(shù)
print("Model's state_dict:")
for param_tensor in model.state_dict():
print(param_tensor, "\t", model.state_dict()[param_tensor].size())
# 打印優(yōu)化器參數(shù)
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
print(var_name, "\t", optimizer.state_dict()[var_name])
保存和加載模型
保存的文件結(jié)尾一般使用.pt或.pth捅伤。最后的model.eval()表示將drop和batch nromalization層設(shè)置為測試模式劫流。可通過mode.train()轉(zhuǎn)化為訓(xùn)練模式丛忆。
模型參數(shù)
#保存模型參數(shù)
torch.save(model.state_dict(), PATH)
#加載模型參數(shù)
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()
加載不同的模型
加載他人訓(xùn)練的模型祠汇，可能需要忽略部分層。則將load_state_dict方法的strict參數(shù)設(shè)置為False熄诡。則模型金價在modelB中名稱對應(yīng)的層的參數(shù)可很。
torch.save(modelA.state_dict(), PATH)
modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)
加載不同設(shè)備的模型
將由GPU保存的模型加載到CPU上。
將torch.load()函數(shù)中的map_location參數(shù)設(shè)置為torch.device('cpu')
device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))
將由GPU保存的模型加載到GPU上凰浮。確保對輸入的tensors調(diào)用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調(diào)用input = input.to(device)方法导坟。map_location是將模型加載到GPU上，model.to(torch.device('cuda'))是將模型參數(shù)加載為CUDA的tensor圈澈。最后保證使用.to(torch.device('cuda'))方法將需要使用的參數(shù)放入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)
