Pytorch模型檔案`*.pt`與`*.pth` 的儲存與載入

华小电發表於2024-07-19
PyTorch模型

1.*.pt檔案

.pt檔案儲存的是模型的全部，在載入時可以直接賦值給新變數model = torch.load("filename.pt")。

具體操作:

  • （1）. 模型的儲存
torch.save(model,"Path/filename.pt")
  • （2）. 模型的載入
model = torch.load("filename.pt")

注意：torch.load()的引數使用字串引數。

2. .pth檔案

.pth儲存的是模型引數，透過字元字典進行儲存，在載入該類檔案時應該先例項化一個具體的模型，然後對新建立的空模型，進行引數賦予。

具體操作:

  • （1）. 模型的儲存
torch.save(model.state_dict(), PATH)
  • （2）. 模型的載入
model = nn.Module() # 這裡要先例項化模型
model.load_state_dict(torch.load("filename.pth"))

操作例項

  1. 首先定義一個模型作為例子
# Define model
class TheModelClass(nn.Module):
    # 類的初始化
    def __init__(self):
        # 繼承父類 nn.Module 的屬性和方法
        super(TheModelClass, self).__init__()
        # Inputs_channel, Outputs_channel, kernel_size
        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):
        # 卷積作用後，使用relu進行非線性化，最後使用池化操作進行特徵個數，引數量的降低
        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])

2. 現在開始進行模型的儲存與載入

PATH = "/home/深度學習/model"
#  第一種模型儲存和載入方式
torch.save(model.state_dict(), PATH+"/TheModuleClass.pth")
model = TheModelClass()
model.load_state_dict(torch.load("/home/深度學習/model/TheModuleClass.pth"))

for param_tensor in model.state_dict():
    print(f"{param_tensor}<<<{model.state_dict()[param_tensor].size()}")
print(model)

# 輸出結果
'''
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])
TheModelClass(
  (conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
  (fc1): Linear(in_features=400, out_features=120, bias=True)
  (fc2): Linear(in_features=120, out_features=84, bias=True)
  (fc3): Linear(in_features=84, out_features=10, bias=True)
)
'''
#  第二種模型儲存和載入方式

torch.save(model, PATH + "/the_module_class.pt")
model = torch.load(PATH + "/the_module_class.pt")

for param_tensor in model.state_dict():
    print(f"{param_tensor} <<< {model.state_dict()[param_tensor].size()}")
print(model)

#  輸出結果
'''
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])
TheModelClass(
  (conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
  (fc1): Linear(in_features=400, out_features=120, bias=True)
  (fc2): Linear(in_features=120, out_features=84, bias=True)
  (fc3): Linear(in_features=84, out_features=10, bias=True)
)
'''

總結

  • 這裡推薦使用第二種方法，因為儲存和載入檔案簡單，而且生成的二進位制檔案區分程度高。
  • torch.save() 儲存模型的引數，為以後模型推理核模型恢復提供了更加方便更加靈活的方法。
  • 一定要在模型評估時， 關閉批次規範化和丟棄法, 僅僅在模型訓練時有用，模型推理時一定要關閉(所謂模型推理，指是使用模型進行的實際應用)
  • 載入.pth 要先例項化，再進行引數的承接。

相關文章