儲存和加載模型 pytorch

# 儲存和加載模型
# 當儲存和加載模型時，需要熟悉三個核心功能：
# 1、torch.save：将序列化對象儲存到磁盤。此函數使用Python的pickle子產品進行序列化。使用此函數可以儲存如模型、tensor、
# 字典等各種對象。
# 2、torch.load：使用pickle的unpickling功能将pickle對象檔案反序列化到記憶體。此功能還可以有助于裝置加載資料。
# 3、torch.nn.Module.load_state_dict：使用反序列化函數 state_dict 來加載模型的參數字典。

# 一、什麼是狀态字典：state_dict?
# 在PyTorch中，torch.nn.Module模型的可學習參數（即權重和偏差）包含在模型的參數中，（使用model.parameters()可以進行
# 通路）。 state_dict是Python字典對象，它将每一層映射到其參數張量。注意，隻有具有可學習參數的層（如卷積層，線性層等）
# 的模型 才具有state_dict這一項。目标優化torch.optim也有state_dict屬性，它包含有關優化器的狀态資訊，以及使用的超參數。
# 因為state_dict的對象是Python字典，是以它們可以很容易的儲存、更新、修改和恢複，為PyTorch模型和優化器添加了大量子產品。
# 下面通過從簡單模型訓練一個分類器中來了解一下state_dict的使用。
# 定義模型
import torch
import torch.nn as nn
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])

# 二、儲存和加載推理模型
# 2.1 儲存/加載state_dict（推薦使用）
# 儲存 torch.save(model.state_dict(), PATH)
# 加載
# model = TheModelClass(*args, **kwargs)
# model.load_state_dict(torch.load(PATH))
# model.eval()
1此部分儲存/加載過程使用最直覺的文法并涉及最少量的代碼。以 Python `pickle 子產品的方式來儲存模型。這種方法的缺點
是序列化資料受限于某種特殊的類而且需要确切的字典結構。這是因為pickle無法儲存模型類本身。相反，它儲存包含類的文
件的路徑，該檔案在加載時使用。 是以，當在其他項目使用或者重構之後，您的代碼可能會以各種方式中斷。
2在 PyTorch 中最常見的模型儲存使用‘.pt’或者是‘.pth’作為模型檔案擴充名。
3請記住，在運作推理之前，務必調用model.eval()設定 dropout 和 batch normalization 層為評估模式。如果不這麼做，可能
導緻模型推斷結果不一緻。

三、儲存和加載 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()
# - or -
model.train()
1當儲存成 Checkpoint 的時候，可用于推理或者是繼續訓練，儲存的不僅僅是模型的 state_dict 。儲存優化器的 
state_dict 也很重要, 因為它包含作為模型訓練更新的緩沖區和參數。你也許想儲存其他項目，比如最新記錄的訓練損失，
外部的torch.nn.Embedding層等等。
2要儲存多個元件，請在字典中組織它們并使用torch.save()來序列化字典。PyTorch 中常見的儲存checkpoint 是使用 .tar 
檔案擴充名。
3要附加元件目，首先需要初始化模型和優化器，然後使用torch.load()來加載本地字典。這裡,你可以非常容易的通過簡單查詢字
典來通路你所儲存的項目。
4請記住在運作推理之前，務必調用model.eval()去設定 dropout 和 batch normalization 為評估。如果不這樣做，有可能
得到不一緻的推斷結果。 如果你想要恢複訓練，請調用model.train()以確定這些層處于訓練模式。

四、在一個檔案中儲存多個模型
儲存
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()
1當儲存一個模型由多個torch.nn.Modules組成時，例如GAN(對抗生成網絡)、sequence-to-sequence (序列到序列模型), 
或者是多個模型融合, 可以采用與儲存正常檢查點相同的方法。換句話說，儲存每個模型的 state_dict 的字典和相對應的
優化器。如前所述，可以通過簡單地将它們附加到字典的方式來儲存任何其他項目，這樣有助于恢複訓練。
2 PyTorch 中常見的儲存 checkpoint 是使用 .tar 檔案擴充名。
3 要附加元件目，首先需要初始化模型和優化器，然後使用torch.load()來加載本地字典。這裡，你可以非常容易的通過簡單查
詢字典來通路你所儲存的項目。
4 請記住在運作推理之前，務必調用model.eval()去設定 dropout 和 batch normalization 為評估。如果不這樣做，有可
能得到不一緻的推斷結果。 如果你想要恢複訓練，請調用model.train()以確定這些層處于訓練模式。

五、使用在不同模型參數下的熱啟動模式
儲存
torch.save(modelA.state_dict(), PATH)
加載
modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)
1在遷移學習或訓練新的複雜模型時，部分加載模型或加載部分模型是常見的情況。利用訓練好的參數，有助于熱啟動訓練過程，
并希望幫助你的模型比從頭開始訓練能夠更快地收斂。
2無論是從缺少某些鍵的 state_dict 加載還是從鍵的數目多于加載模型的 state_dict , 都可以通過在load_state_dict()函
數中将strict參數設定為 False 來忽略非比對鍵的函數。
3如果要将參數從一個層加載到另一個層，但是某些鍵不比對，主要修改正在加載的 state_dict 中的參數鍵的名稱以比對要在
加載到模型中的鍵即可。

六、通過裝置儲存/加載模型
儲存到 CPU、加載到 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.device('cpu')傳遞給torch.load()函數中的map_location參數.在這種情況下，
使用 map_location參數将張量下的存儲器動态的重新映射到CPU裝置。

儲存到 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)
# 確定在你提供給模型的任何輸入張量上調用input = input.to(device)
在GPU上訓練并把模型儲存在GPU，隻需要使用model.to(torch.device('cuda'))，将初始化的 model 轉換為 CUDA 優化模型。
另外，請務必在所有模型輸入上使用.to(torch.device('cuda'))函數來為模型準備資料。請注意，
調用my_tensor.to(device)會在GPU上傳回my_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"))  # Choose whatever GPU device number you want
model.to(device)
# 確定在你提供給模型的任何輸入張量上調用input = input.to(device)
在CPU上訓練好并儲存的模型加載到GPU時，将torch.load()函數中的map_location參數設定為cuda:device_id。這會将模型
加載到 指定的GPU裝置。接下來，請務必調用model.to(torch.device('cuda'))将模型的參數張量轉換為 CUDA 張量。最後，
確定在所有模型輸入上使用  .to(torch.device('cuda'))函數來為CUDA優化模型。請注意，調用my_tensor.to(device)會
在GPU上傳回my_tensor的新副本。它不會覆寫my_tensor。 是以， 請手動覆寫張量
my_tensor = my_tensor.to(torch.device('cuda'))。

儲存 torch.nn.DataParallel 模型
儲存
torch.save(model.module.state_dict(), PATH)
加載
# 加載任何你想要的裝置
torch.nn.DataParallel是一個模型封裝，支援并行GPU使用。要普通儲存 DataParallel 模型, 請儲存
model.module.state_dict()。 這樣，你就可以非常靈活地以任何方式加載模型到你想要的裝置中。