Pytorch 模型的存儲與加載

本文主要內容來自Pytorch官方文檔推薦的一篇英文博客, 本文主要介紹了在Pytorch中模型的存儲方法, 以及存儲形式, 以及Pytorch存儲模型正真存儲的是模型的什么結構. 以及加載模型的時候, 模型的哪些數據會被加載. 以及加載后的形式.

首先大致講下三個最主要的函數的功能:

torch.save: 將序列化的對象存儲到硬盤中.此函數使用Python的pickle實用程序進行序列化. 對於數據類型都可以進行序列化存儲, 模型, 張量, 以及字典, 等各種數據對象都可以使用該函數存儲.

torch.load: 該函數使用的是 pickle 的階序列化過程, 並將結果存如內存中, 該函數也促進設備加載數據.

torch.nn.Module.load_state_dict: 使用反序列化的 state_dict 加載模型的參數字典

模型的加載

state_dict 是什么

在一個Pytorch模型中, 通常是 torch.nn.module , 模型中可學習的參數被包含在模型的參數中, 通常是可以使用 model.parameters() 函數訪問, 通常都是使用該方法訪問的. state_dict只是一個Python字典對象，它將每個圖層映射到其參數張量, 這個字典的 key 是圖層的 'name', 注意, 只有該層有可學習的參數的層, 也就是可以通過反向傳播優化的層, 以及 registered buffers (batchnorm’s running_mean) 才會在 state_dict 中有存儲條目. 優化器對象（torch.optim）也具有state_dict，其中包含有關優化器狀態以及所用超參數的信息. state_dict 的本質是對模型進行了字典化.

state_dict的字典形式使得對模型的操作更加的靈活, 例如直接導出模型, 修改其中的參數信息, 或者對層數進行修改等, 然后繼續將模型保留. 還是使用一個簡單的模型舉個例子:

class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        # 這里卷積核的大小是 5, 個數是 6, 輸入的 width 是 3
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # 兩次卷積的結果應該是 5x5x16 的矩陣
        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

# Initialize model
model = TheModelClass()

# Initialize optimizer
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# Print model's state_dict
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
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 存儲與加載模型

save:

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

Load 模型:

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

從模型存儲的角度, 存儲模型的時候, 唯一需要存儲的是該模型訓練的參數, torch.save() 函數也可以存儲模型的 state_dict. 使用該方法進行存儲, 模型被看做字典形式, 所以對模型的操作更加靈活. 在這種形式下常見的PyTorch約定是使用.pt或.pth文件擴展名保存模型.

注意, 加載模型之后, 並不能直接運行, 需要使用 model.eval() 函數設置 Dropout 與層間正則化. 另一方面, 該方法在存儲模型的時候是以字典的形式存儲的, 也就是存儲的是模型的字典數據, Pytorch 不能直接將模型讀取為該形式, 必須先 torch.load() 該模型, 然后再使用 load_state_dict().

將模型作為整體存儲與加載

Save:

torch.save(model, PATH)

Load:

# Model class must be defined somewhere
model = torch.load(PATH)
model.eval()

使用該方法相當於跳過了對模型的 state_dict 描述的過程, 而是直接使用 python 的 pickle 包, 這種方法的缺點是, 模型的存儲形式與加載形式十分固定, 這樣做的原因是因為pickle不會保存模型類本身. 而是存出來包含該文件的路徑，該路徑在加載時使用. 因此，在其他項目中使用或重構后，代碼可能會以各種方式中斷. 但是這種方法存儲的文件的類型與前面的方法一樣. 同樣, 以該方法加載模型運行之前需要調用 model.eval() .

存儲與加載一般的 Checkpoint

Save:

torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)

Load:

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

可以看出 checkpoints 是模型主要內容的一個字典, 基本包含了模型各種數據, 例如上面的例子模型的參數使用的是 optimizer.state_dict().

存儲 checkpoints 主要目的是為了方便加載模型繼續訓練, 將所有的信息存儲, 加載模型繼續訓練的時候就會更加方便. 為了存儲一個訓練過程的多種信息, 最好的方式是使用 dictionary 進行序列化, 這樣存儲一個訓練模型的形式是 .tar, 要加載項目，首先初始化模型和優化器，然后使用torch.load() 在本地加載字典.從這里開始, 只需按期望查詢字典即可輕松訪問已保存的項目. 請記住，在運行推理之前，必須調用model.eval() 來將 Dropout 和 Batch 正則化設置為評估模式, 不這樣做將產生不一致的推斷結果. 如果恢復訓練，那么調用model.train() 以確保這些層處於訓練模式.

在一個文件中存儲多個模型

save:

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)

Load:

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

保存包含多個 torch.nn.Modules 的模型（例如GAN，序列到序列模型或模型集合）時，將采用與保存常規檢查點相同的方法。 換句話說，保存每個模型的state_dict和相應的優化器的字典. 如前所述，您可以保存任何其他可以幫助您恢復培訓的項目，只需將它們添加到字典中即可. 使用該方法存儲的文件也是 .tar 形式的, 要加載模型，請首先初始化模型和優化器，然后使用torch.load（）在本地加載字典。 從這里，您只需按期望查詢字典即可輕松訪問已保存的項目.

跨平台模型保存與加載

GPU 到 CPU

Save:

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

Load:

device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))

Save on GPU, Load on GPU

Save:

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

Load:

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)
# Make sure to call input = input.to(device) on any input tensors that you feed to the model

Save on CPU, Load on GPU

Save:

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

Load:

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)
# Make sure to call input = input.to(device) on any input tensors that you feed to the model

加載部分模型

舉個例子:

# build
        encoder = TransformerModel(params, dico, is_encoder=True, with_output=False)  # TODO: only output when necessary - len(params.clm_steps + params.mlm_steps) > 0
        decoder = TransformerModel(params, dico, is_encoder=False, with_output=True)
        
# reload pretrained word embeddings
        if params.reload_emb != '':
        	# 表示加載預訓練模型
            word2id, embeddings = load_embeddings(params.reload_emb, params)
            set_pretrain_emb(encoder, dico, word2id, embeddings)
            set_pretrain_emb(decoder, dico, word2id, embeddings)
            set_pretrain_emb(model2, dico, word2id, embeddings)

        # reload a pretrained model
        if params.reload_model != '':
            enc_path, dec_path = params.reload_model.split(',')
            assert not (enc_path == '' and dec_path == '')

            # reload encoder
            if enc_path != '':
                enc_reload = torch.load(enc_path, map_location=lambda storage, loc: storage.cuda(params.local_rank))
                # 預訓練模型是在 GPU 上訓練的
                enc_reload = enc_reload['model' if 'model' in enc_reload else 'encoder']
                # 導入存儲的文件的模型
                if all([k.startswith('module.') for k in enc_reload.keys()]):
                    enc_reload = {k[len('module.'):]: v for k, v in enc_reload.items()}
                # 這個過程相當於將model 反序列化為 state_dict的形式
                encoder.load_state_dict(enc_reload, strict=False)
                # 這個后面的 strict=False 就是對 encoder 與 enc_reload.state_dict之間差異進行處理, 如果encoder 的模型結構與 enc_reload模型結構
                # 不一樣的時候, 就會向 encoder 轉化, 也就是 encoder 不包含的層就不會導入, 例如這里 enc_reload 就是一個完整的 Transformer 模型, 但是
                # encoder 是不包含輸出部分的, 所以就不會加載這部分

對於該部分, 本文只是做了個簡單的例子介紹, 更詳細的內容參見傳送門 . 對於這個傳送門的例子, 如果我們先存儲一個大模型, 將大模型加載到小模型的時候, 使用:

path = 'xxx.pth'
model = Net()
model.load_state_dict(t.load(path), strict=False)
for module in model.named_modules():
    print(module)
for name, param in model.named_parameters():
    print(name, param)

從輸出可以看出模型向下兼容,