### PyTorch 中加载模型的方法 在 PyTorch 中，可以通过两种主要方式来保存和加载模型：一种是保存整个模型（包括架构和参数），另一种是仅保存 `state_dict`（只包含模型的权重）。以下是详细的说明： #### 1. 加载整个模型 当保存的是整个模型时，可以直接通过 `torch.load()` 方法将其加载回来。这种方式的优点在于简单易用，缺点则是灵活性较低，因为一旦保存了完整的模型，就很难对其进行修改。 ```python import torch from torchvision.models import resnet18 # 假设我们之前已经保存了一个 ResNet-18 模型 model = torch.load('model.pth') # 加载整个模型 model.eval() # 设置为评估模式 ``` 上述代码展示了如何加载一个预先保存好的完整模型[^1]。 --- #### 2. 加载模型的 `state_dict` 如果只是保存了模型的状态字典 (`state_dict`)，那么需要先重新定义模型的架构，然后再加载状态字典到该模型中。这种方法更加灵活，允许对模型进行调整或扩展。 ##### 2.1 定义模型架构并加载 `state_dict` 假设有一个简单的神经网络模型如下所示： ```python import torch import torch.nn as nn class SimpleModel(nn.Module): def __init__(self): super(SimpleModel, self).__init__() self.fc1 = nn.Linear(784, 128) self.relu = nn.ReLU() self.fc2 = nn.Linear(128, 10) def forward(self, x): x = self.fc1(x) x = self.relu(x) x = self.fc2(x) return x # 创建模型实例 model = SimpleModel() # 假设已保存 state_dict 到文件 'model_state_dict.pth' state_dict = torch.load('model_state_dict.pth') # 将 state_dict 载入模型 model.load_state_dict(state_dict) model.eval() # 设置为评估模式 ``` 这段代码演示了如何加载模型的 `state_dict` 并应用到自定义模型上[^2]。 --- #### 3. 只加载部分参数 有时可能只需要加载模型的部分参数而不是全部。这通常用于迁移学习场景下，其中某些层被冻结而其他层则继续训练。实现方法是对齐键名并将特定层的参数加载进去。 ```python pretrained_dict = torch.load('partial_model_state_dict.pth') model_dict = model.state_dict() # 过滤掉不需要的参数 pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} # 更新现有模型的参数 model_dict.update(pretrained_dict) # 加载更新后的参数 model.load_state_dict(model_dict) ``` 此代码片段展示了一种策略，即从预训练模型中提取所需的关键参数，并将其应用于目标模型中。 --- ### 总结 PyTorch 提供了多种保存和加载模型的方式，具体选择取决于实际需求。对于快速部署或者不需更改的情况可以选择保存整个模型；而对于更复杂的任务比如微调，则推荐使用 `state_dict` 的方式进行操作。