脚本

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

训练和验证脚本从 PyTorch Imagenet Examples 的早期版本演变而来。随着时间的推移，我添加了重要的功能，包括基于 NVIDIA’s APEX Examples 的 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 将强制使用它们。

验证 / 推理脚本

验证和推理脚本的用法类似。一个在验证集上输出指标，另一个在 csv 中输出 topk 类 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

训练示例

EfficientNet-B2 与 RandAugment - 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

MixNet-XL 与 RandAugment - 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+ 个 epoch。对于更好的 GPU 或启用 AMP 的情况，按比例增加批量大小和 LR。这些参数适用于 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

EfficientNet-B3 与 RandAugment - 81.5 top-1, 95.7 top-5

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

EfficientNet-B0 与 RandAugment - 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

ResNet50 与 JSD 损失和 RandAugment（clean + 2x RA augs）- 79.04 top-1, 94.39 top-5

在两张较旧的 1080Ti 显卡上训练，这花了一段时间。与我第一次 78.99 的良好 AugMix 训练相比，ImageNet 验证结果仅略好，且不具有统计学意义。但是，这些权重在 ImageNetV2、ImageNet-Sketch 等测试中更稳健。与我的第一次 AugMix 运行不同，我启用了 SplitBatchNorm，禁用了 clean 分割上的随机擦除，并提高了 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

EfficientNet-ES (EdgeTPU-Small) 与 RandAugment - 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

ResNeXt-50 32x4d 与 RandAugment - 79.762 top-1, 94.60 top-5

这些参数也适用于 SE-ResNeXt-50 和 SK-ResNeXt-50，可能也适用于 101。我将它们用于 SK-ResNeXt-50 32x4d，我在 2 个 GPU 上训练它，每个有效批量大小使用稍高的 LR（lr=0.18，每个 GPU b=192）。下面的 cmd 命令行针对 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 上更新