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

在PyTorch-Image-Models（timm）库中使用Vision Transformer模型时，开发者经常需要根据实际任务调整分类头的数量。本文将以vit_intern300m_patch14_448模型为例，深入探讨如何正确修改分类头数量而不影响模型加载。

问题现象

当开发者尝试通过create_model函数创建模型并指定num_classes=3参数时，会遇到如下错误：

RuntimeError: Error(s) in loading state_dict for VisionTransformer:
Missing key(s) in state_dict: "head.weight", "head.bias"

原因分析

这个问题的根源在于模型加载机制。当直接使用checkpoint_path参数加载预训练权重时，timm库会严格匹配检查点中的模型结构。原始vit_intern300m_patch14_448模型的分类头维度与用户指定的3类不匹配，导致无法直接加载权重。

解决方案

方法一：使用pretrained参数

更推荐的做法是使用pretrained=True参数，让timm库自动处理模型适配：

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

这种方式下，timm会自动：

1. 下载预训练权重
2. 根据指定的分类数量调整模型结构
3. 智能地初始化新的分类头

方法二：pretrained_cfg_overlay

对于高级用户，可以使用pretrained_cfg_overlay参数更精细地控制模型配置：

model = create_model(
    'vit_intern300m_patch14_448',
    checkpoint_path="/path/to/model.safetensors",
    pretrained_cfg_overlay={'num_classes': 3}
)

技术原理

timm库的模型加载分为两个阶段：

1. 模型创建阶段：根据参数构建模型结构
2. 权重加载阶段：将预训练权重加载到匹配的结构中

当直接使用checkpoint_path时，库会跳过模型适配逻辑，直接尝试严格匹配权重。而使用pretrained参数时，会触发完整的适配流程，包括分类头的调整。

最佳实践

1. 优先使用pretrained=True而非直接指定checkpoint_path
2. 如需修改分类数量，确保使用正确的参数组合
3. 对于自定义模型路径，考虑先创建适配后的模型，再单独加载权重

通过理解这些机制，开发者可以更灵活地在timm库中使用各种Vision Transformer模型，同时根据任务需求调整模型结构。