# 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](https://github.com/pytorch/examples). I have added significant functionality over time, including CUDA specific performance enhancements based on
[NVIDIA's APEX Examples](https://github.com/NVIDIA/apex/tree/master/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](hparams#mobilenetv4-conv-medium) as `mobilenetv4.yaml`:

```bash
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](hparams#adapting-a-recorded-config).

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

```bash
# 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 `--`:

```bash
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`:

```bash
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:

```bash
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:

```bash
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):

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

To run inference from a checkpoint:

```bash
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](hparams#selected-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](https://gist.github.com/rwightman/f6705cb65c03daeebca8aa129b1b94ad#file-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](https://huggingface.co/timm/mobilenetv4_conv_medium.e500_r256_in1k).

```bash
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](https://gist.github.com/rwightman/fb37c339efd2334177ff99a8083ebbc4#file-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](https://huggingface.co/timm/vit_medium_patch16_reg4_gap_256.sbb_in1k).

```bash
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](https://gist.github.com/rwightman/ee0b02c1e99a0761264d1d1319e95e5b#file-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](https://huggingface.co/timm/convnext_tiny.in12k_ft_in1k).

```bash
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](https://huggingface.co/datasets/timm/plant-pathology-2021) 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]`.

```bash
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](https://www.tensorflow.org/datasets/catalog/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](https://arxiv.org/abs/2009.14119). 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](https://staff.aist.go.jp/takumi.kobayashi/publication/2023/CVPR2023.pdf),
  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](https://arxiv.org/abs/2208.02955), a hyperparameter-free multi-label extension of softmax cross-entropy with 0 as the built-in logit threshold.
* `db`: [Distribution-Balanced Loss](https://arxiv.org/abs/2007.09654) 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](https://arxiv.org/abs/1901.05555) 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):

```bash
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](hparams#config-file-gists).

#### 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:

```bash
./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:

```bash
./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:

```bash
./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](https://github.com/michaelklachko) achieved these results with the command line for B2 adapted for larger batch size, with the recommended B0 dropout rate of 0.2.

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

```bash
./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](https://github.com/andravin) with 8 V100 cards. Model EMA was not used, final checkpoint is the average of 8 best checkpoints during training.

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

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

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