日本免费高清视频-国产福利视频导航-黄色在线播放国产-天天操天天操天天操天天操|www.shdianci.com

學(xué)無(wú)先后，達(dá)者為師

網(wǎng)站首頁(yè) 編程語(yǔ)言 正文

PyTorch模型保存與加載實(shí)例詳解_python

作者：LXYTSOS ? 更新時(shí)間： 2022-06-25 編程語(yǔ)言
  • torch.save：保存序列化的對(duì)象到磁盤，使用了Python的pickle進(jìn)行序列化，模型、張量、所有對(duì)象的字典。
  • torch.load：使用了pickle的unpacking將pickled的對(duì)象反序列化到內(nèi)存中。
  • torch.nn.Module.load_state_dict：使用反序列化的state_dict加載模型的參數(shù)字典。

state_dict 是一個(gè)Python字典，將每一層映射成它的參數(shù)張量。注意只有帶有可學(xué)習(xí)參數(shù)的層（卷積層、全連接層等），以及注冊(cè)的緩存（batchnorm的運(yùn)行平均值）在state_dict 中才有記錄。state_dict同樣包含優(yōu)化器對(duì)象，存儲(chǔ)了優(yōu)化器的狀態(tài)，所使用到的超參數(shù)。

一個(gè)簡(jiǎn)單的例子

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

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

# 打印優(yōu)化器的 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()

要注意這個(gè)細(xì)節(jié)，如果使用nn.DataParallel在一臺(tái)電腦上使用了多個(gè)GPU，那么加載模型的時(shí)候也必須先進(jìn)行nn.DataParallel。

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

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

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

保存/加載整個(gè)模型

保存：

torch.save(model, PATH)

加載：

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

這種保存/加載模型的過程使用了最直觀的語(yǔ)法，所用代碼量少。這使用Python的pickle保存所有模塊。這種方法的缺點(diǎn)是，保存模型的時(shí)候，序列化的數(shù)據(jù)被綁定到了特定的類和確切的目錄。這是因?yàn)閜ickle不保存模型類本身，而是保存這個(gè)類的路徑，并且在加載的時(shí)候會(huì)使用。因此，當(dāng)在其他項(xiàng)目里使用或者重構(gòu)的時(shí)候，加載模型的時(shí)候會(huì)出錯(cuò)。

一般來(lái)說(shuō)，PyTorch的模型以.pt或者.pth文件格式保存。

一定要記住在評(píng)估模式的時(shí)候調(diào)用model.eval()來(lái)固定dropout和批次歸一化。否則會(huì)產(chǎn)生不一致的推理結(jié)果。

保存加載用于推理的常規(guī)Checkpoint/或繼續(xù)訓(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()
# - 或者 -
model.train()

在保存用于推理或者繼續(xù)訓(xùn)練的常規(guī)檢查點(diǎn)的時(shí)候，除了模型的state_dict之外，還必須保存其他參數(shù)。保存優(yōu)化器的state_dict也非常重要，因?yàn)樗四Ｐ驮谟?xùn)練時(shí)候優(yōu)化器的緩存和參數(shù)。除此之外，還可以保存停止訓(xùn)練時(shí)epoch數(shù)，最新的模型損失，額外的torch.nn.Embedding層等。

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

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

同樣，評(píng)估模型的時(shí)候一定不要忘了調(diào)用model.eval()。

保存多個(gè)模型到一個(gè)文件

保存：

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()

保存的模型包含多個(gè)torch.nn.Modules時(shí)，比如GAN，一個(gè)序列-序列模型，或者組合模型，使用與保存常規(guī)檢查點(diǎn)的方式來(lái)保存模型。也就是說(shuō)，保存每個(gè)模型的state_dict和對(duì)應(yīng)的優(yōu)化器到一個(gè)字典中。我們可以保存任何能幫助我們繼續(xù)訓(xùn)練的東西到這個(gè)字典中。

使用其他模型來(lái)預(yù)熱當(dāng)前模型

保存：

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

加載：

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

在遷移學(xué)習(xí)或者訓(xùn)練新的復(fù)雜模型時(shí)，加載部分模型是很常見的。利用經(jīng)過訓(xùn)練的參數(shù)，即使只有少數(shù)參數(shù)可用，也將有助于預(yù)熱訓(xùn)練過程，并且使模型更快收斂。

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

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

跨設(shè)備保存與加載模型

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))

當(dāng)在CPU上加載一個(gè)GPU上訓(xùn)練的模型時(shí)，在torch.load()中指定map_location=torch.device('cpu')，此時(shí)，map_location動(dòng)態(tài)地將tensors的底層存儲(chǔ)重新映射到CPU設(shè)備上。

上述代碼只有在模型是在一塊GPU上訓(xùn)練時(shí)才有效，如果模型在多個(gè)GPU上訓(xùn)練，那么在CPU上加載時(shí)，會(huì)得到類似如下錯(cuò)誤：

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

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

# 原始通過DataParallel保存的文件
state_dict = torch.load('myfile.pth.tar')
# 創(chuàng)建一個(gè)不包含`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
# 加載參數(shù)
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)
# 往模型中輸入數(shù)據(jù)的時(shí)候不要忘記在任意tensor上調(diào)用input = input.to(device)

在把GPU上訓(xùn)練的模型加載到GPU上時(shí)，只需要使用model.to(torch.devie('cuda'))將初始化的模型轉(zhuǎn)換為CUDA優(yōu)化模型。同時(shí)確保在模型所有的輸入上使用.to(torch.device('cuda'))。注意，調(diào)用my_tensor.to(device)會(huì)返回一份在GPU上的my_tensor的拷貝。不會(huì)覆蓋原本的my_tensor，因此要記得手動(dòng)將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)

?總結(jié)

原文鏈接：https://blog.csdn.net/LXYTSOS/article/details/90639524

欄目分類
最近更新