指令碼

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 上更新