1. 项目概述：为什么保存与加载是深度学习的“存档点”？

在PyTorch的世界里折腾过一阵子的朋友，大概率都经历过这种“血压飙升”的时刻：花了几个小时甚至几天训练好的模型，因为程序意外退出、内核重启或者想换个环境测试，结果模型和好不容易跑出来的中间结果全没了，一切又得从头开始。或者，当你精心调参的模型在测试集上表现惊艳，想要分享给同事或者部署到生产环境时，却发现不知道如何把这个“黑盒子”完整地打包带走。这些痛点，恰恰指向了PyTorch乃至所有深度学习框架中一个至关重要，却又容易被新手忽视的基础技能——模型的保存与加载。

简单来说，PyTorch中的保存与加载，就是为你的张量（Tensor）和模型（Model）创建“存档点”。它解决的远不止是“防止丢失”的问题。想象一下游戏存档：你可以随时从存档点继续，避免重复劳动；你可以把存档发给朋友，让他直接体验你后期的游戏内容；你还可以备份多个存档，尝试不同的剧情分支。在深度学习中，保存与加载机制扮演着同样的角色：模型持久化、训练断点续训、模型分享与部署以及迁移学习。无论是torch.save()一个简单的Tensor，还是处理复杂的包含优化器状态、学习率调度器的完整训练状态，其核心逻辑都是将内存中的Python对象（通常是状态字典state_dict）序列化到磁盘文件（通常是.pt或.pth后缀），并在需要时反序列化加载回内存。

从网络热词如“ad崩溃没保存”、“上次任务中保存”引发的普遍共鸣，到“pytorch安装”、“anaconda配置pytorch环境”后必然要面对的实际操作，再到“构建模型”、“开源模型质变”分享环节的刚需，掌握这套“存档”与“读档”的规范操作，是从深度学习入门迈向熟练应用的必经之路。本教程将彻底拆解PyTorch中保存与加载的每一个细节，让你不仅能应对常规场景，更能优雅处理那些容易踩坑的复杂情况。

2. 核心概念解析：state_dict——模型状态的“身份证”

在深入实操之前，必须理解一个核心概念：state_dict（状态字典）。这是PyTorch保存与加载机制的基石，不理解它，后面的所有操作都将是空中楼阁。

2.1 什么是state_dict？

你可以把state_dict理解为一个Python字典对象，它精确地映射了模型或优化器内部所有可学习参数（learnable parameters）和持久缓冲区（persistent buffers）的名字到其对应的Tensor值。

• 可学习参数：就是模型通过训练要更新的那些权重（weights）和偏置（biases）。例如，线性层（nn.Linear）中的weight和bias。
• 持久缓冲区：是模型的一部分，需要被保存和加载，但不会被优化器更新。例如，批归一化层（nn.BatchNorm2d）中用于推理的 running mean（running_mean）和 running variance（running_var）。

state_dict不包含模型的结构信息（比如用了几个层，层与层之间如何连接），它只关心这些层里具体的参数值。这种设计实现了模型架构与参数的解耦，带来了巨大的灵活性。

2.2 查看state_dict

让我们通过一个简单的例子直观感受一下：

import torch import torch.nn as nn # 定义一个简单的模型 class SimpleNet(nn.Module): def __init__(self): super().__init__() self.fc1 = nn.Linear(10, 5) self.bn = nn.BatchNorm1d(5) self.fc2 = nn.Linear(5, 2) def forward(self, x): x = self.fc1(x) x = self.bn(x) x = torch.relu(x) x = self.fc2(x) return x model = SimpleNet() print(“模型的 state_dict:“) for param_tensor in model.state_dict(): print(f“{param_tensor:20} | Shape: {model.state_dict()[param_tensor].size()}“) # 输出示例： # fc1.weight | Shape: torch.Size([5, 10]) # fc1.bias | Shape: torch.Size([5]) # bn.weight | Shape: torch.Size([5]) # 这是BatchNorm的gamma # bn.bias | Shape: torch.Size([5]) # 这是BatchNorm的beta # bn.running_mean | Shape: torch.Size([5]) # bn.running_var | Shape: torch.Size([5]) # bn.num_batches_tracked | Shape: torch.Size([]) # 跟踪的batch数 # fc2.weight | Shape: torch.Size([2, 5]) # fc2.bias | Shape: torch.Size([2])

同样，优化器也有自己的state_dict，它包含了优化器的状态（如动量缓存）以及其关联的参数组信息。

optimizer = torch.optim.SGD(model.parameters(), lr=0.01) print(“\n优化器的 state_dict 键:“) print(optimizer.state_dict().keys()) # 输出：dict_keys([‘state‘, ‘param_groups‘])

注意：model.state_dict()中的键（如fc1.weight）与模型类定义中属性的名字严格对应。这要求我们在定义模型时，给每个子模块（nn.Module）起一个唯一且清晰的名字，否则加载时会出现找不到对应键的错误。

2.3 为什么是state_dict，而不是直接保存整个模型？

这是新手常有的疑问。PyTorch也提供了torch.save(model, ‘model.pth‘)这种保存整个模型对象的方式（称为“pickle”整个模型）。但这种方式不被推荐，原因如下：

1. 兼容性差：保存的模型文件与定义模型的源代码文件强绑定。如果后续修改了模型类的代码（比如类名、结构），即使state_dict没变，也可能无法加载。
2. 安全性：pickle模块在反序列化时会执行存储的字节码，如果模型文件来自不可信的来源，可能存在安全风险。
3. 灵活性低：无法轻松地将参数加载到结构不同但部分层名匹配的模型中（这在迁移学习中很常见）。

因此，最佳实践始终是保存和加载state_dict。

3. 基础保存与加载操作详解

掌握了state_dict的概念后，我们就可以开始实战了。PyTorch提供了非常简洁的API：torch.save()和torch.load()。

3.1 保存与加载单个Tensor

这是最简单的场景，常用于保存预处理后的数据、中间特征或标签。

# 保存一个Tensor x = torch.randn(3, 4) torch.save(x, ‘tensor.pt‘) # 通常使用 .pt 或 .pth 后缀 # 加载一个Tensor x_loaded = torch.load(‘tensor.pt‘) print(torch.allclose(x, x_loaded)) # 输出: True

3.2 保存与加载模型参数（state_dict）

这是最常用、最推荐的方式。

# 1. 保存模型参数 model = SimpleNet() # ... 这里通常会有训练过程，更新model的参数 ... torch.save(model.state_dict(), ‘model_weights.pth‘) # 2. 加载模型参数 # 首先，必须实例化一个与保存时结构完全相同的模型 model_new = SimpleNet() # 然后，将保存的参数加载到这个新实例中 model_new.load_state_dict(torch.load(‘model_weights.pth‘)) # 最后，通常需要将模型设置为评估模式（如果用于推理） model_new.eval()

关键点：load_state_dict()函数要求目标模型（model_new）的state_dict键必须与加载的字典键完全匹配（包括名字和Tensor的形状）。如果不匹配，会抛出错误。你可以通过设置strict=False参数来忽略不匹配的键，但这通常意味着部分参数没有被加载，需要谨慎使用。

3.3 保存与加载整个训练检查点（Checkpoint）

在长时间训练（尤其是训练大模型）时，我们不仅需要保存模型参数，还需要保存优化器状态、当前的epoch、损失值等信息，以便从中断处恢复训练。这通常通过保存一个字典来实现。

# 定义一个检查点字典 checkpoint = { ‘epoch‘: 10, ‘model_state_dict‘: model.state_dict(), ‘optimizer_state_dict‘: optimizer.state_dict(), ‘loss‘: 0.05, ‘lr_scheduler_state_dict‘: scheduler.state_dict() # 如果有学习率调度器 } # 保存检查点 torch.save(checkpoint, ‘checkpoint_epoch_10.pth‘) # 加载检查点，恢复训练 checkpoint_loaded = torch.load(‘checkpoint_epoch_10.pth‘) model.load_state_dict(checkpoint_loaded[‘model_state_dict‘]) optimizer.load_state_dict(checkpoint_loaded[‘optimizer_state_dict‘]) epoch = checkpoint_loaded[‘epoch‘] loss = checkpoint_loaded[‘loss‘] scheduler.load_state_dict(checkpoint_loaded[‘lr_scheduler_state_dict‘]) # 恢复训练模式 model.train()

实操心得：检查点文件名最好包含关键信息，如checkpoint_epoch_{epoch}_loss_{loss:.4f}.pth。这样在文件夹里一眼就能看出哪个检查点性能最好、训练到了哪一步。对于超大规模训练，定期保存检查点并保留最好的几个是标准操作。

4. 进阶场景与疑难问题排查

实际项目中，情况往往比基础教程复杂。下面这些场景和问题，几乎每个PyTorch开发者都会遇到。

4.1 设备不匹配问题：CPU vs GPU

这是加载模型时最常见的错误之一。错误信息常类似于：

RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location=torch.device(‘cpu‘) to map your storages to the CPU.

或者键名中带有module.前缀（多GPU训练导致）。

问题根源：模型可能在GPU上训练并保存，但尝试在只有CPU的环境加载；或者反之。此外，使用DataParallel或DistributedDataParallel包装后，模型参数名会被自动加上module.前缀。

解决方案：torch.load()的map_location参数是你的救星。

# 场景1：在CPU上加载一个在GPU上保存的模型 model = SimpleNet() # 方法：指定map_location为‘cpu‘ state_dict = torch.load(‘gpu_trained_model.pth‘, map_location=torch.device(‘cpu‘)) model.load_state_dict(state_dict) # 场景2：强制加载到指定设备（更通用的写法） device = torch.device(‘cuda:0‘ if torch.cuda.is_available() else ‘cpu‘) state_dict = torch.load(‘model.pth‘, map_location=device) model.load_state_dict(state_dict) model.to(device) # 场景3：处理多GPU训练保存的模型（带module.前缀），在单GPU或CPU上加载 state_dict = torch.load(‘ddp_model.pth‘, map_location=device) # 如果键名有‘module.‘前缀，而当前模型没有，需要手动去除 new_state_dict = {k.replace(‘module.‘, ‘‘): v for k, v in state_dict.items()} model.load_state_dict(new_state_dict)

4.2 模型结构不完全匹配与迁移学习

你有一个在ImageNet上预训练好的ResNet模型，想用它来做自己特定的医学图像分类任务。你的模型在ResNet基础上，修改了最后的全连接层（分类头）。这时直接加载就会报错，因为最后的fc.weight形状不匹配。

解决方案：使用load_state_dict()的strict=False参数，并手动处理不匹配的键。

import torchvision.models as models # 1. 加载预训练权重 pretrained_dict = torch.load(‘pretrained_resnet.pth‘) # 2. 创建你的模型（例如，ResNet50，但分类头输出类别数改为10） model = models.resnet50(pretrained=False) # 先不加载官方预训练 num_ftrs = model.fc.in_features model.fc = nn.Linear(num_ftrs, 10) # 修改分类头 # 3. 获取当前模型的state_dict model_dict = model.state_dict() # 4. 筛选预训练字典中，键名在当前模型中也存在的部分 # 这步会过滤掉因为修改fc层而不匹配的键 pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and v.size() == model_dict[k].size()} # 5. 更新当前模型的字典 model_dict.update(pretrained_dict) # 6. 加载（strict=False允许不匹配的键存在） model.load_state_dict(model_dict, strict=False) print(‘成功加载了预训练层，并初始化了新的分类头。‘)

4.3 保存与加载自定义复杂对象

有时，我们想保存的不仅仅是模型和优化器，还可能包括数据集、词汇表等自定义类的实例。只要这些对象可以被Python的pickle模块序列化，torch.save()就能处理。

class MyDataset: def __init__(self, data): self.data = data self.vocab = {‘hello‘: 0, ‘world‘: 1} # 一个简单的自定义属性 dataset = MyDataset(torch.randn(100, 10)) checkpoint = { ‘model_state‘: model.state_dict(), ‘dataset‘: dataset # 保存整个自定义对象 } torch.save(checkpoint, ‘full_checkpoint.pth‘) # 加载时，自定义对象也会被恢复 loaded_checkpoint = torch.load(‘full_checkpoint.pth‘) loaded_dataset = loaded_checkpoint[‘dataset‘] print(loaded_dataset.vocab) # 输出: {‘hello‘: 0, ‘world‘: 1}

注意事项：自定义类必须定义在模块的顶层（不能被嵌套在其他函数或类中），并且其源代码在加载时必须可用（即可以被导入），否则pickle可能无法正确反序列化。对于非常复杂的或包含外部资源（如文件句柄）的对象，更稳妥的做法是只保存其核心数据（如vocab字典），在加载时用这些数据重新构造对象。

5. 工程化最佳实践与性能优化

当项目从个人实验走向团队协作或生产部署时，保存与加载的规范性和效率就变得尤为重要。

5.1 文件格式与序列化性能

PyTorch默认使用Python的pickle协议。从PyTorch 1.6开始，引入了基于zipfile的新的存储格式（通过torch.save(..., _use_new_zipfile_serialization=True)），它支持更高效的存储和随机访问，尤其是在处理大型张量时。在较新版本中，这通常是默认或推荐行为。

# 显式使用新的zipfile序列化（PyTorch 1.6+） torch.save(model.state_dict(), ‘model_new_format.pth‘, _use_new_zipfile_serialization=True) # 加载方式不变 state_dict = torch.load(‘model_new_format.pth‘)

对于超大规模的模型（如数十GB的LLM），直接使用torch.save/load可能内存压力巨大。此时可以考虑：

1. 分片保存：将state_dict按层或按模块拆分成多个文件保存和加载。
2. 使用流式加载：对于非常大的单个Tensor，可以结合torch.load的map_location参数和mmap（内存映射）技术进行部分加载。
3. 专用格式：考虑转换为如ONNX、TorchScript(torch.jit.save) 或使用torch.save配合pickle协议4+，这些格式可能在特定部署场景下更高效。

5.2 版本控制与兼容性管理

模型文件应该被纳入版本控制系统（如Git LFS）进行管理。同时，强烈建议在检查点中保存元数据。

checkpoint = { ‘model_state_dict‘: model.state_dict(), ‘metadata‘: { ‘pytorch_version‘: torch.__version__, ‘model_class‘: model.__class__.__name__, ‘model_config‘: {‘input_size‘: 224, ‘num_classes‘: 1000}, # 模型结构配置 ‘creation_time‘: ‘2023-10-27‘, ‘git_commit_hash‘: ‘abc123def‘, # 关联代码版本 ‘performance‘: {‘val_acc‘: 0.945, ‘val_loss‘: 0.12} } } torch.save(checkpoint, ‘model_with_meta.pth‘)

这样，在未来即使代码迭代，也能清晰地知道这个模型文件是在什么环境下、用什么代码、达到什么性能产出的，极大降低了维护成本。

5.3 安全加载与错误处理

永远不要加载来源不明的模型文件。在生产环境中，加载模型时应加入健全性检查。

def safe_load_model(model, checkpoint_path, expected_keys=None, device=‘cpu‘): “““安全加载模型，包含基本检查和错误处理”“” if not os.path.exists(checkpoint_path): raise FileNotFoundError(f“检查点文件不存在: {checkpoint_path}“) try: checkpoint = torch.load(checkpoint_path, map_location=device) except Exception as e: raise IOError(f“加载模型文件失败，文件可能已损坏: {e}“) if ‘model_state_dict‘ not in checkpoint: raise KeyError(“检查点中未找到 ‘model_state_dict‘ 键”) state_dict = checkpoint[‘model_state_dict‘] # 可选：检查关键层是否存在 if expected_keys: missing_keys = [k for k in expected_keys if k not in state_dict] if missing_keys: print(f“警告: 状态字典中缺少预期键: {missing_keys}“) # 非严格模式加载，并记录信息 missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) if missing_keys: print(f“警告: 以下键在模型中缺失，未被加载: {missing_keys}“) if unexpected_keys: print(f“警告: 以下键在模型中没有对应，被忽略: {unexpected_keys}“) print(“模型加载完成。“) return checkpoint.get(‘metadata‘, {}) # 返回元数据 # 使用示例 metadata = safe_load_model(model, ‘trained_model.pth‘, expected_keys=[‘conv1.weight‘, ‘fc.bias‘])

6. 常见问题排查速查表

在实际操作中，你可能会遇到各种各样的问题。下面这个表格汇总了典型问题、可能的原因及解决方案。

掌握这些排查技巧，能让你在遇到问题时快速定位，而不是盲目地重试或搜索。最后，一个最朴素也最重要的建议：在覆盖任何重要模型文件之前，先做备份。无论是通过代码自动备份最好的几个检查点，还是手动复制，这个习惯能挽救你无数个小时的劳动成果。模型保存与加载，这项看似简单的技能，其稳定性和可靠性，是支撑起所有复杂深度学习项目从实验走向落地的基石。