PyTorch模型保存格式详解：.pkl vs .pth/.pt

核心特性

• 文件扩展名：.pkl

• 底层技术：Python Pickle模块的序列化格式

• 序列化原理：

  import pickle
  # 保存模型
  with open('model.pkl', 'wb') as f:
    pickle.dump(model, f)
  
  # 加载模型
  with open('model.pkl', 'rb') as f:
    model = pickle.load(f)

优势与局限

适用场景

• 实验性代码快速保存
• 多组件流水线打包（如预处理+模型）
• 短期本地调试

with open(model_export_path, 'wb') as file:
    d = {
        "clf": clf,  # 分类器（如MultinomialNB）
        "vectorizer": vectorizer,  # 文本特征提取器（如CountVectorizer）
        "tfidftransformer": tfidftransformer,  # TF-IDF权重转换器
    }
    pickle.dump(d, file)  # 序列化到.pkl文件

核心作用

1. 全状态保存：完整保留对象结构、参数及依赖关系，加载时无需重新初始化
2. 多对象打包：支持将预处理、模型、转换器等流水线组件统一存储，确保推理时逻辑一致性
3. Python生态友好：与Scikit-learn、PyTorch等库无缝集成，适合快速实验和本地调试

🔐 .pth/.pt 文件格式

核心特性

• 文件扩展名：.pth 或 .pt

• PyTorch标准格式：

  # 保存状态字典（推荐）
  torch.save(model.state_dict(), 'model.pth')
  
  # 加载状态字典
  model = MyModel()
  model.load_state_dict(torch.load('model.pth'))
  
  # 完整模型保存（不推荐）
  torch.save(model, 'model_complete.pt')

最佳实践

1. 优先保存状态字典（state_dict()）

2. 配合模型定义代码使用：

   class MyModel(nn.Module):
      def __init__(self):
          super().__init__()
          self.linear = nn.Linear(10, 2)
   
      def forward(self, x):
          return self.linear(x)

核心优势

• 🔒 更高安全性（无代码执行风险）
• 🔄 更好的版本兼容性
• 🚀 PyTorch生态原生支持
• 📦 更小的文件体积

⚖️ 格式对比与推荐

🏆 行业推荐方案

# 标准保存方式
torch.save(model.state_dict(), 'model.pth')

# 标准加载方式
model = MyModel()  # 需预先定义结构
model.load_state_dict(torch.load('model.pth'))

关键结论：

• 生产中禁用.pkl：存在安全漏洞和兼容性问题
• 始终使用.pth/.pt：配合state_dict()实现安全可靠的模型持久化
• 完整模型保存仅用于调试：避免生产环境使用torch.save(model)