保存和加載模型
在完成60分鐘入門之后浑塞，接下來有六節(jié)tutorials和五節(jié)關(guān)于文本處理的tutorials借跪。爭取一天一節(jié)。不過重點是關(guān)注神經(jīng)網(wǎng)絡(luò)構(gòu)建和數(shù)據(jù)處理部分酌壕。
本節(jié)主要是用于解決模型的保存和加載掏愁。會的直接跳過就好歇由。我也只是做記錄，這篇搞定就直接進(jìn)入NLP部分果港。
三個核心函數(shù)：
- torch.save：將序列化的對象保存在硬盤上沦泌，使用Python的pickle來序列化。
- torch.load：使用pickle的拆包功能將硬盤上的序列化文件導(dǎo)入內(nèi)存中辛掠。
- torch.nn.Module.load_state_dict：加載一個模型的參數(shù)字典谢谦。
主要目錄
- 什么是state_dict?
- Saving & Loading Model for Inference
- Saving & Loading a General Checkpoint
- Saving Multiple Models in One File
- Warmstarting Model Using Parameters from a Different Model
- Saving & Loading Model Across Devices
1.什么是state_dict?
在PyTorch中，torch.nn.Module的可學(xué)習(xí)參數(shù)(即權(quán)重和偏差)公浪，模塊模型包含在model's參數(shù)中(通過model.parameters()訪問)。state_dict是個簡單的Python dictionary對象船侧，它將每個層映射到它的參數(shù)張量欠气。
注意，只有具有可學(xué)習(xí)參數(shù)的層(卷積層镜撩、線性層等)才有model's state_dict中的條目预柒。優(yōu)化器對象(connector .optim)也有一個state_dict，其中包含關(guān)于優(yōu)化器狀態(tài)以及所使用的超參數(shù)的信息袁梗。
可以通過遍歷模型中的state_dict每一個tensor來查看宜鸯。
簡單構(gòu)建一個模型，這里我用了兩種寫法遮怜，注釋掉的那一種是比較通用的淋袖，也是建議使用的。
import torch
import torch.nn as nn
import torch.nn.functional as F
# Define model
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 farward(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
# class TheModelClass(nn.Module):
# def __init__(self):
# super(TheModelClass, self).__init__()
# self.conv = torch.nn.Sequential()
# self.conv.add_module('conv1', nn.Conv2d(3, 6, 5))
# self.conv.add_module('pool', nn.MaxPool2d(2, 2))
# self.conv.add_module('conv2', nn.Conv2d(6, 16, 5))
# self.dense = torch.nn.Sequential()
# self.dense.add_module('fc1', nn.Linear(16 * 5 * 5, 120))
# self.dense.add_module('fc2', nn.Linear(120, 84))
# self.dense.add_module('fc3', nn.Linear(84, 10))
#
# def forward(self, x):
# conv_out = self.conv(x)
# res = conv_out.view(conv_out.size(0), -1)
# out = self.dense(res)
# return out
# Initialize model
model = TheModelClass()
# Initialize optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)
print(model)
查看以下模型的state_dict
print('Model state dict:')
for param in model.state_dict():
print(param, "\t", model.state_dict()[param].size())
print('Optimizer state dict:')
for var in optimizer.state_dict():
print(var, "\t", optimizer.state_dict()[var])
TheModelClass(
(conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
(pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
(fc1): Linear(in_features=400, out_features=120, bias=True)
(fc2): Linear(in_features=120, out_features=84, bias=True)
(fc3): Linear(in_features=84, out_features=10, bias=True)
)
Model state dict:
conv1.weight torch.Size([6, 3, 5, 5])
conv1.bias torch.Size([6])
conv2.weight torch.Size([16, 6, 5, 5])
conv2.bias torch.Size([16])
fc1.weight torch.Size([120, 400])
fc1.bias torch.Size([120])
fc2.weight torch.Size([84, 120])
fc2.bias torch.Size([84])
fc3.weight torch.Size([10, 84])
fc3.bias torch.Size([10])
Optimizer state dict:
state {}
param_groups [{'lr': 0.0001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [2761682742728, 2761682742872, 2761682742944, 2761682743016, 2761682743088, 2761682743160, 2761682743232, 2761682796616, 2761682796688, 2761682796760]}]
2. 保存和加載模型的推理
2.1 保存和加載state_dict
torch.save(model.state_dict(), PATH)
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()
需要注意的一點是：必須調(diào)用model.eval()锯梁，以便在運行推斷之前將dropout和batch規(guī)范化層設(shè)置為評估模式即碗。如果不這樣做，將會產(chǎn)生不一致的推斷結(jié)果陌凳。
2.2 保存和加載整個模型
torch.save(model, PATH)
model=torch.load(PATH)
model.eval()
3. 保存和加載一般檢查點以進(jìn)行推理和/或恢復(fù)訓(xùn)練
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss,
...
}, PATH)
加載會稍微麻煩一些
model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)
checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']
model.eval()
# - or -
# model.train()
保存通用檢查點時剥懒，用于推理或恢復(fù)訓(xùn)練，必須保存的不僅僅是模型的state_dict合敦。保存優(yōu)化器的state_dict也很重要初橘，因為它包含作為模型訓(xùn)練更新的緩沖區(qū)和參數(shù)。您可能想要保存的其他項目是您停止使用的紀(jì)元充岛，最新記錄的訓(xùn)練損失保檐，外部torch.nn.Embedding等。
要保存多個組件崔梗，請在字典中組織它們并使用torch.save()來序列化字典展东。常見的PyTorch約定是使用.tar文件擴(kuò)展名保存這些檢查點。
要加載項目炒俱，首先初始化模型和優(yōu)化器盐肃，然后使用torch.load()在本地加載字典爪膊。可以通過簡單地查詢字典來輕松訪問已保存的項目砸王。
請記住推盛，在運行推理之前，必須調(diào)用model.eval()將dropout和批處理規(guī)范化層設(shè)置為評估模式谦铃。如果不這樣做耘成，將導(dǎo)致不一致的推理結(jié)果。
如果希望恢復(fù)訓(xùn)練驹闰，請調(diào)用model.train()以確保這些圖層處于訓(xùn)練模式瘪菌。
