PyTorch-Image-Models中自定义分类头数量的注意事项

在PyTorch-Image-Models（简称timm）库中使用Vision Transformer模型时，用户可能会遇到一个常见问题：当尝试修改预训练模型的分类头数量时，会出现状态字典加载错误。本文将深入分析这一问题的原因，并提供解决方案。

问题现象

当用户尝试通过create_model函数创建vit_intern300m_patch14_448模型并指定num_classes=3时，系统会抛出RuntimeError，提示状态字典中缺少"head.weight"和"head.bias"这两个关键参数。然而，如果不指定num_classes参数，模型却能正常加载。

原因分析

这一现象的根本原因在于模型加载机制的设计差异：

1. checkpoint_path参数：当使用此参数时，timm库会严格按照检查点文件中的状态字典来加载模型，不做任何适配性修改。这意味着模型结构必须与检查点完全一致，包括分类头的维度。

2. pretrained参数：当使用此参数时，timm库会执行更智能的模型适配逻辑。它会自动处理分类头维度的变化，即使预训练检查点的分类数量与当前需求不同。

解决方案

针对这一问题，有以下几种解决方法：

1. 使用pretrained参数替代checkpoint_path：

model = create_model('vit_intern300m_patch14_448', pretrained=True, num_classes=3)

2. 使用pretrained_cfg_overlay进行配置覆盖：

model = create_model('vit_intern300m_patch14_448', 
                    checkpoint_path=checkpoint_path,
                    pretrained_cfg_overlay={'num_classes': 3})

3. 手动修改分类头（不推荐）：

model = create_model('vit_intern300m_patch14_448', checkpoint_path=checkpoint_path)
model.reset_classifier(3)  # 修改分类头维度

最佳实践建议

1. 对于大多数迁移学习场景，推荐使用pretrained=True参数而非直接指定检查点路径。

2. 如果需要精确控制模型加载行为，可以结合使用pretrained_cfg和pretrained_cfg_overlay参数。

3. 在修改分类头数量后，通常需要重新训练模型，因为预训练的分类头权重已经不再适用。

4. 对于大型Transformer模型，修改分类头后可能需要调整学习率等训练参数。

通过理解timm库的模型加载机制，用户可以更灵活地应用预训练模型到各种计算机视觉任务中，同时避免常见的配置错误。