Scripts

A train, validation, inference, and checkpoint cleaning script included in the github root folder. Scripts are not currently packaged in the pip release.

The training and validation scripts evolved from early versions of the PyTorch Imagenet Examples. I have added significant functionality over time, including CUDA specific performance enhancements based on NVIDIA’s APEX Examples.

Training Script

The variety of training args is large and not all combinations of options (or even options) have been fully tested. For the training dataset folder, specify the folder to the base that contains a train and validation folder.

For a single-node run, use torchrun with one process per GPU. For example, after saving the MobileNetV4 config as mobilenetv4.yaml:

torchrun --standalone --nproc_per_node=2 train.py \
  --config mobilenetv4.yaml --data-dir /data/imagenet --experiment mobilenetv4-cm

./distributed_train.sh 2 ... is a convenience wrapper around torchrun. For one GPU, use python train.py ... and adjust the batch size and learning rate as described in HParams.

`--amp` uses PyTorch native mixed precision, with float16 by default. Use `--amp-dtype bfloat16` on supported hardware. The examples below use native PyTorch AMP and DDP; Apex is not required.

Config files

Any training argument can be given in a YAML file via --config, keyed by the argument name with underscores (--weight-decay becomes weight_decay). Command line arguments override the file. Each run writes the full set of effective arguments to args.yaml in its output folder, and that file is itself a valid --config, so a run can be reproduced with python train.py --config output/train/my-run/args.yaml --experiment my-run-2. The file records output, experiment and resume too, so without a new experiment name the command writes into the original folder, and with --resume auto continues the original run rather than repeating it. Config keys are validated against the script’s arguments: unknown keys warn, invalid choices or non-numeric values for numeric arguments error.

Resuming and milestone checkpoints

Each run writes last.pth.tar (always the latest epoch, with model, optimizer, EMA and AMP scaler state), model_best.pth.tar, and a ranked history of checkpoint-N.pth.tar files limited by --checkpoint-hist.

--resume PATH restores that state and continues the schedule from the checkpoint’s epoch. This is the usual way to start a new experiment folder from an existing checkpoint.
--resume auto resumes in place: with --output and --experiment set, it continues from last.pth.tar in that experiment folder if one exists and otherwise starts a fresh run. The same command can be re-run after an interruption. The checkpoint history, best metric tracking and summary.csv all carry on consistently, so --checkpoint-hist eviction and model_best stay correct across the restart.
--stop-epoch N stops after N completed epochs without changing the schedule set by --epochs. The final epoch is always evaluated and checkpointed, so a run can be trained in segments.
--save-epochs N [N ...] keeps immutable milestone-N.pth.tar snapshots (N is a completed epoch count) outside the ranked history. They are hardlinked from last.pth.tar where the filesystem allows so they cost no extra space, and an evaluation is forced at those epochs.

Together these make warmup-stable-decay (WSD) experiments easy to run by hand: train a stable trunk once, then branch independent cooldowns from its milestones with full optimizer state.

# stable trunk through 30 completed epochs, keeping milestones at 18, 24 and 30. The 1 epoch decay is never reached.
python train.py --config configs/vit.yaml --sched wsd --sched-on-updates --epochs 31 --decay-epochs 1 \
    --stop-epoch 30 --save-epochs 18 24 30 --output output/wsd --experiment stable --resume auto

# a 6 epoch 1-sqrt cooldown branched from epoch 24, i.e. a 30 epoch run that shares the first 24 epochs with the trunk
python train.py --config configs/vit.yaml --sched wsd --sched-on-updates --epochs 30 --decay-epochs 6 --lr-decay-fn 1-sqrt \
    --resume output/wsd/stable/milestone-24.pth.tar --output output/wsd --experiment decay-p24-d6

The decay occupies the last --decay-epochs of the schedule, so a cooldown of D epochs from P completed epochs is --epochs P+D --decay-epochs D (subtract --warmup-epochs from --epochs when --warmup-prefix is used). sweep.py below automates this.

Hyperparameter Sweeps

sweep.py (PyTorch 2.5 or newer) plans and runs sweeps over any train.py arguments, optionally with WSD branching. It writes one complete train.py --config file per job plus a plan.json with the job dependency graph, then runs ready jobs on a pool of local GPUs. The plan stays fixed while workers run. Every job is an ordinary train.py invocation with its own experiment folder, and each generated run.sh can be run by hand or handed to another scheduler.

A quick sweep needs no spec file. Search axes come from --sweep, the base configuration from the trailing train.py arguments after --:

python sweep.py run --name lr-wd --sweep lr=1e-4,3e-4,1e-3 --sweep weight_decay=0.05,0.1 --gpus 0,1,2,3 \
    -- --model vit_small_patch16_224 --data-dir /data/imagenet --sched cosine --epochs 100 --amp

--gpus lists the GPUs to run on, as indices into CUDA_VISIBLE_DEVICES when it is set (in a Slurm allocation, for example).

A WSD sweep shares one stable trunk per hyperparameter combination across many cooldowns. For each learning rate below, the trunk runs to 30 epochs with milestones at 18, 24 and 30, and a 6 epoch cooldown is branched from each milestone, 4 trunks and 12 cooldowns in total instead of 12 full runs. Cooldowns start as soon as their milestone exists, while the trunk continues. Jobs needing more than one GPU are launched with torchrun:

python sweep.py run --name vit-wsd --sweep lr=1e-3,5e-4,2e-4,1e-4 \
    --wsd-points 18 24 30 --wsd-decay-epochs 6 --wsd-decay-fn 1-sqrt --gpus-per-job 2 --gpus 0,1 \
    -- --model vit_small_patch16_224 --data-dir /data/imagenet --warmup-epochs 3 --warmup-prefix --amp

The planner owns the following flags, so they can’t be given on the command line or swept.

In every sweep:

--output
--experiment
--resume
--stop-epoch
--save-epochs

Additionally, in a WSD sweep:

--sched (it may be given as wsd, which the planner sets anyway)
--epochs
--decay-epochs
--lr-decay-fn

A YAML spec adds random search with distributions, paired (zipped) axes, repeats over seeds and a base config file, see the docstring at the top of sweep.py for the schema. python sweep.py plan spec.yaml writes the configs and graph.dot without running anything.

The sweep lives in output/sweeps/<name>/, relative to the spec file (or the working directory for a command line sweep), with one folder per trial, named after the swept values (trial-000-lr0.001), each containing the job experiment folders. Every kept checkpoint is a full copy with optimizer state, so sweep jobs keep one ranked checkpoint besides last.pth.tar and model_best.pth.tar (--checkpoint-hist 1), and WSD sweeps step the schedule per update (--sched-on-updates). Setting either in the base configuration overrides it. While it runs, or afterwards:

python sweep.py status output/sweeps/vit-wsd    # job states and progress
python sweep.py results output/sweeps/vit-wsd   # ranked endpoint metrics, written to results.csv

Results compare trials at equal budgets. A WSD sweep shows one row per trial and one column per cooldown (p24-d6 is a 6 epoch cooldown from 24 epochs) with the final metric, averaged over repeats, and the best of each column marked. A plain sweep is ranked by its final metric within each epoch budget. results.csv holds every job, trunk segments included, with the best epoch of each. Without a validation set the ranking falls back to the training loss, and the endpoint checkpoints can be evaluated with validate.py.

Running sweep.py run again on an interrupted sweep resumes it: finished jobs are skipped and a job with a partial run continues from its own last.pth.tar. Workers on several machines sharing the sweep folder can each run it. They coordinate through flock locks under state/locks, which every training process of a job holds as well, so a job whose worker died is neither started again nor replanned until its last process has exited. A sweep shared by several machines needs a filesystem whose flock works across all of them, such as NFSv4 or Lustre mounted with -o flock. Mounts with local-only locking don’t coordinate across machines, so check locking from each client. Stopping a worker (Ctrl-C, or SIGTERM from scancel) stops its jobs, and running a job’s run.sh by hand bypasses the coordination.

Lock files under state/locks are permanent: never delete or replace them to clear a busy sweep. Stop or wait for its workers and training processes instead. A running (gone) status means the training lock is free but the attempt has not been reconciled; run a worker to settle it and resume unfinished work. status and results do not repair job state.

A worker reads the plan once when it starts. To change a sweep, stop all of its workers and training processes, plan again with --force, then start them again. Before any job has started, anything about the plan can change. Afterwards its training jobs can only be extended: more trials, cooldowns, or WSD points after the existing ones, where each trunk continues from where it stopped. Changes to existing training jobs need a new sweep. The generated paths are fixed when the plan is written, so plan again under a new location rather than moving a sweep folder.

If planning is interrupted, workers detect configs or launchers from an incomplete plan write and refuse to start. Re-run the original planning command with --force to finish or repair the plan, then restart the workers. Each successful plan write publishes a matching generation in plan.json, the configs and the launchers; an unchanged, complete plan writes nothing.

For example, to extend the CLI WSD sweep above to 36 stable epochs, stop its workers and jobs, then repeat the original planning arguments with the later point appended:

python sweep.py plan --force --name vit-wsd --sweep lr=1e-3,5e-4,2e-4,1e-4 \
    --wsd-points 18 24 30 36 --wsd-decay-epochs 6 --wsd-decay-fn 1-sqrt --gpus-per-job 2 \
    -- --model vit_small_patch16_224 --data-dir /data/imagenet --warmup-epochs 3 --warmup-prefix --amp
python sweep.py run output/sweeps/vit-wsd --gpus 0,1

For a YAML sweep, append the point in the original spec and run python sweep.py plan spec.yaml --force before restarting. Keep wsd.decay_epochs explicit and unchanged ([6] in this example): its default is 20% of the final point, so extending the horizon can otherwise change existing cooldowns. Preserve existing trial IDs; reordering or expanding grid axes can renumber trials. With lr_noise, the trunk’s schedule depends on its horizon, so adding a later point to a started sweep requires a new sweep.

Training settings belong in the plan: use the spec’s train block, or arguments after -- when creating a CLI sweep. There are no runtime training overrides. Reporting settings, such as the ranking metric, and output settings (checkpoint_hist, log_wandb, wandb_project, wandb_tags) may be changed offline with --force without resetting job state or changing training identity. Output settings apply to future attempts; completed jobs are not rerun.

sweep.py requires Unix flock: tested on Linux, expected to work but untested on macOS. Native Windows is unsupported; use WSL2. Locks don’t survive a network partition that outlasts an NFS lease: a disconnected process may keep training after losing its locks, so another machine can start an overlapping attempt.

Validation / Inference Scripts

Validation and inference scripts are similar in usage. One outputs metrics on a validation set and the other outputs topk class ids in a csv. Specify the folder containing validation images, not the base as in training script.

To validate with the model’s pretrained weights (if they exist):

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

To run inference from a checkpoint:

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

Training Examples

ImageNet training recipes

These commands correspond to the selected HParams recipes, where you can find copyable YAML configs, source gists, and published checkpoint results. They assume ImageFolder data under /data/imagenet/train/ and /data/imagenet/validation/. --batch-size is per GPU. The explicit learning rates in the MobileNetV4 and ConvNeXt examples apply to the stated global batch sizes; they do not scale automatically if the number of GPUs changes.

MobileNetV4 Conv Medium, ImageNet-1k

From mnv4_cm_r256_e500_gpu2.yaml: 2 GPUs, global batch 1024, 256px training, and 20 warmup + 500 cosine epochs. The corresponding weights are mobilenetv4_conv_medium.e500_r256_in1k.

torchrun --standalone --nproc_per_node=2 train.py --data-dir /data/imagenet \
  --model mobilenetv4_conv_medium.e500_r256_in1k --img-size 256 --crop-pct 0.95 --batch-size 512 \
  --opt adamw --lr 0.002 --weight-decay 0.1 --clip-grad 5.0 \
  --sched cosine --sched-on-updates --epochs 500 --warmup-epochs 20 --warmup-prefix \
  --warmup-lr 0 --min-lr 0 --cooldown-epochs 0 \
  --aa rand-m8-inc1-mstd1.0 --mixup 0.2 --mixup-prob 0.8 --smoothing 0.1 \
  --reprob 0.25 --remode pixel --drop 0.2 --drop-path 0.1 \
  --model-ema --model-ema-decay 0.9998 --model-ema-warmup --dist-bn reduce \
  --amp --amp-dtype bfloat16 --torchcompile inductor --experiment mobilenetv4-cm

Use --amp-dtype float16 if your GPU does not support bfloat16; --torchcompile inductor is optional.

ViT Medium with registers, ImageNet-1k

From train-vit_medium_reg4-4gpu-in1k.yaml: 4 GPUs, global batch 1024, and 600 epochs including warmup + 10 cooldown epochs. The corresponding weights are vit_medium_patch16_reg4_gap_256.sbb_in1k.

torchrun --standalone --nproc_per_node=4 train.py --data-dir /data/imagenet \
  --model vit_medium_patch16_reg4_gap_256.sbb_in1k --model-kwargs fix_init=True --head-init-scale 0 \
  --img-size 256 --crop-pct 0.95 --batch-size 256 \
  --opt adamw --opt-eps 1e-8 --lr-base 0.0008 --lr-base-size 512 --lr-base-scale sqrt \
  --weight-decay 0.07 --clip-grad 5.0 --sched cosine --sched-on-updates \
  --epochs 600 --warmup-epochs 20 --warmup-lr 5e-7 --min-lr 5e-7 --cooldown-epochs 10 \
  --aa rand-m7-inc1-mstd1.0-n3 --aug-repeats 3 --color-jitter 0 \
  --mixup 0.8 --cutmix 1.0 --smoothing 0.1 --reprob 0.2 --remode pixel --drop-path 0.2 \
  --seed 0 --amp --amp-dtype float16 --torchcompile inductor --experiment vit-medium-sbb

This run uses AdamW without EMA. Its peak LR is approximately 0.001131 at global batch 1024.

ConvNeXt Tiny, ImageNet-12k to ImageNet-1k fine-tuning

From hparams_84_00-convnext_tiny.in12k-args.yaml: 8 GPUs, global batch 4096, and 75 epochs including warmup + 10 cooldown epochs. Load the ImageNet-12k weights and set --num-classes 1000 to replace the classifier for fine-tuning. The resulting checkpoint family is convnext_tiny.in12k_ft_in1k.

torchrun --standalone --nproc_per_node=8 train.py --data-dir /data/imagenet \
  --model convnext_tiny.in12k --pretrained --num-classes 1000 \
  --img-size 224 --crop-pct 1.0 --batch-size 512 --grad-checkpointing \
  --opt adamw --lr 1e-4 --weight-decay 0.05 --layer-decay 0.75 --clip-grad 3.0 \
  --sched cosine --epochs 75 --warmup-epochs 10 --warmup-lr 1e-6 --min-lr 5e-7 \
  --cooldown-epochs 10 --lr-noise 0.1 1.0 \
  --aa rand-m8-inc1-mstd101 --mixup 0.8 --cutmix 1.0 --smoothing 0.1 \
  --reprob 0.3 --remode pixel --drop-path 0.1 --model-ema --model-ema-decay 0.9998 \
  --amp --amp-dtype float16 --experiment convnext-tiny-ft

Multi-label classification

Use --task multilabel for images with multiple independent labels.

--num-classes sets the vocabulary size.
--target-format indices (default) reads the target column as a list of positive, zero-based class IDs. Labels are converted to float multi-hot targets before batching, and empty lists represent images with no positive labels.
--target-format multihot reads a dense multi-hot vector of length --num-classes (bool, int, or float values in [0, 1]). The vector is used as-is, so soft values train directly with BCE and are binarized at 0.5 for the evaluation metrics. Class maps for multi-hot targets must be keyed by source class index.

Without the flag every list is read as class indices, since an integer multi-hot vector is indistinguishable from an index list.

The plant pathology dataset provides this format: an image column and a labels column containing a sequence of ClassLabel values. Its six classes are complex, frog_eye_leaf_spot, healthy, powdery_mildew, rust, and scab. For example, [5, 1, 0] becomes [1, 1, 0, 0, 0, 1].

python train.py --dataset hfds/timm/plant-pathology-2021 \
  --task multilabel --target-key labels --num-classes 6 \
  --model resnet50 --pretrained --batch-size 32 --epochs 30 \
  --opt adamw --lr 1e-4 --weight-decay 0.01 --amp

python validate.py --dataset hfds/timm/plant-pathology-2021 \
  --task multilabel --target-key labels --num-classes 6 \
  --model resnet50 --checkpoint ./output/train/<experiment>/model_best.pth.tar

hfids/timm/plant-pathology-2021 supports streaming with the same target format. Reader support:

Hugging Face hfds and hfids support multi-label targets.
WebDataset (wds) supports lists of class IDs stored in each sample’s .json sidecar under the field named by --target-key. Without --target-key, the .cls file supplies a single class index.
Folder and tar readers yield single class indices.

Datasets that store one binary field per label, such as Wake Vision, use these settings:

--target-key person,depiction,body_part selects a comma separated list of fields. Each field becomes one class in the order listed, and values greater than zero count as positive. This works for the hfds, hfids, and tfds readers, and for WebDataset when the fields are in the .json sidecar.
--num-classes must equal the number of fields.
--target-format stays at its default, since multi-field keys define their own format.
Nested field paths use /, for example --input-key input --target-key output/pollen_output,output/wasps_output for the TFDS bee_dataset. Single target keys accept the same paths. A literal field name containing / (such as TFDS image/filename) takes precedence over a path.

Standard and NaFlex loaders support multi-label targets, including mixup and cutmix; NaFlex training currently requires a map-style dataset such as hfds.

The task trains with binary cross-entropy with logits. Smoothing defaults to zero for multi-label training; setting --smoothing s moves each binary target to target * (1 - s) + 0.5 * s, applied by Mixup/CutMix when they are active and by the loss otherwise. Existing --bce-sum and --bce-target-thresh options also apply. Distillation and JSD loss are not supported by this task.

--loss selects other losses (bce is the multi-label default):

asl: Asymmetric Loss. It follows the same smoothing and Mixup/CutMix behavior as BCE and is computed in float32 so it is safe with --amp.
- --asl-gamma-neg sets focusing for negative targets (default 4).
- --asl-gamma-pos sets focusing for positive targets (default 1).
- --asl-clip sets the probability margin for negatives (default 0.05).
- --asl-reduction defaults to batchmean, summing the loss over classes and averaging over the batch. sum restores the original ASL reduction, which also sums over the batch so the effective learning rate scales with batch size.
twoway: Two-way Multi-Label Loss, contrasting positive and negative logits across the classes of each sample and the samples of each class. Its logits are ranking scores, so mAP is the meaningful metric; F1 at a sigmoid threshold is not calibrated.
- --twoway-tp sets the positive logit temperature (default 4).
- --twoway-tn sets the negative logit temperature (default 1).
zlpr: ZLPR, a hyperparameter-free multi-label extension of softmax cross-entropy with 0 as the built-in logit threshold.
db: Distribution-Balanced Loss for long-tailed multi-label data. BCE with per-sample re-balancing weights that account for label co-occurrence, negative-tolerant regularization of the logits, and a focal term. It supports Mixup/CutMix and smoothing, with defaults from the reference DB configuration:
- --db-neg-scale sets the negative logit scale (default 2).
- --db-init-bias sets the class-prior logit shift factor (default 0.05).
- --db-focal-gamma sets focal focusing (default 2).
- --class-stats supplies the required training label counts as a class stats file (below).

twoway and zlpr need binary targets, so they do not support Mixup/CutMix or --smoothing.

For skewed label distributions, bce and asl support two weighting flags:

--loss-pos-weight weights the positive term of each class, balancing positives against negatives within a class. This also shifts predictions towards positive (higher recall at a fixed threshold, less calibrated probabilities).
--loss-class-weight weights the whole loss of each class, balancing the classes against each other and keeping calibration.

Each flag takes explicit weights (a single value, a comma separated list with one value per class, or a file) or a method name that derives the weights from the training label counts given with --class-stats:

--loss-pos-weight neg_pos: negatives / positives per class.
--loss-class-weight inv_freq: inverse class frequency, normalized to a mean of 1.
--loss-class-weight effective_num: class-balanced weights from the effective number of samples, normalized to a mean of 1.

Derived-weight options:

--loss-weight-power sets the exponent for neg_pos and inv_freq (default 1), with 0.5 for a softer square root re-weighting.
--loss-weight-beta sets beta for effective_num weights (default 0.999).
--loss-weight-max caps the derived weights.

Both weighting flags can use the same stats. class_weights.py counts the labels of a split, with the same dataset arguments as train.py, and writes the class stats file (--method previews the weights of a method):

python class_weights.py --dataset hfds/timm/plant-pathology-2021 --split train \
  --task multilabel --target-key labels --num-classes 6 --output stats.json

python train.py --dataset hfds/timm/plant-pathology-2021 \
  --task multilabel --target-key labels --num-classes 6 \
  --class-stats stats.json --loss-pos-weight neg_pos --loss-weight-power 0.5 ...

Deprecated flags:

Old flag	Replacement
`--bce-loss`	`--loss bce`
`--jsd-loss`	`--loss jsd`
`--bce-pos-weight`	`--loss-pos-weight`

Validation reports BCE loss, map (mean average precision), micro_f1, macro_f1, and sample_f1. Metrics are percentages and use the complete validation set, including predictions gathered across distributed ranks. Ordered distributed sampler padding is excluded. Exact mAP retains predictions on CPU until evaluation finishes, so its memory use scales with the number of examples and classes. Classes without positives contribute zero AP; F1 values with a zero denominator are zero.

--eval-metric map is the default for checkpoint selection and scheduling. F1 uses a sigmoid probability threshold of 0.5, configurable with --multilabel-threshold; mAP is independent of this threshold. Model outputs remain logits, so apply sigmoid for independent class probabilities when using the trained model.

Historical ImageNet recipes

These earlier runs remain useful references for RMSProp, SGD, and AugMix/JSD training. Hardware and reported metrics describe the original runs. For newer AdamW recipes, start with the examples above or the gist collections.

EfficientNet-B2 with RandAugment - 80.4 top-1, 95.1 top-5

The original run used dual Titan RTX cards with NVIDIA Apex; current --amp uses native PyTorch AMP:

./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 with RandAugment - 80.5 top-1, 94.9 top-5

The original run used dual Titan RTX cards with NVIDIA Apex; current --amp uses native PyTorch AMP:

./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 and SE-ResNeXt-26-T

These hparams (or similar) work well for a wide range of ResNet architecture, generally a good idea to increase the epoch # as the model size increases… ie approx 180-200 for ResNe(X)t50, and 220+ for larger. Increase batch size and LR proportionally for better GPUs or with AMP enabled. These params were for 2 1080Ti cards:

./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 with RandAugment - 81.5 top-1, 95.7 top-5

The training of this model started with the same command line as EfficientNet-B2 w/ RA above. After almost three weeks of training the process crashed. The results weren’t looking amazing so I resumed the training several times with tweaks to a few params (increase RE prob, decrease rand-aug, increase ema-decay). Nothing looked great. I ended up averaging the best checkpoints from all restarts. The result is mediocre at default res/crop but oddly performs much better with a full image test crop of 1.0.

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

Michael Klachko achieved these results with the command line for B2 adapted for larger batch size, with the recommended B0 dropout rate of 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 with JSD loss and RandAugment (clean + 2x RA augs) - 79.04 top-1, 94.39 top-5

Trained on two older 1080Ti cards, this took a while. Only slightly, non statistically better ImageNet validation result than my first good AugMix training of 78.99. However, these weights are more robust on tests with ImageNetV2, ImageNet-Sketch, etc. Unlike my first AugMix runs, I’ve enabled SplitBatchNorm, disabled random erasing on the clean split, and cranked up random erasing prob on the 2 augmented paths.

./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 --loss jsd --dist-bn reduce

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

Trained by Andrew Lavin with 8 V100 cards. Model EMA was not used, final checkpoint is the average of 8 best checkpoints during training.

./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 w/ RandAugment - 79.762 top-1, 94.60 top-5

These params will also work well for SE-ResNeXt-50 and SK-ResNeXt-50 and likely 101. I used them for the SK-ResNeXt-50 32x4d that I trained with 2 GPU using a slightly higher LR per effective batch size (lr=0.18, b=192 per GPU). The cmd line below are tuned for 8 GPU training.

./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

Update on GitHub