★ FEATURED ARTICLE 我的入坑经历两年前我第一次在一个图像分类项目里做预训练模型选型当时的做法还是去各仓库翻权重文件下载之后还要自己对着源码找预处理参数、看归一化均值方差、试输入尺寸光是让一个模型跑通前向就花了大半天。后来同事甩给我一句“直接用 PyTorch Image Models”也就是现在大家口里的 timm我才意识到原来视觉模型的获取、使用、微调可以这么省心。timm 本质上是一个 PyTorch 生态下的视觉模型库官方说法是 image models包含了大量预训练模型、训练脚本和图像变换工具。它最大的价值不是某一个模型很强而是把几百个模型架构、成千上万个预训练权重、以及训练时必备的数据增强和调度策略全部统一到了一套 API 下面。平时做分类、做迁移学习、做特征提取甚至做对比实验timm 都可以直接当“标准工具箱”来用。这篇文章会围绕实际项目中使用 timm 的真实过程把模型加载、权重管理、微调改造、数据增强、训练加速和部署导出这些环节逐一拆开来说。如果你刚接触视觉模型会被各种仓库的琐碎细节劝退timm 能帮你省掉大量“适配”工作如果你已经用 torchvision 或者手写训练脚本timm 里的增强和训练配套工具也值得借鉴如果你要快速验证多个模型的效果timm 统一接口带来的便利更是立竿见影1. timm 到底帮你省掉了哪些麻烦很多人第一眼看到 timm 的模型列表会有点懵HuggingFace 上动辄上千个模型卡片各种后缀和变体不知道从哪下手。我的建议是先别管全貌理解 timm 在设计上解决的三类问题后面所有的用法都会自然串起来。1.1 统一模型接口带来的巨大便利在没有 timm 的日子里你想换一个模型意味着换一套加载代码。不同作者的实现风格千差万别有的构造函数要传num_classes有的要传pretrained有的还要求你先实例化一个 backbone 再套 head。项目一旦做大了这种碎片化很折磨人。timm 把这件事收敛成了非常少的核心调用import timm # 一行代码创建模型 model timm.create_model(resnet50, pretrainedTrue, num_classes10) # 一行代码查看模型默认输入尺寸 print(model.default_cfg[input_size]) # (3, 224, 224)timm.create_model的背后逻辑可以理解成一个模型工厂。你传字符串给model_name库内部去查注册表找到对应的构造函数然后根据你的pretrained、num_classes、dropout、global_pool等参数把模型拼好。这个过程对你完全透明也不需要关心不同模型架构的命名差异。这样做的好处是实验脚本可以做得非常干净。我一般会把模型名做成配置项跑实验的时候像传参一样切换结构python train.py --model convnext_tiny --epochs 30 --lr 3e-4 python train.py --model efficientnet_b0 --epochs 30 --lr 3e-4 python train.py --model vit_base_patch16_224 --epochs 30 --lr 3e-5换模型不需要改训练代码的逻辑只需要改参数这对对比实验和快速迭代来说收益极大。1.2 权重管理很少再为了下载模型发愁预训练权重是另一个痛点。很多开源仓库把权重放在自己的服务器上下载速度不稳定版本迭代后哈希也会变甚至有些权重下下来才发现框架版本不匹配。timm 的权重管理延续了模型工厂的设计思路首次使用某个预训练模型时它会尝试从指定的权重地址下载默认落在一个缓存目录里。这个缓存路径可以通过环境变量控制。比如我想把权重放到当前项目的pretrained_weights目录下export TIMM_HUB_ROOT/data/weights/timm在分布式或者多个机器共用数据盘的场景里这个设置特别实用。只要把缓存目录共享出去其他机器第一次加载权重时就能命中缓存不再重复下载十几 GB 的模型。如果你所在环境无法直接访问外部下载源timm 也支持手动放置权重文件。你只需要把事先下载好的权重放到缓存目录下并且文件名和它预期的保持一致。这个一致性检查不用担心运行时会打印出实际查找路径按日志来放就行。我还在部分实验环境里用过离线方式配合timm.list_models()先看有哪些模型可用再针对性准备权重。提示生产环境建议用固定的 timm 版本并锁住权重缓存目录。不同版本之间模型注册表会变化不锁版本容易遇到权重文件级联更新的情况。2. create_model 的关键参数与应用场景create_model看起来很简单但实际用下来有不少细节决定了你后续工作顺不顺利。我见过很多人在这一步踩了坑然后误以为是模型本身出了问题。2.1 必须吃透的基础参数常用参数看起来不多但每个背后都有实际用途我整理了一份实际使用中形成的速查表。参数作用实际建议model_name模型结构字符串标记如resnet50、vit_base_patch16_224用timm.list_models(*关键字*)搜索具体有哪些可用模型pretrained是否加载预训练权重微调默认True特征提取时也通常为Truenum_classes分类头输出维度需要做自定义分类时务必显式指定不要让模型默认用 1000global_pool特征图池化方式可选avg、max、avgmax、token等默认avgViT 这类带 token 时可以选token或avgdrop_rate线性分类头前的 dropout 比例数据量少时可以适当调高默认通常是 0.0pretrained_cfg权重来源配置可覆盖默认权重地址对加载权重地址有特殊要求时使用in_chans输入通道数默认 3做 1 通道灰度图或者 5 通道多光谱时需要改features_only只返回中间层特征不包含分类头做分割、检测、特征提取时需要features_onlyTrue是一个非常高性价比的参数。比如你想拿某模型做 backbone希望把特征图提取出来而不关心分类头直接设置即可model timm.create_model( resnet50, pretrainedTrue, features_onlyTrue, out_indices[2, 3, 4], ) # 得到的 model 会返回一个特征列表每个元素对应一个 stage在医学影像、遥感图像这类自定义输入通道场景里设置in_chans也值得留意。比如输入只有单通道权重是预训练的三通道模型timm 会做首层通道数的自适应通常是把第一层卷积的权重做均值化处理。这样你不需要自己提前处理权重文件。2.2 通过 default_cfg 查看模型默认配置很多人忽略了model.default_cfg这个字段但其实它记录了模型训练时的很多信息输入尺寸、归一化均值方差、插值方式、是否包含 token、原始模型来源等。cfg model.default_cfg print(cfg[mean]) # 归一化均值比如 [0.485, 0.456, 0.406] print(cfg[std]) # 归一化方差 print(cfg[input_size]) # 模型的默认输入尺寸 print(cfg[interpolation]) # 预测时数据预处理可以完全复用这份配置像 ViT 这类模型输入尺寸通常不是 224 就是 384归一化的 mean/std 也可能因为训练用的数据集不一样而不同。如果这些细节不匹配预训练模型的表现会大打折扣。你在做推理服务时可以直接从default_cfg取配置而不是在代码里写死一套常数。注意不同模型的default_cfg结构基本一致但个别字段会有差异。比如 vision transformer 类模型通常会有crop_pct和num_tokensCNN 类模型则没有。以实际打印结果为准。2.3 模型搜索与命名规律timm 的模型命名有比较强的规律结构名 变体 输入分辨率。例如vit_base_patch16_224表示 ViT Basepatch 大小为 16输入尺寸为 224“_in21k”后缀通常表示在 ImageNet-21k 上预训练的 checkpoint适合进一步微调“_ft_in1k”则表示先在 21k 上预训练又在 1k 上微调过可以直接迁移使用。这听起来是小事但很多新手在选模型时往往只盯着最高的测试准确率忽略了输入尺寸和模型尺寸对部署的影响。我的经验是先做一轮小范围筛选再在项目约束下做二次选择。3. 微调实操改造模型与迁移学习的关键步骤timm 的预训练模型数量庞大但大多数下游任务不是拿来做原始 1000 类分类而是做成自定义分类、特征提取、分割检测的骨干网。这一章节讲的是我在实际微调中高频使用的几个改造方案。3.1 替换分类头的正确方式当我们传入num_classes时timm 已经自动更换了分类层。但如果你需要在更高级别上改动模型结构比如去掉分类头直接拿特征或者把分类头换成一个自定义的多层 MLP就需要直接操作模型的模块。以 ResNet 系为例预训练模型一般会以.fc、.head或.classifier命名分类层ViT 和大部分 Transformer 系则是.head。找到之后替换即可import torch import torch.nn as nn import timm model timm.create_model(resnet50, pretrainedTrue, num_classes0) # num_classes0 表示不要分类头只保留主干 model.fc nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(2048, 512), nn.ReLU(), nn.Dropout(0.3), nn.Linear(512, 10), )对于 ViT 这类模型num_classes0后通常还保留着pre_logits出来的特征向量你可以从这里接一个自定义的 MLP head。model_vit timm.create_model(vit_base_patch16_224, pretrainedTrue, num_classes0) # model_vit 的输出是 patch token 的平均池化特征 model_vit.head nn.Linear(model_vit.embed_dim, 10)这里要提醒一下timm 不同的模型在num_classes0时的返回值并不完全相同。有的返回的是经过全局池化的特征有的是原始的特征状态。实在不确定时用一个虚拟输入跑一遍打印形状成本很低。3.2 冻结主干与梯度分组策略微调时最常见的是“冻结主干只训分类头”或者“主干用较小学习率、头用较大学习率”。前者适合数据量比较少、防止过拟合的场景后者适合目标任务和预训练任务差异较大的情况。model timm.create_model(resnet50, pretrainedTrue, num_classes10) # 冻结主干参数只让分类头参与训练 for name, param in model.named_parameters(): if not name.startswith(fc) and not name.startswith(head): param.requires_grad False但单纯冻结或全部放开并不是最优的。我通常会把参数分成两组完全微调时用一个低学习率跑主干用一个相对较高的学习率跑顶层。这个做法在迁移学习效果上明显好于“统一切换”尤其是在数据量小但目标和预训练任务形态差异较大的项目里。param_groups [ {params: [p for n, p in model.named_parameters() if fc not in n and head not in n], lr: 1e-5}, {params: [p for n, p in model.named_parameters() if fc in n or head in n], lr: 1e-3}, ] optimizer torch.optim.AdamW(param_groups)3.3 输入尺寸调整与位置编码插值很多视觉模型尤其 ViT对输入尺寸是有约束的。vit_base_patch16_224直接喂 384×384 分辨率进去如果不对位置编码做插值程序会直接报错或者特征分布被破坏。PyTorch 的生态里这个问题通常通过插值位置编码来解决timm 也内置了相关支持。部分新版本模型通过pos_embed的 resample 逻辑自动处理输入尺寸变化。旧版本则可能需要手动让输入尺寸必须匹配模型的预设分辨率。PiT、DeiT这类模型对分辨率相对灵活CNN 类模型更是天然可以应对任何尺寸只是需要把全局池化和分类头处理好。我遇到过不少同学以为模型输入越大效果一定越好结果直接把 ViT 的 224 模型改成 448 后遇到各种诡异报错又不知道是位置编码的问题。所以我的建议是输入分辨率调整优先选 CNN 类模型如果一定要用 Transformer 类模型就选本身支持 384 或 448 的变体或者确认位置编码被正确插值处理后再做。3.4 特征提取模式下关闭不需要的模块做特征提取时timm 仍会默认带上一些你并不需要的模块比如训练用的DropPath、分类 head 里的 dropout。虽然它们对前向特征结果没有直接影响但会让输出产生轻微随机性而且浪费一点点计算量。如果只是提取特征不参与训练最稳妥的方式是切到 eval 模式并关闭梯度model.eval() with torch.no_grad(): features model(x)如果你做的是需要保存特征到磁盘的批任务建议固定随机种子并确认每个 batch 的结果可重复。我在多进程提取特征时就遇到过因为DropPath随机性导致前后两次提取结果不一致的问题排查到最后是忘了把模型切到 eval 模式。4. 数据增强与训练管线的进阶配置timm 除了模型外还附带了一套相当实用的训练工具核心是数据增强和训练加速。如果只是用模型不学这些等于浪费了一大半资源。4.1 timm 自带的数据增强组件timm.data里提供了丰富的增强策略最常见的是RandAugment、AutoAugment、Mixup、CutMix等。这些做法在今天看来已经是视觉训练的标准套路问题是怎么把它们组合得当。我习惯把增强分两类几何/颜色类增强和混合类增强。前者包括随机裁剪、翻转、RandAugment后者包括 Mixup 和 CutMix。两类可以叠加但不要过度否则模型训练初期会很难收敛。from timm.data import create_transform from timm.data.transforms import RandAugment img_size 224 train_transform create_transform( input_sizeimg_size, is_trainingTrue, no_augFalse, hflip0.5, vflip0.0, color_jitter0.3, auto_augmentrand-m9-mstd0.5-inc, re_prob0.25, re_modepixel, re_count1, std[0.229, 0.224, 0.225], )上面这段是训练用的增强。auto_augment参数传递了一个策略字符串表示 RandAugment 的参数re_prob是 Random Erasing 的概率。create_transform会根据传入的 is_training 自动在训练和验证模式下切换验证时自动做 resize center crop normalize不用你手动维护两套变换。这其中的关键是别小看Random Erasing。在少量数据场景下随机擦除对提升鲁棒性的帮助非常明显。很多基础教程里不会讲这一点但 timm 把它作为默认选项之一说明它对训练质量的提升是经过验证的。提示如果你的任务是不需要强增强的领域比如一些对称性很强的工业质检图auto_augment建议从较弱策略开始调别上来就开最强的。盲目堆增强反而会让小数据集训练困难。4.2 Mixup 与 CutMix 的正确姿势Mixup 是“按比例融合两张图及其标签”CutMix 是“用一块矩形区域替换另一张图并混合标签”。二者都是缓解过拟合的重要工具timm 中通过timm.data.mixup.Mixup实现。实际使用中我总结出的经验是这样的分类标签要用 soft label而不是 hard label否则 Mixup 就失去了软化边界的意义通常mixup_alpha0.8、cutmix_alpha1.0两者以一定概率同时启用训练前几个 epoch 不要开启 mixup等模型基础收敛后再逐渐引入能减少早期震荡from timm.data.mixup import Mixup mixup_fn Mixup( mixup_alpha0.8, cutmix_alpha1.0, prob1.0, modebatch, label_smoothing0.1, num_classes10, ) samples, targets next(iter(train_loader)) samples, targets mixup_fn(samples, targets)上面的label_smoothing是另一个细节。它本质上的作用是告诉模型“有些标签并不是绝对可信”在噪声标签场景下很有用。Mixup 本身已经产生了软标签如果模型训练时再用 label smoothing相当于双重平滑效果在不同数据集上表现不同我一般会分别跑一版对比。4.3 为什么要顺手用 timm 的 EMA 和调度器timm 内置了ModelEmaV2在做较大模型训练时指数移动平均能让最终模型稳定性明显提升。原因是模型参数在优化后期会在最优解附近震荡EMA 相当于对参数轨迹做平滑尤其适合训练过程噪声比较大的场景。from timm.utils import ModelEmaV2 ema_model ModelEmaV2(model, decay0.9999) # 在每个训练 step 后调用 ema_model.update(model)学习率调度器方面timm 提供了CosineLRScheduler等实现。和 PyTorch 自带的余弦退火相比它额外支持 warmup、cosine 重启等参数设置。我用的比较多的是 warmup 后接 cosine 衰减from timm.scheduler import CosineLRScheduler scheduler CosineLRScheduler( optimizer, t_initialtotal_epochs, warmup_t5, warmup_lr_init1e-6, lr_min1e-6, t_in_epochsFalse, # 如果按 step 调度记得设为 False )如果你只是调模型做对比不追求纪录直接沿用这个组合效率很高。4.4 训练脚本里值得照抄的实用钩子在真实项目中只改模型不动训练流程其实不够。我建议从 timm 的训练脚本里抽几个钩子加到你自己的代码里定期把验证集 top-1/top-5 结果记录下来在 epoch 或 step 维度都能回溯训练中断后能够从 checkpoint 恢复 optimizer、scheduler、EMA 的状态是否使用 AMP 自动混合精度作为开关选项# 训练中检查是否需要存 EMA 版本的模型 def save_checkpoint( model, ema_model, optimizer, scheduler, epoch, filepath ): state { model: model.state_dict(), ema_model: ema_model.state_dict() if ema_model else None, optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict(), epoch: epoch, } torch.save(state, filepath)这些内容本身不大但在你的项目从“实验脚本”变成“可维护训练系统”的过程中会省掉很多返工成本。5. 模型推理与部署时容易忽略的细节训练阶段的模型和部署阶段的模型往往是两回事。很多时候本地跑得欢到了服务化推理却发现精度掉、速度慢。这里整理了我实际部署 timm 模型时踩过的几个关键点。5.1 输入预处理的一致性你训练时用的 mean/std推理时也要保持一致。这个一致性不仅是数值上的还包括图像解码方式、缩放算法、插值方式。timm 的default_cfg里记录了interpolation比如bicubic还是bilinear。同样一张图用不同插值方式缩放模型输出差异可能不小。在工业级服务里客户端传上来的图片质量良莠不齐我建议统一用服务器端解码和预处理不要依赖前端处理后的图片。from timm.data import resolve_data_config from timm.data.transforms import create_transform config resolve_data_config(model.default_cfg) transform create_transform(**config) # 服务端对原始图片做统一预处理 tensor transform(Image.open(sample.jpg).convert(RGB))resolve_data_config会根据模型的 default_cfg 自动读到正确的 size、mean、std、interpolation不需要自己手动维护。5.2 导出 ONNX 或 TorchScript 时注意动态尺寸部署到生产环境时经常需要把模型导出为 ONNX 或者 TorchScript。CNN 类模型的输入尺寸一般比较灵活导出时可以直接指定动态的 H、WViT 和带固定位置编码的模型则需要在导出时固定输入尺寸否则位置编码维度不匹配。还有一个很实际的坑导出前要确保模型处于eval()状态并且关闭梯度。否则导出的模型里会带着一些只在训练阶段使用的图结构导致服务端推理时内存开销大、速度慢甚至结果不一致。model.eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, model.onnx, input_names[pixel_values], output_names[logits], dynamic_axes{pixel_values: {0: batch}}, opset_version17, )5.3 GPU 上推理速度与批处理大小的权衡timm 模型大多预训练表现很好但不同架构在同一 GPU 上的推理速度差异远超过参数量的差距。比如efficientnet_b0和resnet18的准确率接近前者的推理速度可能慢不少而某些 ViT 变体在小 batch 下反而不如小卷积网络表现稳定。我的做法是导出前先做一轮推理性能基线测试# 用 timm 脚本自带的基准工具大致测一下 python benchmark.py --model resnet50 --batch-size 64如果只是分类场景efficientnet和convnext类模型通常比同样精度的传统残差网络更优如果有实时性要求还是老牌resnet、mobilenetv3更稳。6. 做模型对比实验时的各种坑与经验最后说说我把 timm 当作实验基础设施时积累下的经验。这些坑不一定写在文档里但遇到一次就够你折腾半天。6.1 随机种子和可复现性timm.create_model默认权重是从网上下载的下载过程不受随机种子控制这个没问题。但如果你追求实验完全可复现需要注意三处数据加载时的 shuffle 种子、分布式采样器的随机种子模型前向里的随机增强比如RandAugment、Mixup训练时的torch.backends.cudnn.benchmarkimport random import numpy as np import torch def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed)对比实验时一个常犯的错误是只固定了 torch 的种子没有固定数据增强的随机流。你会发现两次实验差距忽大忽小还以为是模型问题其实是增强流程没有稳定下来。6.2 模型微调时 max_norm 和优化器之间的互动timm 模型的预训练权重通常质量很高但微调数据量不足时容易过拟合一个常用对策是梯度裁剪。梯度裁剪会限制单次梯度的最大范数但如果你用 EMA它更新的是参数的移动平均梯度裁剪对 EMA 的影响相对间接。torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0)我也见过有人在微调时用偏大大的学习率觉得有预训练权重兜底。我的经验是微调阶段主干学习率普遍应该低于从头训练一个数量级尤其是 ViT 这类结构超参数错一点就明显感知到收敛变慢。6.3 对同一份数据的多尺度评估验证时很多人直接用训练时的分辨率来评估但模型在default_cfg中给出的分辨率其实未必是它的最优工作分辨率。比如很多 EfficientNet 变体在高一点的分辨率上会获得明显更好的精度只是随之带来更多计算开销。我平时会做一组多尺度验证把输入尺寸分别调到 224、256、288、320 等几档观察准确的上升曲线。这样能直观地知道部署时应该选哪个分辨率for size in [224, 256, 288]: config {**resolve_data_config(model.default_cfg), input_size: (3, size, size)} # 做验证评估这个做法成本低收益却很高。有一次我用一个 224 预训练模型做推理服务只是把服务端缩放到 288几个类别上准确率直接涨了接近两个点而推理耗时只多了不到三成。6.4 权重缓存是空间大户要定期清理timm 的权重缓存目录存放所有下载过的预训练权重网络模型类动辄几百 MBViT Large 之类可能超过 1 GB。如果你经常尝试不同模型这个目录会膨胀得很快。建议养成两个习惯重要项目里把需要的权重手动放到共享目录避免每个机器各下载一份定期用du -sh $TIMM_HUB_ROOT看一下缓存大小清理早期用不到的权重文件如果你想确认缓存目录在哪Python 里可以直接打印import timm from timm.models.hub import get_cache_dir print(get_cache_dir())6.5 当模型行为奇怪时先查版本最后一条经验也是出现频率最高的模型行为和你预期不一致时先检查 timm 版本。这个库迭代很快同一模型名在不同版本里的细节实现可能有变化比如分类头命名、默认 dropout、甚至正确率水平都会有差异。import timm print(timm.__version__)如果项目里 A 同事和 B 同事报告的结果对不上八成是 timm 或 PyTorch 版本不一致。养成锁定运行环境依赖版本的习惯能省掉很多无意义的排查。最后想说的用 timm 这一年多我最大的感受是它不只是一个模型动物园更是一套把“从模型到训练再到部署”的脏活累活标准化的工作流。无论是快速做 baseline、微调下游任务还是横向对比四五个结构timm 都能让你把注意力放回项目本身而不是反复去适配不同作者写的代码风格。如果你正准备开始一个视觉项目我的建议很直接先用timm.list_models()翻一翻找几个主流的 CNN 和 Transformer 模型做一轮小 baseline再根据精度、速度、参数量权衡。等选定结构之后把create_transform、Mixup、ModelEmaV2这些现成的轮子用起来训练管线也就具备了一个合格项目的雏形。等模型上线后遇到问题也别忘了回来看看default_cfg和部署预处理的细节很多精度波动其实都藏在这类不起眼的地方。
阅读完成 · 觉得有帮助？
