脚本

github 根文件夹中包含一个训练、验证、推理和检查点清理脚本。这些脚本目前未打包在 pip 发布版本中。

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

训练脚本

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

要在 ImageNet 上训练一个 SE-ResNet34，使用本地分布式训练，4 个 GPU，每个 GPU 一个进程，采用余弦调度器，随机擦除概率为 50% 且使用逐像素随机值。

./distributed_train.sh 4 --data-dir /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。从 timm 0.4.3 版本开始，--amp 默认为原生 AMP。如果安装了 APEX 组件，--apex-amp 将强制使用 APEX 组件。

验证/推理脚本

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

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

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

从检查点运行推理

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

训练示例

使用 RandAugment 的 EfficientNet-B2 - 80.4 top-1, 95.1 top-5

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

./distributed_train.sh 2 --data-dir /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

使用 RandAugment 的 MixNet-XL - 80.5 top-1, 94.9 top-5

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

./distributed_train.sh 2 --data-dir /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 个 epoch，对于更大的模型则为 220+ 个。对于更好的 GPU 或启用了 AMP 的情况，应按比例增加批量大小和学习率。这些参数是为 2 块 1080Ti 显卡设置的。

./distributed_train.sh 2 --data-dir /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

使用 RandAugment 的 EfficientNet-B3 - 81.5 top-1, 95.7 top-5

该模型的训练始于与上述 EfficientNet-B2 w/ RA 相同的命令行。经过近三周的训练后，进程崩溃了。结果看起来不怎么样，所以我多次恢复训练并对一些参数进行了调整（增加 RE 概率，减少 rand-aug，增加 ema-decay）。效果都不理想。最后，我平均了所有重启中最好的检查点。结果在默认分辨率/裁剪下表现平平，但奇怪的是，在完整图像裁剪比例为 1.0 的测试中表现要好得多。

使用 RandAugment 的 EfficientNet-B0 - 77.7 top-1, 95.3 top-5

Michael Klachko 使用适用于更大批量大小的 B2 命令行，并采用推荐的 B0 dropout 率 0.2，取得了这些结果。

./distributed_train.sh 2 --data-dir /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

使用 JSD 损失和 RandAugment 的 ResNet50（clean + 2x RA 增强）- 79.04 top-1, 94.39 top-5

在两块较旧的 1080Ti 显卡上训练，耗时较长。ImageNet 验证结果仅比我第一次成功的 AugMix 训练（78.99）略好，无统计学意义。然而，这些权重在 ImageNetV2、ImageNet-Sketch 等测试中更具鲁棒性。与我第一次 AugMix 运行不同，我启用了 SplitBatchNorm，在 clean split 上禁用了随机擦除，并在 2 个增强路径上加大了随机擦除的概率。

./distributed_train.sh 2 --data-dir /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

使用 RandAugment 的 EfficientNet-ES (EdgeTPU-Small) - 78.066 top-1, 93.926 top-5

由 Andrew Lavin 使用 8 块 V100 显卡训练。未使用模型 EMA，最终检查点是训练过程中 8 个最佳检查点的平均值。

./distributed_train.sh 8 --data-dir /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 - 75.766 top-1, 92,542 top-5

./distributed_train.sh 2 /--data-dir 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

使用 RandAugment 的 ResNeXt-50 32x4d - 79.762 top-1, 94.60 top-5

这些参数也适用于 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 --data-dir /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 上更新