仅通过.pt文件重训PyTorch模型的方法咨询

核心结论 #

必须编写与.pt文件对应模型完全一致的模型类，不能省略。因为绝大多数情况下，.pt文件仅保存模型的参数张量（即state_dict），不包含模型的结构定义信息。

具体操作步骤 #

1. 还原模型类定义
   你编写的MyModel类必须和原模型的类名、层结构、层命名、参数维度完全一致，否则加载参数时会出现键不匹配的错误。举个例子，原模型的定义如果是：

   import torch.nn as nn
   
   class MyModel(nn.Module):
       def __init__(self):
           super().__init__()
           self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
           self.pool = nn.MaxPool2d(2, 2)
           self.fc1 = nn.Linear(64 * 16 * 16, 128)
           self.fc2 = nn.Linear(128, 10)
       
       def forward(self, x):
           x = self.pool(nn.functional.relu(self.conv1(x)))
           x = x.flatten(1)
           x = nn.functional.relu(self.fc1(x))
           x = self.fc2(x)
           return x
   

   你写的类就必须和上面的代码丝毫不差，包括每个层的名称、参数设置。

2. 加载参数并启动重训
   实例化模型类后，用load_state_dict方法加载.pt文件中的参数，之后即可像训练新模型一样设置优化器、损失函数，启动重训：

   import torch
   import torch.optim as optim
   import torch.nn as nn
   
   # 先定义好和原模型一致的MyModel类（代码同上）
   
   # 实例化模型并加载参数
   model = MyModel()
   # 如果是多GPU训练保存的模型，需要添加map_location参数或去掉module前缀
   # model.load_state_dict(torch.load("your_model.pt", map_location=torch.device('cpu')))
   model.load_state_dict(torch.load("your_model.pt"))
   
   # 设置重训所需的组件
   criterion = nn.CrossEntropyLoss()
   optimizer = optim.Adam(model.parameters(), lr=1e-4)
   # 若只想微调部分层，可冻结其他层参数
   # for param in model.conv1.parameters():
   #     param.requires_grad = False
   
   # 重训循环示例
   for epoch in range(20):
       model.train()
       total_loss = 0.0
       for inputs, labels in your_training_dataloader:
           optimizer.zero_grad()
           outputs = model(inputs)
           loss = criterion(outputs, labels)
           loss.backward()
           optimizer.step()
           total_loss += loss.item()
       print(f"Epoch {epoch+1} | 训练损失: {total_loss/len(your_training_dataloader):.4f}")
   

3. 特殊情况说明
   若原模型是通过torch.save(model, "model.pt")保存的完整模型对象（而非state_dict），可以直接用model = torch.load("model.pt")加载完整模型，无需重新定义类。但这种方法兼容性极差，一旦PyTorch版本、依赖库或代码环境变化，很可能加载失败，因此不推荐使用。

内容的提问来源于stack exchange，提问作者Sir Art