TorchGeo中ResNet/ViT预训练模型在features_only模式下的问题解析

问题背景

在计算机视觉领域，TorchGeo作为一个专注于地理空间数据的PyTorch库，提供了多种预训练模型支持。其中，ResNet和Vision Transformer(ViT)是两种常用的骨干网络架构。TorchGeo允许用户通过设置features_only=True参数来仅提取中间特征，而不使用最后的全连接层分类头。

问题现象

当用户尝试加载Satlas预训练的ResNet152模型并设置features_only=True时，会遇到AssertionError错误。这是由于模型检查点中包含全连接层('fc.weight'和'fc.bias')的参数，而features_only模式下这些参数不会被加载，导致PyTorch的模型加载机制认为存在"意外键"。

技术原理

在PyTorch中，模型参数加载是通过load_state_dict()方法实现的。该方法会检查提供的状态字典与模型架构的匹配程度。默认情况下，任何不匹配的键都会被视为错误。在TorchGeo的实现中，当前代码严格检查所有键都必须匹配，这在features_only模式下会导致问题。

解决方案

正确的处理方式应该是允许忽略全连接层的参数。具体来说，可以修改断言条件，只检查非全连接层的意外键。例如：

assert set(unexpected_keys) <= {'fc.weight', 'fc.bias'}

这种修改既保持了参数加载的严格性，又兼容了features_only模式的使用场景。

影响范围

这个问题不仅影响ResNet152-Satlas预训练模型，还涉及所有基于timm.create_model创建的模型，包括：

• ResNet系列模型
• Swin Transformer系列模型
• Vision Transformer(ViT)系列模型

临时解决方案

在官方修复发布前，用户可以采用以下临时解决方案：

1. 手动将模型的全连接层替换为恒等映射：

model.fc = nn.Identity()

2. 修改本地TorchGeo源码中的断言条件

最佳实践建议

对于需要使用features_only模式的用户，建议：

1. 明确了解该模式下模型输出的特征图尺寸和通道数
2. 注意不同层级特征图的语义信息差异
3. 考虑特征金字塔网络(FPN)等结构来融合多尺度特征
4. 在微调时，根据下游任务调整特征提取的层级

总结

这个问题揭示了深度学习模型设计中接口一致性的重要性。TorchGeo团队已经意识到这个问题，并将在后续版本中修复。对于地理空间分析任务，正确使用预训练模型的特征提取能力可以显著提升模型性能，特别是在数据量有限的情况下。理解并正确处理这类技术细节，是构建高效地理空间分析系统的关键。