PyTorch模型保存与加载全指南：从state_dict到checkpoint实战避坑

搞PyTorch这几年，被问得最多的一个问题就是“模型怎么保存怎么加载”。说实话，这话题看着简单，但真到了实战里，坑一点都不少。我印象最深的一次，是跑了一个两天的训练任务，结果第三天迭代代码的时候发现原脚本被覆盖，模型文件也没存对位置，整个结果直接没了。后来我养成一个习惯：不管项目多小，训练一启动就先写好保存逻辑，甚至先空跑一遍确认能存、能读、能恢复。就是这个“无聊”的习惯，帮我省下了不知道多少个重训的夜晚。

这篇文章不整那些虚的，直接围绕“Pytorch保存和加载模型”这个主题，把两种核心保存方式、checkpoint恢复、分布式训练的坑、模型文件体积优化、以及版本升级带来的兼容性问题全部梳理一遍。不管你是刚入门PyTorch的新手，还是写了不少训练脚本但没系统整理过保存策略的老手，这篇文章都能让你少踩几个坑，把时间花在真正有用的地方。

1. 先搞清楚两种保存方式，再动手写代码

很多教学帖一上来就给代码，但没解释为什么有两种保存方式，区别是什么。这一步搞不清楚，后面遇到问题就是两眼一抹黑。

1.1 state_dict：只存参数，不存结构

state_dict是PyTorch中一个非常核心的概念。简单说，它就是模型里所有可学习参数（weight、bias、BN层的running_mean等）组成的一个Python字典对象，key是参数名，value是Tensor。保存模型时，推荐的做法是只保存这个字典，而不是把整个模型对象都序列化。

为什么推荐只存state_dict？它的数据量小、结构清晰、加载灵活，而且和代码结构解耦。你训练时定义了一个五层卷积的模型，只要在加载时用相同的类定义重新实例化一个模型，然后调用load_state_dict，参数就能一一对上。这就好比你搬家时只搬一箱衣服（参数），而不是把整个衣柜（模型结构）都搬走。只要到了新家重新组装一个衣柜，衣服直接放进去就行。

保存state_dict的代码非常简单：

python复制# 保存
torch.save(model.state_dict(), 'model_weights.pth')

# 加载
model = MyModel()  # 需要先实例化一个结构一致的模型
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

这段代码几乎出现在所有PyTorch教程里，但实际操作中会衍生出很多细节，比如模型类定义位置变了怎么办、实例化时参数不一致怎么办、加载后需不需要调model.eval()。这些细节我放在后面专门讲，先记住一个原则：state_dict保存方式的关键前提是“模型结构在加载时已知且一致”。

1.2 保存完整模型：简单但脆弱

另一种方式是直接保存整个模型对象：

python复制# 保存
torch.save(model, 'model_full.pth')

# 加载
model = torch.load('model_full.pth')
model.eval()

这种方式的优点是非常省事，不需要在加载时重新实例化模型，模型结构也跟着文件走。但代价是文件体积大（会把很多附属信息也存进去）、兼容性差（PyTorch版本升级后很可能加载失败）、而且容易把模型类定义里的临时状态一并序列化，导致加载后模型行为异常。

我自己的经验是，这种方式适合快速演示、或者模型结构本身非常简单且后续不会再改代码的情况。一旦进入正式项目和长期迭代，我强烈建议用state_dict方式，因为模型结构是代码的一部分，应该由代码管理，而不是被埋在序列化文件里。

1.3 一张表看清两者的取舍

注意：无论是哪种方式，保存之前最好先确认目标保存目录存在，否则会直接抛出FileNotFoundError。通常是提前用os.makedirs(save_dir, exist_ok=True)处理一下，这点小细节天天有人踩。

2. 保存模型的标准姿势与工程化细节

知道两种方式之后，接下来要做的是把保存动作做成一套规范流程，而不是训练结束随手保存一下。工程化的保存策略能让你在断点恢复、调参回溯、多轮实验对比时不被动。

2.1 基础保存代码与目录规划

先看一个工程化一点的保存写法：

python复制import torch
import os

class Trainer:
    def __init__(self, model, save_dir="./checkpoints"):
        self.model = model
        self.save_dir = save_dir
        os.makedirs(save_dir, exist_ok=True)

    def save_model(self, epoch, filename=None):
        if filename is None:
            f