天天看點
傳回

pytorch模型加載測試_pytorch儲存和加載模型

pytorch模型加載測試_pytorch儲存和加載模型

torch.save：将序列化的對象儲存到磁盤。此函數使用Python的pickle進行序列化。使用此功能可以儲存各種對象的模型，tensor和dict。

torch.load：使用pickle 的unpickling将目标檔案反序列化到記憶體中。該功能也能夠将資料加載到某個裝置（Saving & Loading Model Across Devices）

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

state_dict

 是什麼？

在PyTorch中，torch.nn.Module模型的可學習參數（即權重和偏差）包含在模型的參數中（可通過model.parameters()擷取）。 state_dict 隻是一個Python字典對象，它将每個圖層映射到其參數tensor。請注意，隻有具有可學習參數的層（卷積層，線性層等）和已注冊的緩沖區（batchnorm的running_mean）才在模型的state_dict中存在。torch.optim也有state_dict，其中包含有關優化器狀态以及所用超參數的資訊。由于state_dict對象是Python詞典，是以可以輕松地儲存，更新，更改和還原它們，進而為PyTorch模型和優化器增加了很多子產品化。

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

output：

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]}]

Saving & Loading Model for Inference

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 and batch normalization layers設定為eval() 模式，不這樣做将産生不一緻的推斷結果。

請注意，load_state_dict（）函數采用字典對象，而不是儲存對象的路徑。這意味着您必須先反序列化儲存的state_dict，然後再将其傳遞給load_state_dict（）函數。即，不能使用model.load_state_dict（PATH）進行加載。

Saving & Loading a General Checkpoint for Inference and/or Resuming Training

Sava：

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

儲存用于檢查或繼續訓練的正常checkpoint時，要儲存的不隻是模型的state_dict，儲存優化器的state_dict也很重要，因為它包含随着模型訓練而更新的緩沖區和參數，可能要儲存的其他是您停止的epoch，最新記錄的訓練損失，外部torch.nn.Embedding層等。要儲存多個元件時，把他們組織字典格式，在外面用torch.sava()儲存。常見的PyTorch約定是使用.tar檔案擴充名儲存這些checkpoint。

要附加元件目，請首先初始化模型和優化器，然後使用torch.load()在本地加載字典。從這裡，您可以通過簡單地查詢字典來輕松通路儲存的項目。

請記住，在運作推理之前，必須調用model.eval()來将dropout 和 batch normalization layers設定為eval模式。不這樣做将産生不一緻的推斷結果。

如果您希望恢複訓練，請調用model.train()以確定這些層處于訓練模式。

Saving Multiple Models in One File

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，序列到序列模型或模型集合）時，将采用與儲存正常checkpoint相同的方法。換句話說，儲存每個模型的state_dict和相應的優化器的字典。如前所述，您可以儲存任何其他可以幫助您恢複訓練的項，隻需将它們添加到字典中即可。常見的PyTorch約定是使用.tar檔案擴充名儲存這些checkpoint。要加載模型，請首先初始化模型和優化器，然後使用torch.load()在本地加載字典。從這裡，您可以通過簡單地查詢字典來輕松通路儲存的項目。

請記住，在運作推理之前，必須調用model.eval()來将dropout 和 batch normalization layers設定為評估模式，不這樣做将産生不一緻的推斷結果。

如果您希望恢複訓練，請調用model.train()将這些層設定為訓練模式。

Warmstarting Model Using Parameters from a Different Model

Save：

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

Load：

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

在transfer learning或訓練新的複雜模型時，部分加載模型是常見方案。利用經過訓練的參數，即使隻有少數幾個可用的參數，也将有助于熱啟動訓練過程，并希望與從頭開始訓練相比，可以更快地收斂模型。無論是從缺少某些key（這裡說的key,是字典中的key-value）的部分state_dict加載，還是加載比要加載的模型更多的key的state_dict，都可以在load_state_dict()函數中将strict參數設定為False，以忽略不比對的keys。如果要将參數從一層加載到另一層，但是某些鍵不比對，隻需更改要加載的state_dict中參數key的名稱，以比對要加載到的模型中的key。

Saving & Loading Model Across Devices

1、Save on GPU, Load on 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))

2、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

3、Save on CPU, Load on GPU

Save:

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

Save:

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

4、Saving

torch.nn.DataParallel

 Models

Save:

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

Load:

# Load to whatever device you want

torch.nn.DataParallel 是一個模型包裝程式，可以并行使用GPU，要以一般方式儲存DataParallel模型，請儲存model.module.state_dict()。這樣，您可以靈活地将所需的模型加載到所需的任何裝置。

pytorch模型加載測試
上一篇: pytorch之模型的儲存與加載
下一篇: 模型的儲存與加載（pytorch）