TransXNet实战:用D-Mixer替代注意力机制,高效完成图像分类微调
简介围绕TransXNet在图像分类任务中的实战应用资源核心是使用transxnet_t模型完成植物分类任务并在数据集上实现了96%以上的准确率。TransXNet在ImageNet-1K上相比Swin-T有更高的top-1准确率和更低的计算成本具备较好的模型扩展性与泛化能力。面向计算机视觉初学者、想了解TransXNet落地方法的开发者提供了从模型配置、数据准备到训练推理的完整参考。压缩包内含2000个文件以1978张png格式的植物图像数据集为主体配合6个py训练/推理脚本、6个xml标注或工程配置、2个json含class.json类别映射、1个pth预训练权重等整体约785.92MB目录结构清晰便于按需取用。已有454人学习下载。通过该资源可以了解TransXNet的D-Mixer模块设计思路、模型配置文件写法以及微调训练流程适合需要快速复现高精度图像分类任务、或在此基础上迁移到其他数据集的读者参考实践。1. TransXNet 图像分类一个用 D-Mixer 换掉注意力机制的高性价比选择还记得第一次看到 TransXNet 的对比数据时我确实犹豫了一下TransXNet-T 相比 Swin-T 在 ImageNet-1K 上 top-1 高出 0.3%计算成本还更低这种“精度更高、开销更小”的组合通常只在理想实验里出现。为了验证它落到真实数据集上还能不能打我用 transxnet_t 在一个植物幼苗分类任务上完整跑了一遍从预训练权重加载、微调到推理输出最终 ACC 稳定在 96% 以上。整个过程不是一路顺畅的依赖版本、权重键名、数据增强强度都让我翻过车。这篇实战笔记就是把这些细节拆开告诉你哪些参数值得调、哪些坑必须避开。2. 选型与环境transxnet_t 为什么够用训练前要准备的三样东西2.1 三档模型怎么选TransXNet-T/S/B 精度与计算成本的权衡TransXNet 系列里 T、S、B 三档模型用的核心模块都是 D-Mixer。这个模块的巧妙之处在于它把输入特征图按通道分成两个分支一个分支走轻量化的全局建模路径另一个分支走局部细节增强路径最后再融合起来。和传统 Transformer 里完整的自注意力机制相比D-Mixer 避免了对整张特征图做二次复杂度的全局注意力计算因此在保持感受野的同时把计算量压了下来。对于图像分类这种任务尤其是我手里的植物幼苗数据集叶片纹理、边缘形状这些中高层特征比全局依赖更重要所以 D-Mixer 的混合策略从原理上就是合适的。三档模型的选型可以从实际需求倒推。我的数据集规模不大训练集大概一万出头类别有 12 类很多类别之间的差异集中在叶片颜色和纹理上并不需要特别大的模型容量去拟合细粒度差异。这种情况下TransXNet-B 的 84.6% top-1 固然好看但参数量和推理耗时对我来说是过剩的TransXNet-S 的 83.8% 精度表现很有吸引力可它的计算量也逼近了我的显卡预算上限反倒是 TransXNet-T在 ImageNet 上已经有和 Swin-T 对标的精度计算成本又是三档里最低的自然成了首选。模型档位ImageNet-1K top-1计算成本定位适合场景TransXNet-T约 81.6%对比 Swin-T 高 0.3%低中小数据集、单卡训练、快速迭代TransXNet-S83.8%中中等规模数据、细粒度分类TransXNet-B84.6%高大数据集、追求极致精度把这张表打印出来贴在工位上的好处是每次有人问我“为什么不用更大的模型”我都能直接指着计算成本那一列说清楚。做工程选型不是挑精度最高的而是挑“精度够用且训练周期可控”的。transxnet_t 对单卡 24G 显存以下的场景非常友好默认输入尺寸 224 时batch size 可以放到 64 甚至更大这在后面的训练脚本里会体现出来。2.2 环境准备虚拟环境、依赖库与预训练权重确定用 transxnet_t 之后先把环境隔离出来。我习惯用 conda 管理不同模型的依赖尤其 TransXNet 这类更新较快的项目依赖版本经常互相打架共用环境会把人逼疯。创建环境时 Python 版本选择 3.9 或 3.10 都比较稳妥PyTorch 建议装 2.1 以上后续如果需要用 torch.compile 加速训练会方便很多。conda create -n transxnet python3.9 -y conda activate transxnet pip install torch2.1.0 torchvision0.16.0 pip install timm opencv-python tqdm tensorboard依赖里 timm 很重要。虽然 transxnet_t 的模型定义不依赖 timm 直接实例化但训练时的数据增强、CosineAnnealing 策略和一些辅助工具函数都是从 timm 里调的少了它后面会到处报错。opencv-python 负责读取图像和部分预处理tensorboard 用来盯训练曲线tqdm 纯粹是让训练过程不那么难熬。接着是预训练权重。在模型发布页面找到 transxnet_t 对应的权重文件后我习惯建一个 weights 目录单独放它并且用重命名的规范写清楚来源。仓库里自带的 TransXNet_Demo.iml 是编辑器模块文件用 PyCharm 打开项目时它会自动识别项目结构省去手动配置的麻烦这个文件不需要额外处理。mkdir -p weights mv transxnet_t.pth weights/下载完成后先别急着写训练代码用 Python 快速验证一下权重能不能正常加载这一步能省掉后面很多排错时间。常见的情况是权重文件下载不完整、编码格式不对或者模型定义里 class 数量和预训练头不一致。先跑一个最小加载代码确认权重本身是好的再进入下一步。import torch ckpt torch.load(weights/transxnet_t.pth, map_locationcpu) if state_dict in ckpt: keys list(ckpt[state_dict].keys()) else: keys list(ckpt.keys()) print(权重键数量:, len(keys)) print(前三个键名:, keys[:3])这里的 map_locationcpu 是为了避免本机没有 GPU 时加载报错如果权重文件里嵌套了 state_dict 键说明作者在保存时包了一层后面加载模型要对应取出来。打印前三个键名能帮你确认权重里是否包含分类头一般最后一个键名会带有 head 或 fc 字样以及键名格式是 model.xxx 还是直接裸的 xxx这直接决定了 load_state_dict 时要不要去掉前缀。2.3 训练前固定随机种子与数据划分让每次实验都有后悔药模型选好了环境也通了训练前的最后一道工序是固定随机种子和划分数据集。很多初学者不重视这一步结果就是同一个脚本每次跑出来的精度都不一样换个机器更是完全对不上根本没法判断哪个改动真正有效。我一般会在 config 文件里统一管理种子和路径保证复现不是玄学。# config.py import os import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) torch.backends.cudnn.deterministic True torch.backends.cudnn.benchmark Falsedef split_dataset(data_root, val_ratio0.15, seed42): train_dirs, val_dirs [], [] for cls_name in sorted(os.listdir(data_root)): cls_path os.path.join(data_root, cls_name) if not os.path.isdir(cls_path): continue imgs os.listdir(cls_path) rng random.Random(seed) rng.shuffle(imgs) split_idx int(len(imgs) * (1 - val_ratio)) train_dirs.extend([(cls_name, img) for img in imgs[:split_idx]]) val_dirs.extend([(cls_name, img) for img in imgs[split_idx:]]) return train_dirs, val_dirsset_seed 里两个关键的开关是 cudnn.deterministic 和 cudnn.benchmark。前者告诉 cuDNN 只使用确定性算法保证卷积计算在多次运行间结果一致后者在关闭状态时会放弃自动选择最优卷积算法速度上会有轻微损失但换来的是可复现性。split_dataset 里我用 Random(seed) 而不是全局 random这样即使其他地方调用了随机数数据划分的结果也不会被影响。数据划分比例的取舍也可以说一句15% 的验证集比例对一万张规模的数据集是合理的留出来大约 1500 张做验证既能看出模型真实水平又不会让训练集变得太局促。如果数据集只有两三千张验证集比例可以降到 10%但对应的评估指标波动会变大需要结合多次实验结果一起看。3. 训练一支植物分类器config 参数与两阶段微调的完整过程3.1 数据集目录结构从散图到标准格式植物幼苗数据集最初是从多个来源收集的散图文件名混乱、类别大小不一直接喂给训练脚本会出问题。我常用的做法是把数据整理成标准目录结构train 和 val 两个根目录下面每个类别一个子文件夹图片按类别归位。这样配合 torchvision 的 ImageFolder 接口可以直接读不需要手写 Dataset 读取逻辑。from torchvision.datasets import ImageFolder train_dataset ImageFolder(data/plant/train, transformtrain_transform) val_dataset ImageFolder(data/plant/val, transformval_transform) print(类别映射:, train_dataset.class_to_idx)ImageFolder 会自动扫描子目录并生成类别到索引的映射类别名就是文件夹名。这一行代码省掉了手写样本读取的麻烦还顺带帮你确认了目录结构对不对。class_to_idx 里的映射是后续生成 class.json 的依据一般会看到类似 {apple_leaf: 0, grape_leaf: 1, tomato_leaf: 2} 这样的形式。需要注意的是ImageFolder 要求 train 和 val 下面的类别文件夹名字保持一致否则两边索引会对不上。整理完目录后还要看一眼每个类别的样本数。症状是如果是按类别比例划分的训练集和验证集通常还算均衡但如果是从某个分类网站上直接爬的图很可能某个类别只有二三十张另一个类别上千张这种不平衡会直接反映在最终混淆矩阵里。我的建议是先把类别样本分布打印出来再决定要不要做数据增强补偿。3.2 训练脚本transxnet_t 的关键参数逐条说明训练脚本是整个项目里最值得花时间的部分。transxnet_t 和普通 CNN 分类模型的训练流程在骨架上是类似的但它的 D-Mixer 结构和预训练分布决定了参数需要做一些针对性调整。下面是我用的完整训练配置每一条都标注了为什么这么设。model_cfg dict( model_nametransxnet_t, num_classes12, input_size224, drop_path_rate0.1, ) train_cfg dict( epochs150, batch_size64, lr5e-4, min_lr1e-6, warmup_epochs5, weight_decay0.05, label_smoothing0.1, )drop_path_rate 是 TransXNet 这类带 Transformer 结构模型特有的正则化参数它随机丢弃网络中某些路径模拟多模型集成效果对防止过拟合很有帮助。0.1 是一个在中小数据集上比较安全的起点如果你发现训练集准确率很高而验证集跟不上可以往 0.2 方向加。label_smoothing 设为 0.1 则会让损失函数不单纯追求把正确类别的概率推向 1留一点余量给其他类别对提升验证集准确率有实际帮助尤其是植被叶片这种类间有相似性的任务。from model.transxnet import TransXNet model TransXNet(**model_cfg) ckpt torch.load(weights/transxnet_t.pth, map_locationcpu) if state_dict in ckpt: ckpt ckpt[state_dict] missing_keys, unexpected_keys model.load_state_dict(ckpt, strictFalse) print(缺失键:, missing_keys) print(意外键:, unexpected_keys)加载预训练权重时我没有用默认的 strictTrue这是故意的。因为预训练模型的分类头通常是 1000 类而我们的任务是 12 类加载时分类头必然不匹配。用 strictFalse 可以先加载主干部分的权重缺失的键正好是分类头层的参数后续随机初始化后由训练来填充。打印出来的缺失键和意外键信息要仔细看一遍如果意外键太多可能是权重结构不对需要检查是否要加model.前缀如果缺失键里既有分类头又有主干里的层则要考虑权重本身是否完整。3.3 两阶段微调先训头再解冻的稳健策略正式训练前我先做两阶段微调。第一阶段把主干部分冻结只训练分类头目的是让新初始化的分类头快速适应 D-Mixer 提取的特征分布。第二阶段再把主干解冻用较小的学习率整体微调让预训练特征往植物叶片这个特定领域稍微偏移。这种方案在数据集与预训练分布差距较大的场景下很有效能避免一开始就大范围更新参数导致特征崩塌。for name, param in model.named_parameters(): if head not in name and fc not in name: param.requires_grad False stage1_optimizer torch.optim.AdamW( filter(lambda p: p.requires_grad, model.parameters()), lr1e-3, weight_decay0.05 )冻结判断条件里我同时匹配了 head 和 fc因为不同版本代码里分类头的命名可能不同宁可多匹配一个也不会漏。filter 确保冻结的层不会进入优化器否则 AdamW 会对不需要梯度的参数计算负梯度轻则警告重则报错。Stage 1 一般跑 10 到 15 个 epoch 就够了看验证集准确率不再上升就停。for name, param in model.named_parameters(): param.requires_grad True stage2_optimizer torch.optim.AdamW( model.parameters(), lr5e-4, weight_decay0.05 ) scheduler torch.optim.lr_scheduler.CosineAnnealingLR( stage2_optimizer, T_max150, eta_min1e-6 )Stage 2 解冻全部参数学习率从 5e-4 开始按余弦曲线下降到 1e-6。这里的 T_max 是余弦退火的半个周期长度我设成与剩余训练 epoch 数一致这样学习率曲线会在最后一个 epoch 恰好降到最低点。为什么用余弦退火而不是逐步下降D-Mixer 这类混合架构在训练后期对学习率很敏感余弦退火在接近收敛时给出平滑且缓慢的学习率衰减能比阶梯下降更稳定地逼近最优解。训练迭代部分的代码本身不复杂核心就是五个动作取 batch、前向计算、算损失、反向传播、更新参数。值得注意的细节是梯度裁剪这对 Transformer 类的模型几乎是必须的。scaler torch.cuda.amp.GradScaler() for batch_idx, (inputs, targets) in enumerate(train_loader): inputs, targets inputs.cuda(), targets.cuda() optimizer.zero_grad() with torch.cuda.amp.autocast(): outputs model(inputs) loss criterion(outputs, targets) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0) scaler.step(optimizer) scaler.update()混合精度训练能显著降低显存占用和加快训练速度但它会改变梯度的数值范围所以梯度裁剪必须放在 scaler.unscale_ 之后否则裁剪的数值不是真实的梯度。clip_grad_norm_ 的 max_norm 设成 5.0 是为了防止个别 batch 的异常梯度把整个模型参数推飞。训练 D-Mixer 时我遇到过几次 loss 突然变成 nan 的情况加上这一步后再也没有出现过。3.4 训练过程监控loss 曲线与验证集准确率的联动判断训练过程中我一般用 TensorBoard 同时盯训练集 loss、验证集 loss 和验证集准确率三条曲线。植物分类任务在 150 个 epoch 的训练里曲线大致会是这样的走势前 10 个 epoch 由于在训分类头验证集准确率快速从随机水平爬到 60% 左右第 11 个 epoch 解冻主干后准确率会有一个明显跳升冲到 85% 以上然后进入慢速爬坡期每一步可能只有 0.1% 到 0.3% 的提升到最后一个 epoch 附近准确率曲线基本走平。import json best_epoch, best_acc 0, 0.0 train_history, val_history [], [] for epoch in range(1, train_cfg[epochs] 1): train_loss run_one_epoch(model, train_loader, optimizer, criterion, scheduler, scaler, is_trainTrue) val_acc evaluate(model, val_loader) train_history.append({epoch: epoch, loss: train_loss}) val_history.append({epoch: epoch, acc: val_acc}) if val_acc best_acc: best_acc val_acc best_epoch epoch torch.save(model.state_dict(), checkpoints/best_model.pth) print(fEpoch {epoch}: 新最佳 ACC {val_acc:.4f}) if epoch % 10 0: print(fEpoch {epoch}/{train_cfg[epochs]} | loss {train_loss:.4f} | val_acc {val_acc:.4f}) result { model: transxnet_t, best_epoch: best_epoch, best_val_acc: float(best_acc), train: train_history, val: val_history, } with open(result.json, w, encodingutf-8) as f: json.dump(result, f, indent2, ensure_asciiFalse)保存策略这里也值得一提只有当验证集准确率刷新历史最佳时才覆盖保存权重这样最后手头一定是最优模型而不是最后一个 epoch 的模型。result.json 里我刻意把 train 和 val 两个历史记录都存下来后续画曲线或写报告直接用这个文件不用重新训练。文件里 best_epoch 和 best_val_acc 是最终要汇报的核心指标其他字段是回溯用的。看曲线时要有一个基本判断训练集 loss 和验证集 loss 之间的差距如果在拉大就是过拟合信号验证集准确率在某个 epoch 后开始波动但训练集准确率还在上涨说明正则化强度不够可以加 drop_path_rate 或者增强数据增强。如果验证集和训练集准确率同时停滞在低水平那大概率是学习率或模型容量的问题不要盲目加正则化。4. 推理验证与结果文件class.json 与 result.json 的字段含义4.1 推理脚本从一张单图到 Top-5 置信度分布训练完成后用验证集里随机抽出的单张图片走一遍推理流程是检验模型是否真的可用的第一道关。推理脚本和训练脚本的预处理必须保持一致这一点印象特别深我最初推理时用了和训练不同的标准化参数结果模型在验证集上 acc 只有 40%排查了半天才发现是 Normalize 的 mean 和 std 写错了这类低级错误只要养成对比预处理代码的习惯就不会犯。from PIL import Image import torchvision.transforms as T val_transform T.Compose([ T.Resize(248), T.CenterCrop(224), T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) img Image.open(samples/5e4d1ee0d.png).convert(RGB) x val_transform(img).unsqueeze(0).cuda() model.eval() with torch.no_grad(): logits model(x) prob torch.softmax(logits, dim-1) top5_prob, top5_idx torch.topk(prob, k5, dim-1)Resize 到 248 再 CenterCrop 到 224 是很多预训练模型的默认推理策略它比直接 Resize 到 224 多了一点中心裁剪的鲁棒性。unsqueeze(0) 补上一个 batch 维度因为模型期望的输入是四维张量。model.eval() 必须在推理前调用否则模型里如果有 Dropout 或 DropPath 层它们会保持训练模式导致推理结果带有随机性。仓库里的那些以哈希命名的 png 样例图比如 5e4d1ee0d.png就是用来做快速验证的直接替换图片路径即可。打印 Top-5 结果时要把索引映射回类别名。这就要用到 class.json 了。4.2 class.json 与 result.json推理结果的两个关键文件class.json 是从训练集的类别目录生成的本质是 class_to_idx 的反向映射。我在项目里见过有人直接把 ImageFolder 返回的索引当成类别名去展示结果界面上显示的都是数字完全没法看。正确做法是生成 class.json 后在推理脚本里加载它把索引转成可读的类别名称。{ 0: apple_leaf, 1: grape_leaf, 2: tomato_leaf, 3: potato_leaf, 4: corn_leaf, 5: cucumber_leaf, 6: strawberry_leaf, 7: pepper_leaf, 8: eggplant_leaf, 9: soybean_leaf, 10: pumpkin_leaf, 11: watermelon_leaf }注意 class.json 的键是字符串形式的数字这是 JSON 格式本身没有整数键导致的。读取时要先 int(k) 再取类别名这个细节在第一版推理脚本里让我卡了十分钟显示出来全是 None 才意识到类型问题。import json with open(class.json, r, encodingutf-8) as f: idx_to_class {int(k): v for k, v in json.load(f).items()} for rank, (p, idx) in enumerate(zip(top5_prob[0], top5_idx[0])): print(fTop-{rank 1}: {idx_to_class[int(idx.item())]:20s} {p.item():.4f})result.json 则是训练完成后自动生成的记录的是训练过程和最终指标。这个文件在最终交付或写项目报告时直接就能用。它的结构在第 3.4 节的训练代码里已经给出了best_epoch 表示验证集准确率最高点出现的轮次best_val_acc 是对应最高准确率train 和 val 列表分别记录了每个 epoch 的 loss 和准确率。查看 result.json 时如果发现 best_epoch 特别靠后说明模型在最后阶段还在持续提升可以考虑延长训练时长如果 best_epoch 很早但后面波动大就说明学习率后期衰减不够平滑或者训练集和验证集分布有偏差。4.3 一次翻车排查背景干扰如何让置信度错乱模型训练完acc 96% 看着很不错结果测试新采集的图片时被一记闷棍打醒。有一张葡萄叶子的照片叶片本身拍得很清晰但背景里有几片褐色的枯叶模型给出的 top-1 结果是番茄叶置信度还高得离谱接近 0.93。模型在验证集上明明表现很好的为什么换到真实场景就翻车了排查步骤从数据和模型两个方向展开。先确认这张翻车图是否在数据分布内新采集图片和训练集图片的区别在于背景复杂度、光照条件和拍摄角度。训练集里的样本大多是实验室背景下拍摄的背景干净单一而新照片在自然环境下拍摄枯叶、泥土、甚至手指都可能入镜。模型并没有真正学到“葡萄叶片的纹理特征”而是学到了“绿色背景特定形状叶片”这个组合当背景中出现类似叶片纹理的干扰物时它就会激活错误的特征通道。对应的解法有两个层次。短期方案是增加数据增强的随机性在训练时加入 RandomResizedCrop 的尺度范围调整和多尺度训练让模型见过更多局部裁切和尺度变化长期方案是收集更多真实场景下的负样本把带枯叶、带杂草背景的照片加入训练集。我采取的落地操作是先生成一张低置信度样本列表把所有验证集里置信度低于 0.85 的图片全部导出让人工过一遍标注是否正确。这一步操作虽然原始但能快速定位模型失效的模式。import csv low_conf_samples [] for path, target, prob in zip(val_paths, val_targets, all_probs): if prob 0.85: low_conf_samples.append((path, idx_to_class[target], float(prob))) with open(low_conf_samples.csv, w, newline) as f: writer csv.writer(f) writer.writerow([path, true_label, confidence]) writer.writerows(low_conf_samples)5. 常见问题避坑五个训练中的高频翻车点和修复方法5.1 预训练权重键名不匹配KeyError 与 strict 加载失败现象运行加载权重的代码时直接抛 KeyError或者模型初始化后所有参数都停留在随机值训练第一篇 loss 就炸到很高。原因不同来源的 TransXNet 权重保存格式不统一。有一些仓库在权重文件里多包了一层 model 前缀另一些则直接把所有 state_dict 平铺。加上分类头名称在不同实现里可能是 head 也可能是 fc.weight导致本地模型定义和权重文件里的键名对不上。解决先把权重里的键名打印出来看看是否有统一前缀再决定加载方式。state_dict torch.load(weights/transxnet_t.pth, map_locationcpu) if state_dict in state_dict: state_dict state_dict[state_dict] new_state_dict {} for k, v in state_dict.items(): if k.startswith(module.): k k[7:] if k.startswith(model.): k k[6:] new_state_dict[k] v model.load_state_dict(new_state_dict, strictFalse)5.2 数据增强过度导致欠拟合准确率卡在 75% 怎么都上不去现象训练集和验证集准确率同步停滞在 75% 左右loss 也在高位浮动增加训练轮次毫无改善。原因同时开了 RandAugment、MixUp、CutMix 三种增强策略对一万张规模的植物数据集来说增强强度过猛。模型在训练时看到的图片被破坏得太厉害叶片纹理被严重遮挡或拼接学到的特征不够稳定。解决小数据集上我最终只保留了 RandomResizedCrop、RandomHorizontalFlip 和轻微的 ColorJitter。先跑通 baseline确认模型能过拟合训练集再逐步增强增强策略的强度。如果训练集准确率能在 30 个 epoch 内冲到 98%说明模型容量和增强强度是匹配的再开始尝试 MixUp 也不迟。5.3 学习率设置失败warmup 结束就开始发散现象训练前 5 个 epoch 的 loss 平稳下降但从第 6 个 epoch 开始 loss 突然变 nan 或者验证集准确率断崖下跌。原因transxnet_t 这类带 D-Mixer 的混合架构对初始学习率比较敏感。如果 warmup 阶段起始学习率设置过高模型在刚开始时参数更新幅度过大会破坏预训练权重的特征分布等到 warmup 结束时问题彻底爆发。解决把 warmup 的起始学习率降到 1e-7让模型在 warmup 阶段从几乎不更新逐步过渡到正常学习率。具体实现可以用 timm 里现成的 LambdaLR。def warmup_lambda(epoch): if epoch warmup_epochs: return (epoch 1) / warmup_epochs * 0.1 return 1.0 scheduler torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambdawarmup_lambda )5.4 类别不平衡导致弱势类准确率低大类和少类差距悬殊现象整体验证集准确率 96%但看每个类别的单独准确率叶子特征明显的类别能到 98%而在样本数少的类别上只有 62%在混淆矩阵里大量被预测成大类别。原因植物数据集中有些类别采集难度高比如某种病害叶片只有几十张样本。模型在训练中见过的大类别样本多决策边界会偏向大类别少类别样本即便被正确分类置信度也被压低。解决给 DataLoader 接入加权采样器让样本数少的类别每个 epoch 被抽到的概率更大。注意 sampling_weight 需要按类别样本数量反比计算。from torch.utils.data.sampler import WeightedRandomSampler class_counts torch.bincount(torch.tensor(train_dataset.targets)) class_weights 1.0 / class_counts.float() sample_weights class_weights[train_dataset.targets] sampler WeightedRandomSampler(sample_weights, num_sampleslen(sample_weights), replacementTrue) train_loader DataLoader(train_dataset, batch_size64, samplersampler)5.5 训练中途卡死DataLoader 的 num_workers 设置不合理现象训练进行到第 20 个 epoch 时进度条突然不动GPU 利用率掉到 0%终端也没有任何报错输出像是整个进程被挂起。原因num_workers 设置过大同时配合某些数据集文件存储在高延迟的机械硬盘上数据读取线程之间相互阻塞进程死锁。另一个常见原因是 Windows 系统下 num_workers 超过 0 时在子进程里再次调用数据集读取逻辑导致句柄冲突。解决先把 num_workers 降到 4并加上 persistent_workersTrue 让子进程跨 epoch 常驻。这个方法在 Linux 和 Windows 下都能稳定运行。如果数据集特别大可以考虑先做一次预取缓存把所有图片读取并缩放到固定尺寸后保存成二进制格式训练时只读缓存文件能大幅减少 IO 等待。6. 进阶技巧用特征图可视化验证 D-Mixer 的判别依据再决定要不要调大模型模型 acc 到了 96% 之后我遇到一个新问题剩下的 4% 错误到底是数据问题还是模型能力天花板如果盲目换 TransXNet-S训练时间翻倍提升未必明显。所以先用特征图可视化确认 D-Mixer 到底学到了什么。用 PyTorch 的 forward hook 把靠近 D-Mixer 输出的特征图抓出来叠加到原图上生成热力图。activation {} def make_hook(name): def hook_fn(module, input, output): activation[name] output.detach() return hook_fn model.blocks[0].register_forward_hook(make_hook(block0)) model.eval() with torch.no_grad(): _ model(x) feat activation[block0][0, 0].cpu().numpy() feat (feat - feat.min()) / (feat.max() - feat.min() 1e-8)这个热力图能从直观上判断模型是盯着叶片纹理做判断还是在靠背景色蒙答案。如果热力图高亮区域集中在叶片边缘和叶脉说明模型学到了可靠的判别特征如果高亮区域散落在背景里前面的数据增强和清洗就得再走一遍。基于这个发现我给推理流程加了一个低置信度拦截机制当 top-1 置信度低于 0.8 时不再直接输出类别而是把样本进入待人工复核队列。threshold 0.8 if top5_prob[0][0] threshold: print(该样本置信度不足建议人工复核) else: print(预测结果:, idx_to_class[int(top5_idx[0][0])])从那以后我每次训练完成都会强制走一遍低置信度样本检查把 image path、true_label、pred_label、confidence 四列信息并排打印出来逐个确认是模型问题还是标注问题。检查完再回头清洗数据模型在植物分类任务上最终从 96.1% 推到了 96.4%。这四个百分点的提升不是调参调出来的而是把模型置信度最低的样本一张张看出来的这种笨功夫在图像分类项目里是最值的投入。希望这套 TransXNet 实战流程里的选型逻辑、训练参数和踩坑记录能帮你少走几段弯路。本文还有配套的精品资源点击获取