脚本

GitHub 根文件夹中包含的训练、验证、推理和检查点清理脚本。脚本目前未打包在 pip 版本中。

训练和验证脚本从 PyTorch Imagenet 示例的早期版本演变而来。随着时间的推移，我添加了重要的功能，包括基于 NVIDIA 的 APEX 示例的 CUDA 特定性能增强功能。

训练脚本

训练参数种类繁多，并非所有选项组合（甚至选项本身）都经过了充分测试。对于训练数据集文件夹，请指定包含 train 和 validation 文件夹的基文件夹。

要使用余弦调度、50% 的随机擦除概率和逐像素随机值，在本地分布式环境下、使用 4 个 GPU、每个 GPU 一个进程，在 ImageNet 上训练 SE-ResNet34

./distributed_train.sh 4 /data/imagenet --model seresnet34 --sched cosine --epochs 150 --warmup-epochs 5 --lr 0.4 --reprob 0.5 --remode pixel --batch-size 256 --amp -j 4

建议使用 PyTorch 1.9+ 与 PyTorch 原生 AMP 和 DDP，而不是 APEX AMP。--amp 默认情况下从 timm 版本 0.4.3 开始使用原生 AMP。--apex-amp 将强制使用 APEX 组件（如果已安装）。

验证/推理脚本

验证和推理脚本在用法上类似。一个输出验证集上的指标，另一个输出 CSV 文件中的 topk 类别 ID。指定包含验证图像的文件夹，而不是像训练脚本那样指定基文件夹。

要使用模型的预训练权重（如果存在）进行验证

python validate.py /imagenet/validation/ --model seresnext26_32x4d --pretrained

要从检查点运行推理

python inference.py /imagenet/validation/ --model mobilenetv3_large_100 --checkpoint ./output/train/model_best.pth.tar

训练示例

EfficientNet-B2 使用 RandAugment - top-1 80.4，top-5 95.1

这些参数适用于安装了 NVIDIA Apex 的双 Titan RTX 卡

./distributed_train.sh 2 /imagenet/ --model efficientnet_b2 -b 128 --sched step --epochs 450 --decay-epochs 2.4 --decay-rate .97 --opt rmsproptf --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.3 --drop-path 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --amp --lr .016

MixNet-XL 使用 RandAugment - top-1 80.5，top-5 94.9

这些参数适用于安装了 NVIDIA Apex 的双 Titan RTX 卡

./distributed_train.sh 2 /imagenet/ --model mixnet_xl -b 128 --sched step --epochs 450 --decay-epochs 2.4 --decay-rate .969 --opt rmsproptf --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.3 --drop-path 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.3 --amp --lr .016 --dist-bn reduce

SE-ResNeXt-26-D 和 SE-ResNeXt-26-T

这些超参数（或类似参数）适用于各种 ResNet 架构，一般来说，随着模型大小的增加，增加 epoch 数量是一个好主意……例如，对于 ResNe(X)t50 大约为 180-200，对于更大的模型则为 220+。对于更好的 GPU 或启用 AMP 时，按比例增加批次大小和 LR。这些参数适用于 2 个 1080Ti 卡

./distributed_train.sh 2 /imagenet/ --model seresnext26t_32x4d --lr 0.1 --warmup-epochs 5 --epochs 160 --weight-decay 1e-4 --sched cosine --reprob 0.4 --remode pixel -b 112

EfficientNet-B3 使用 RandAugment - top-1 81

该模型的训练从与上面 EfficientNet-B2 w/ RA 相同的命令行开始。经过近三周的训练，进程崩溃了。结果看起来并不理想，因此我多次调整了一些参数（增加 RE 概率，降低 rand-aug，增加 ema-decay）以恢复训练。但效果仍然不佳。最终，我对所有重启过程中最好的检查点进行了平均。结果在默认的 res/crop 下平庸，但奇怪的是，在 1.0 的全图像测试裁剪下表现要好得多。

EfficientNet-B0 with RandAugment - top-1 77.7，top-5 95.3

Michael Klachko 使用为更大批量大小调整的 B2 命令行，以及推荐的 B0 dropout 率 0.2 实现了这些结果。

./distributed_train.sh 2 /imagenet/ --model efficientnet_b0 -b 384 --sched step --epochs 450 --decay-epochs 2.4 --decay-rate .97 --opt rmsproptf --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.2 --drop-path 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --amp --lr .048

ResNet50 with JSD loss and RandAugment (clean + 2x RA augs) - top-1 79.04，top-5 94.39

这个模型是在两张老的 1080Ti 显卡上训练的，花了很长时间。ImageNet 验证结果只比我的第一个好的 AugMix 训练结果 78.99 略好，但没有统计学意义上的差异。然而，这些权重在 ImageNetV2、ImageNet-Sketch 等测试中更加鲁棒。与我的第一个 AugMix 运行不同，我在 clean 分割上启用了 SplitBatchNorm，禁用了随机擦除，并在两个增强的路径上增加了随机擦除概率。

./distributed_train.sh 2 /imagenet -b 64 --model resnet50 --sched cosine --epochs 200 --lr 0.05 --amp --remode pixel --reprob 0.6 --aug-splits 3 --aa rand-m9-mstd0.5-inc1 --resplit --split-bn --jsd --dist-bn reduce

EfficientNet-ES (EdgeTPU-Small) with RandAugment - top-1 78.066，top-5 93.926

由 Andrew Lavin 使用 8 张 V100 显卡训练。没有使用模型 EMA，最终检查点是训练期间 8 个最佳检查点的平均值。

./distributed_train.sh 8 /imagenet --model efficientnet_es -b 128 --sched step --epochs 450 --decay-epochs 2.4 --decay-rate .97 --opt rmsproptf --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.2 --drop-path 0.2  --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --amp --lr .064

MobileNetV3-Large-100 - top-1 75.766，top-5 92,542

./distributed_train.sh 2 /imagenet/ --model mobilenetv3_large_100 -b 512 --sched step --epochs 600 --decay-epochs 2.4 --decay-rate .973 --opt rmsproptf --opt-eps .001 -j 7 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.2 --drop-path 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --amp --lr .064 --lr-noise 0.42 0.9

ResNeXt-50 32x4d w/ RandAugment - top-1 79.762，top-5 94.60

这些参数也适用于 SE-ResNeXt-50 和 SK-ResNeXt-50，以及可能的 101。我将它们用于 SK-ResNeXt-50 32x4d，使用 2 块 GPU 训练，每个有效批次大小的学习率略高（lr=0.18，每块 GPU b=192）。下面的命令行针对 8 块 GPU 训练进行了调整。

./distributed_train.sh 8 /imagenet --model resnext50_32x4d --lr 0.6 --warmup-epochs 5 --epochs 240 --weight-decay 1e-4 --sched cosine --reprob 0.4 --recount 3 --remode pixel --aa rand-m7-mstd0.5-inc1 -b 192 -j 6 --amp --dist-bn reduce

< > 更新在 GitHub 上