timm 升级指南：从 0.6.x 到 1.x 的模型权重重映射、模块路径迁移与新命名规范

本文为 UPGRADING.md 的扩充解读，面向需要从 timm 0.6.x 旧版本迁移到当前 1.x 版本（本仓库 timm/version.py 显示为 1.0.29.dev0）的开发者。读完本文，你能掌握三件事：checkpoint_filter_fn 权重重映射机制的自动/手动两种调用方式、0.6.x 以来全部 API 变更点（模块迁移、下划线前缀、architecture.pretrained_tag 新命名）及其对应源码位置，以及 Hugging Face Hub 成为权重主源后的下载解析逻辑，从而能安全地复现旧 checkpoint 的加载。

一、timm 的版本兼容策略：尽量保住接口，例外靠映射兜底

UPGRADING.md 开篇明确了维护者的原则：跨版本尽量保持代码接口、尤其是模型权重（weights）的兼容性，但存在例外。这意味着升级 timm 后，两类资产可能失效：

1. 代码层面：从 timm.models.layers、timm.models.helpers 等旧路径导入的符号；
2. 权重层面：旧版本保存的 checkpoint，因模型结构微调（例如为支持 features_only=True 而调整架构）导致 key 或形状对不上。

下面的章节分别给出对应的迁移手段与源码证据。

二、核心机制：checkpoint_filter_fn 权重重映射

2.1 它做什么

checkpoint_filter_fn 定义在各个模型实现模块中，承担两类职责（引自 UPGRADING.md）：

• 将旧版 timm 预训练 checkpoint 重映射到新架构的 key 格式；
• 将第三方（原始作者发布）checkpoint 转换为 timm 格式——即模型引入 timm 时被改造过结构的情况。

2.2 自动调用链：pretrained=True 时如何生效

从源码结构看，自动调用发生在模型构建辅助函数中。build_model_with_cfg 在 pretrained=True 时调用 load_pretrained，并把该模型注册的 pretrained_filter_fn 作为 filter_fn 传入（timm/models/_builder.py#L458-L467）。load_pretrained 内部对 filter 的调用做了向后兼容的签名探测（timm/models/_builder.py#L240-L245）：

if filter_fn is not None:
    try:
        state_dict = filter_fn(state_dict, model)
    except TypeError as e:
        # for backwards compat with filter fn that take one arg
        state_dict = filter_fn(state_dict)

也就是说：新版 filter 函数签名是 filter_fn(state_dict, model)，而 0.6.x 时代只接收 state_dict 单参数的旧 filter 也能继续工作——这正是文档所说“pretrained=True 时自动调用”的底层保证。

2.3 手动调用：把你自己的旧 checkpoint 重映射到新模型

当你加载的是自己保存的旧格式 checkpoint（而非走 pretrained=True 下载），需要手动传入 filter。文档给出的方式是把 checkpoint_filter_fn 传给 timm.models.load_checkpoint(..., filter_fn=...)。load_checkpoint 的完整签名与行为在 timm/models/_helpers.py#L175-L218：

def load_checkpoint(
        model,            # 目标模型
        checkpoint_path,  # checkpoint 文件路径
        use_ema=True,     # 若 checkpoint 中存在 EMA 权重则优先使用
        device='cpu',
        strict=True,
        remap=False,      # 按参数注册顺序整体重排 key（与 filter_fn 二选一）
        filter_fn=None,   # 自定义过滤/重映射函数，签名 (state_dict, model)
        weights_only=True,
) -> Any:

以文档示例中提到的 Swin Transformer V2 为例（UPGRADING.md 原文即以此为 filter_fn=timm.models.swin_transformer_v2.checkpoint_filter_fn 的例子）：

import torch
import timm

model = timm.create_model('swin_v2_base', pretrained=False)

# 方式一：通过 timm.models.load_checkpoint 传入 filter_fn
timm.models.load_checkpoint(
    model,
    'old_swin_v2_base.pt',
    filter_fn=timm.models.swin_transformer_v2.checkpoint_filter_fn,
)

# 方式二：直接对 state dict 应用 filter 后手动 load
state_dict = torch.load('old_swin_v2_base.pt', map_location='cpu')
state_dict = timm.models.swin_transformer_v2.checkpoint_filter_fn(state_dict, model)
model.load_state_dict(state_dict, strict=True)

2.4 一个真实的 checkpoint_filter_fn 长什么样

timm/models/swin_transformer_v2.py#L1075-L1114 是一个可对照学习的完整实现，它依次做了四件事：

1. 解包嵌套结构：兼容 checkpoint 外层包裹 model / state_dict 两种嵌套（L1088-L1089）；
2. 剔除非持久化 buffer：relative_position_index、relative_coords_table、attn_mask 由模型自己重算，直接从 state dict 中跳过（L1094-L1095）；
3. patch embedding 尺寸重采样：当目标模型的 patch_embed.proj 尺寸与 checkpoint 不同时，用双三次插值（带抗锯齿）重采样卷积核（L1097-L1106）；
4. key 重映射：把旧命名 layers.{i}.downsample 偏移为 layers.{i+1}.downsample，并把 head. 替换为 head.fc.（L1108-L1111）——这就是文档所说的“old pretrained checkpoints to new”的具体含义。

仓库中约有 40 个模型模块定义了各自的 checkpoint_filter_fn（如 vision_transformer.py、convnext.py、deit.py 等），升级时只需查看目标模型所在模块是否提供了该函数。

2.5 配套工具：load_state_dict、clean_state_dict、remap_state_dict

升级过程中加载旧 checkpoint，常会用到 timm/models/_helpers.py 中的三个辅助函数（均已从 timm/models/init.py#L129-L130 导出，可直接 timm.models.load_checkpoint 使用）：

函数	位置	作用
load_state_dict	timm/models/_helpers.py#L132-L172	从文件加载权重；支持 .safetensors；按 state_dict_ema → model_ema → state_dict → model 的优先级取嵌套 key
clean_state_dict	timm/models/_helpers.py#L118-L129	去除 DDP 的 module. 前缀与 torch.compile 的 _orig_mod. 前缀
remap_state_dict	timm/models/_helpers.py#L221-L247	按参数注册顺序忽略原 key 整体重排（要求两个模型按相同顺序注册参数），形状不同时可自动 reshape

一个典型的旧 checkpoint 加载组合（例如分布式训练保存、带 EMA 的 timm 训练 checkpoint）：

import torch
import timm

ckpt = torch.load('model-old.pt', map_location='cpu', weights_only=False)
state_dict = timm.models.clean_state_dict(ckpt['state_dict'])  # 去掉 module. 前缀
model = timm.create_model('resnet50', pretrained=False)
model.load_state_dict(state_dict)

checkpoint 加载相关的自动化测试可参考 tests/test_checkpoint_loading.py。

三、从 0.6.x 升级：全部变更点逐项核对

UPGRADING.md 指出：0.6.x 稳定版之后的变化曾在 0.8.x 开发版预览，但并非所有人都完成了迁移。以下是逐项变更与对应的当前仓库实现证据。

3.1 timm.models.layers 迁移到 timm.layers

• 新写法：from timm.layers import DropPath, ConvNormAct, ...
• 旧写法仍可用但有期限：from timm.models.layers import name 通过弃用映射继续工作，但会触发警告；
• 必须修改的写法：import timm.models.layers.module 或 from timm.models.layers.module import name 这种按子模块导入的形式已不可用，必须改为 from timm.layers.<module> import name。

源码证据：timm/models/layers/init.py 本身就是一个转发垫片——第 1 行注释写明 “timm.models.layers is DEPRECATED, please use timm.layers”，它把 timm/layers/ 下的全部公共符号重新导出，并在 第 48-49 行 发出 FutureWarning：

warnings.warn(f"Importing from {__name__} is deprecated, please import via timm.layers", FutureWarning)

因此升级后的第一步操作是：全局搜索代码中的 from timm.models.layers，统一替换为 from timm.layers，消除警告。

3.2 timm.models 下的 builder/helper 模块加上 _ 前缀

原先公开命名的工具模块现在以 _ 前缀表示“内部实现”，文档原文示例为 timm.models.helpers → timm.models._helpers。当前仓库中这些文件是：

旧路径（0.6.x）	新路径（1.x）	主要职责
timm.models.helpers	timm/models/_helpers.py	load_checkpoint、load_state_dict、clean_state_dict、remap_state_dict
timm.models.builder	timm/models/_builder.py	build_model_with_cfg、load_pretrained、resolve_pretrained_cfg
timm.models.registry	timm/models/_registry.py	模型注册与 get_pretrained_cfg 查询
timm.models.manipulate	timm/models/_manipulate.py	输入卷积通道适配等

文档同时提醒：存在临时弃用映射文件，后续会移除。因此直接依赖这些内部路径的代码应尽快改造——对外部使用者而言，推荐的公开入口是 timm.create_model(...) 工厂函数和 timm.models.load_checkpoint(...)，而不是直接导入 builder 内部符号。

3.3 新命名规范：architecture.pretrained_tag

这是 0.6.x 以来最大的 API 变化。所有模型现在支持 architecture.pretrained_tag 形式的命名：

• architecture（架构）：模型结构名，例如 resnet50、vit_base_patch16_224、swin_v2_base；
• pretrained_tag（权重变体标签）：同一架构下不同的预训练权重来源/配置，例如文档示例 resnet50.rsb_a1 中的 rsb_a1；
• 省略 tag 时的默认行为：只写 architecture 会取该架构 default_cfgs 中登记的第一组权重；
• 旧型号名重命名：过去靠完整型号名区分的权重被并入 tag，例如 vit_base_patch16_224_in21k → vit_base_patch16_224.augreg_in21k，旧名仍可通过弃用映射解析。

源码中，tag 的解析发生在 resolve_pretrained_cfg：当传入字符串 tag 时，将其拼成 variant.tag 再查询预训练配置注册表（L366-L370）；查不到 tag 时会回退并打印警告。timm.create_model('resnet50.rsb_a1', pretrained=True) 的调用形式即由此支持。

3.4 为支持 features_only=True 而重映射的 checkpoint

文档说明：一批模型（含 Swin Transformer V2）为更好地支持 features_only=True 提取中间特征，调整了架构，相应地在其模型模块内提供了 checkpoint_filter_fn。对你已保存的旧 checkpoint，只需按第二节的方式把对应模块的 filter 传给 load_checkpoint 即可完成重映射。哪些模块含有该函数，可用 grep -l "def checkpoint_filter_fn" timm/models/*.py 快速确认（当前仓库约 40 个模型模块包含该函数）。

3.5 权重主源迁移到 Hugging Face Hub

UPGRADING.md 声明：Hugging Face Hub 已成为 timm 权重的主源，模型卡中附带论文、原始出处与 license 信息。

源码侧的解析逻辑在 _resolve_pretrained_source：对每个预训练配置，按以下优先级决定下载源——

1. 配置显式指定 source='hf-hub' → 走 Hugging Face Hub 的 hf_hub_id；
2. 显式 source='local-dir' → 从本地目录加载；
3. 默认源（timm）下，按 state_dict（直接传入的权重字典，最高优先级）→ file（本地文件覆盖）→ 旧缓存（需环境变量 TIMM_USE_OLD_CACHE=1）→ hf_hub_id → 传统 url 的顺序解析（L62-L83）。

对升级者的实际影响：0.6.x 时代通过 URL 缓存的旧权重不会自动替换，pretrained=True 将优先从 Hub 拉取新格式的权重（新权重通常已配合 checkpoint_filter_fn 适配好，无需再手动重映射）。

3.6 旧版本的保留方式

0.6.x 旧版仍可通过项目仓库的 0.6.x 分支获取源码，也可以直接用 pip 指定版本号安装旧版。若你的下游项目暂时无法迁移，锁定 0.6.x 版本是最简单的过渡方案；但文档已明确 timm.models.layers 的临时映射“会被移除”，长期应完成迁移。

四、迁移实操清单

按顺序执行以下检查，可完成从 0.6.x 到 1.x 的平滑升级：

1. 升级安装：pip install -U timm（本仓库当前开发版本为 1.0.29.dev0，见 timm/version.py）；
2. 替换 import：将代码中所有 from timm.models.layers import X 改为 from timm.layers import X；删除对 timm.models.layers.<module>、timm.models.helpers 等旧子模块路径的直接导入，改用公开 API（timm.create_model、timm.models.load_checkpoint）；
3. 核对模型名：运行 python -c "import timm; print(timm.list_models('*swin*'))" 之类的方式确认目标型号是否已改用 architecture.tag 形式，旧全名若报 “No pretrained configuration” 警告，切换到新 tag 名；
4. 旧 checkpoint 重映射：
   • 用 timm.models.load_state_dict 加载（自动处理 EMA/嵌套 key）；
   • 需要时先 timm.models.clean_state_dict 去 module. 前缀；
   • 目标模型模块若提供 checkpoint_filter_fn，通过 filter_fn= 传入 timm.models.load_checkpoint 完成 key 重映射；
   • 加载后检查 missing_keys / unexpected_keys 日志（load_pretrained 会对二者分别打印 info/warning，见 timm/models/_builder.py#L283-L290）；
5. 验证精度：用 validate.py 对旧任务做一次评估，确认权重加载无误后再投入训练。

五、升级相关文档与源码索引

内容	位置
升级文档原文	UPGRADING.md
权重加载与 filter 调用链	timm/models/_builder.py
load_checkpoint / clean_state_dict / remap_state_dict	timm/models/_helpers.py
timm.models.layers 弃用垫片	timm/models/layers/init.py
filter 实现示例	timm/models/swin_transformer_v2.py
预训练配置（tag）解析	timm/models/_builder.py
checkpoint 加载测试	tests/test_checkpoint_loading.py

总结：timm 1.x 的升级成本集中在“导入路径”与“型号命名”两处，而权重兼容由 checkpoint_filter_fn + load_checkpoint(filter_fn=...) 这套机制兜底。只要按本文清单逐项替换 import、核对型号 tag、并在加载旧 checkpoint 时挂上对应模块的 filter，旧版本的权重资产可以在新架构下完整复用。