PyTorch预训练模型加载全攻略

PyTorch加载预训练模型是深度学习中的常见操作，本文将详细介绍在PyTorch中加载预训练模型的各种方法，助力开发者快速上手。针对ResNet、VGG等常见模型，推荐使用torchvision.models直接加载预训练版本，并通过weights参数指定权重类型，方便快捷。对于自定义模型，则需手动加载权重，在创建模型后使用load_state_dict函数，并设置strict=False以忽略不匹配部分，灵活应对模型结构的差异。此外，文章还将讲解微调模型时如何冻结部分层，以及推荐使用模型状态字典进行保存和加载，以提高灵活性并避免设备不一致问题。掌握这些关键点，能够有效解决PyTorch中加载预训练模型的各种问题，提升开发效率。

在PyTorch中加载预训练模型的关键在于理解模型结构与权重的匹配关系。1. 对于ResNet、VGG等常见模型，可通过torchvision.models直接加载预训练版本，使用weights参数指定权重类型更推荐；2. 自定义模型需手动加载权重，创建模型后使用load_state_dict并设置strict=False忽略不匹配部分；3. 微调时可冻结部分层，通过设置requires_grad=False锁定参数；4. 推荐仅保存和加载模型状态字典以提高灵活性，并注意设备一致性问题。掌握这些要点即可应对大多数情况。

在PyTorch中加载预训练模型其实并不复杂，关键在于理解模型结构和权重之间的关系。如果你用的是常见的模型架构（比如ResNet、VGG等），通常可以直接从torchvision.models中获取对应的预训练版本。而对于自定义模型或部分修改后的模型，就需要更细致地处理权重加载方式。

使用torchvision快速加载常见预训练模型

对于像ResNet、AlexNet这类经典模型，推荐使用torchvision.models中的接口。这些模型默认提供是否加载预训练权重的选项。

例如：

import torchvision.models as models

model = models.resnet18(pretrained=True)

这样一行代码就能加载一个带预训练权重的ResNet18模型。如果你想换成其他变种，比如resnet50或者vgg16，只需要替换名字即可。

需要注意的是：从 PyTorch 1.8 开始，pretrained参数被逐步弃用，取而代之的是通过weights参数指定具体权重类型。比如：

model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)

这种方式更清晰，也方便你选择不同来源的预训练权重。

自定义模型加载预训练权重

如果你对模型结构做了改动，比如改了最后一层分类器的输出维度，就不能直接使用pretrained=True，而是需要手动加载权重。

一般做法是先创建你的模型结构，然后加载原始权重文件，并忽略不匹配的部分。

步骤如下：

• 创建模型结构
• 加载预训练权重字典
• 使用load_state_dict并设置strict=False

示例代码：

model = MyModifiedResNet()
pretrained_dict = torch.load('resnet18.pth')
model.load_state_dict(pretrained_dict, strict=False)

这里的关键是确保大部分层的名字和结构都能对应上，否则会报错。如果模型结构差异太大，可能需要手动映射某些层。

微调时冻结部分层

微调模型时，有时候我们只想训练最后几层，而保持前面的特征提取层不变。这时候可以通过设置requires_grad来冻结部分参数。

以ResNet为例，假设你想冻结所有层，只训练最后的全连接层：

for param in model.parameters():
    param.requires_grad = False

# 只训练最后的fc层
model.fc = nn.Linear(512, new_num_classes)

这样，在反向传播过程中，只有model.fc的参数会被更新。

当然，也可以有选择地冻结某些层，而不是全部。比如只冻结前几层卷积层，保留后面层可训练。

权重保存与加载格式注意事项

PyTorch模型保存的常见方式有两种：

• 保存整个模型：torch.save(model, PATH)
• 仅保存模型状态字典：torch.save(model.state_dict(), PATH)

推荐使用第二种方式，因为它更灵活，特别是在模型类定义发生变化后仍然能加载权重。

加载时也要注意设备问题。如果保存的是GPU上的模型，而在CPU环境下加载，需要加上map_location='cpu'：

model.load_state_dict(torch.load(PATH, map_location='cpu'))

否则可能会遇到设备不匹配的问题。

基本上就这些操作了。加载预训练模型本身不难，但要注意结构匹配、设备一致性和是否需要冻结部分层。掌握这几个点，大多数情况都能应对。

好了，本文到此结束，带大家了解了《PyTorch预训练模型加载全攻略》，希望本文对你有所帮助！关注golang学习网公众号，给大家分享更多科技周边知识！