PyTorch模型保存與加載的方法

這篇文章主要介紹了PyTorch模型保存與加載的方法的相關知識，內容詳細易懂，操作簡單快捷，具有一定借鑒價值，相信大家閱讀完這篇PyTorch模型保存與加載的方法文章都會有所收獲，下面我們一起來看看吧。

state_dict 是一個Python字典，將每一層映射成它的參數張量。注意只有帶有可學習參數的層（卷積層、全連接層等），以及注冊的緩存（batchnorm的運行平均值）在state_dict 中才有記錄。state_dict同樣包含優化器對象，存儲了優化器的狀態，所使用到的超參數。

一個簡單的例子

# 定義模型
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)

# 打印模型的 state_dict
print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

# 打印優化器的 state_dict
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])

輸出：

Model's 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's state_dict:
state    {}
param_groups     [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072, 4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]

保存/加載 state_dict（推薦）

保存：

torch.save(model.state_dict(), PATH)

加載：

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

要注意這個細節，如果使用nn.DataParallel在一臺電腦上使用了多個GPU，那么加載模型的時候也必須先進行nn.DataParallel。

保存模型的推理過程的時候，只需要保存模型訓練好的參數，使用torch.save()保存state_dict，能夠方便模型的加載。因此推薦使用這種方式進行模型保存。

記住一定要使用model.eval()來固定dropout和歸一化層，否則每次推理會生成不同的結果。

注意，load_state_dict()需要傳入字典對象，因此需要先反序列化state_dict再傳入load_state_dict()

保存/加載整個模型

保存：

torch.save(model, PATH)

加載：

# 模型類必須在別的地方定義
model = torch.load(PATH)
model.eval()

這種保存/加載模型的過程使用了最直觀的語法，所用代碼量少。這使用Python的pickle保存所有模塊。這種方法的缺點是，保存模型的時候，序列化的數據被綁定到了特定的類和確切的目錄。這是因為pickle不保存模型類本身，而是保存這個類的路徑，并且在加載的時候會使用。因此，當在其他項目里使用或者重構的時候，加載模型的時候會出錯。

一般來說，PyTorch的模型以.pt或者.pth文件格式保存。

一定要記住在評估模式的時候調用model.eval()來固定dropout和批次歸一化。否則會產生不一致的推理結果。

保存加載用于推理的常規Checkpoint/或繼續訓練

保存：

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()
# - 或者 -
model.train()

在保存用于推理或者繼續訓練的常規檢查點的時候，除了模型的state_dict之外，還必須保存其他參數。保存優化器的state_dict也非常重要，因為它包含了模型在訓練時候優化器的緩存和參數。除此之外，還可以保存停止訓練時epoch數，最新的模型損失，額外的torch.nn.Embedding層等。

要保存多個組件，則將它們放到一個字典中，然后使用torch.save()序列化這個字典。一般來說，使用.tar文件格式來保存這些檢查點。

加載各個組件，首先初始化模型和優化器，然后使用torch.load()加載保存的字典，然后可以直接查詢字典中的值來獲取保存的組件。

同樣，評估模型的時候一定不要忘了調用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()
# - 或者 -
modelA.train()
modelB.train()

保存的模型包含多個torch.nn.Modules時，比如GAN，一個序列-序列模型，或者組合模型，使用與保存常規檢查點的方式來保存模型。也就是說，保存每個模型的state_dict和對應的優化器到一個字典中。我們可以保存任何能幫助我們繼續訓練的東西到這個字典中。

使用其他模型來預熱當前模型

保存：

torch.save(modelA.state_dict(), PATH)

加載：

modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)

在遷移學習或者訓練新的復雜模型時，加載部分模型是很常見的。利用經過訓練的參數，即使只有少數參數可用，也將有助于預熱訓練過程，并且使模型更快收斂。

在加載部分模型參數進行預訓練的時候，很可能會碰到鍵不匹配的情況（模型權重都是按鍵值對的形式保存并加載回來的）。因此，無論是缺少鍵還是多出鍵的情況，都可以通過在load_state_dict()函數中設定strict參數為False來忽略不匹配的鍵。

如果想將某一層的參數加載到其他層，但是有些鍵不匹配，那么修改state_dict中參數的key可以解決這個問題。

跨設備保存與加載模型

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')，此時，map_location動態地將tensors的底層存儲重新映射到CPU設備上。

上述代碼只有在模型是在一塊GPU上訓練時才有效，如果模型在多個GPU上訓練，那么在CPU上加載時，會得到類似如下錯誤：

KeyError: ‘unexpected key “module.conv1.weight” in state_dict'

原因是在使用多GPU訓練并保存模型時，模型的參數名都帶上了module前綴，因此可以在加載模型時，把key中的這個前綴去掉：

# 原始通過DataParallel保存的文件
state_dict = torch.load('myfile.pth.tar')
# 創建一個不包含`module.`的新OrderedDict
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k[7:] # 去掉 `module.`
    new_state_dict[name] = v
# 加載參數
model.load_state_dict(new_state_dict)

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)
# 往模型中輸入數據的時候不要忘記在任意tensor上調用input = input.to(device)

在把GPU上訓練的模型加載到GPU上時，只需要使用model.to(torch.devie('cuda'))將初始化的模型轉換為CUDA優化模型。同時確保在模型所有的輸入上使用.to(torch.device('cuda'))。注意，調用my_tensor.to(device)會返回一份在GPU上的my_tensor的拷貝。不會覆蓋原本的my_tensor，因此要記得手動將tensor重寫：my_tensor = my_tensor.to(torch.device('cuda'))。

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"))  # 選擇希望使用的GPU
model.to(device)

保存torch.nn.DataParallel模型

保存：

torch.save(model.module.state_dict(), PATH)

關于“PyTorch模型保存與加載的方法”這篇文章的內容就介紹到這里，感謝各位的閱讀！相信大家對“PyTorch模型保存與加載的方法”知識都有一定的了解，大家如果還想學習更多知識，歡迎關注億速云行業資訊頻道。