如何从PyTorch .pt2文件恢复模型类并复用预训练权重？

首先明确：.pt2是PyTorch 2.x的TorchScript模型文件，它本质是序列化的计算图（包含权重和执行逻辑），并非保存的模型类实例（比如用torch.save(model, path)保存的格式）。所以直接从.pt2文件里恢复原始的Python模型类文件是不可能的——因为TorchScript导出时已经把模型结构编译成了静态计算图，丢失了原始类的定义信息（比如类名、自定义方法、注释等）。

不过你可以通过以下方式实现“修改结构后复用原有权重”的需求：

• 从推理代码反推模型结构：既然能访问加载模型的推理代码，通常会有torch.jit.load("model.pt2")的调用。加载模型后，通过print(model)或model.graph可以查看计算图的结构，从中解析出网络层类型、输入输出维度、各层参数形状等信息，据此手动写出对应的Python模型类。
• 映射原有权重到新模型：写出新模型类后，通过original_model.state_dict()获取.pt2模型的权重字典，将这些权重对应到新模型的参数上。如果参数名完全匹配，直接用new_model.load_state_dict(original_weights)加载；若存在参数名差异，可手动调整键名映射，再用new_model.load_state_dict(original_weights, strict=False)（strict=False允许忽略不匹配的参数）。
• 验证权重复用效果：加载权重后，对比新模型和原.pt2模型的推理结果，确保输出一致（或误差在可接受范围内），避免权重映射错误导致性能异常。

简单示例代码：

# 加载原.pt2模型
original_model = torch.jit.load("model.pt2")
# 获取原模型权重字典
original_weights = original_model.state_dict()

# 手动编写反推的模型类
class CustomModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv_layer = nn.Conv2d(3, 64, kernel_size=3)
        self.fc_layer = nn.Linear(64*28*28, 10)
    
    def forward(self, x):
        x = self.conv_layer(x)
        x = x.flatten(1)
        x = self.fc_layer(x)
        return x

# 实例化新模型并加载权重
new_model = CustomModel()
new_model.load_state_dict(original_weights, strict=False)

需要注意的限制：

• 如果原模型包含复杂控制流（比如循环、条件判断），TorchScript会将逻辑固化，反推结构时需要反复核对计算图细节，难度较高。
• 若原模型使用了自定义操作（Custom Operator），需确保新模型能正确调用这些操作，否则权重加载后无法正常运行。

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