The following sections provide greater details of the dataset, running training and inference, and the training results.

Scripts and sample code

To run a non standard configuration use:

• For 1 GPU

  • FP32 python ./main.py --arch resnext101-32x4d -c fanin --label-smoothing 0.1 <path to imagenet> python ./main.py --arch resnext101-32x4d -c fanin --label-smoothing 0.1 --amp --static-loss-scale 256 <path to imagenet>

• For multiple GPUs

  • FP32 python ./multiproc.py --nproc_per_node 8 ./main.py --arch resnext101-32x4d -c fanin --label-smoothing 0.1 <path to imagenet>
  • AMP python ./multiproc.py --nproc_per_node 8 ./main.py --arch resnext101-32x4d -c fanin --label-smoothing 0.1 --amp --static-loss-scale 256 <path to imagenet>

Use python ./main.py -h to obtain the list of available options in the main.py script.

Command-line options:

To see the full list of available options and their descriptions, use the -h or --help command-line option, for example:

python main.py -h

usage: main.py [-h] [--data-backend BACKEND] [--arch ARCH]
               [--model-config CONF] [--num-classes N] [-j N] [--epochs N]
               [--run-epochs N] [-b N] [--optimizer-batch-size N] [--lr LR]
               [--lr-schedule SCHEDULE] [--warmup E] [--label-smoothing S]
               [--mixup ALPHA] [--momentum M] [--weight-decay W]
               [--bn-weight-decay] [--nesterov] [--print-freq N]
               [--resume PATH] [--pretrained-from-file PATH]
               [--static-loss-scale STATIC_LOSS_SCALE] [--dynamic-loss-scale]
               [--prof N] [--amp] [--seed SEED] [--gather-checkpoints]
               [--raport-file RAPORT_FILE] [--evaluate] [--training-only]
               [--no-checkpoints] [--checkpoint-filename CHECKPOINT_FILENAME]
               [--workspace DIR] [--memory-format {nchw,nhwc}]
               DIR

PyTorch ImageNet Training

positional arguments:
  DIR                   path to dataset

optional arguments:
  -h, --help            show this help message and exit
  --data-backend BACKEND
                        data backend: pytorch | syntetic | dali-gpu | dali-cpu
                        (default: dali-cpu)
  --arch ARCH, -a ARCH  model architecture: resnet18 | resnet34 | resnet50 |
                        resnet101 | resnet152 | resnext50-32x4d |
                        resnext101-32x4d | resnext101-32x8d |
                        resnext101-32x8d-basic | se-resnext101-32x4d (default:
                        resnet50)
  --model-config CONF, -c CONF
                        model configs: classic | fanin | grp-fanin | grp-
                        fanout(default: classic)
  --num-classes N       number of classes in the dataset
  -j N, --workers N     number of data loading workers (default: 5)
  --epochs N            number of total epochs to run
  --run-epochs N        run only N epochs, used for checkpointing runs
  -b N, --batch-size N  mini-batch size (default: 256) per gpu
  --optimizer-batch-size N
                        size of a total batch size, for simulating bigger
                        batches using gradient accumulation
  --lr LR, --learning-rate LR
                        initial learning rate
  --lr-schedule SCHEDULE
                        Type of LR schedule: step, linear, cosine
  --warmup E            number of warmup epochs
  --label-smoothing S   label smoothing
  --mixup ALPHA         mixup alpha
  --momentum M          momentum
  --weight-decay W, --wd W
                        weight decay (default: 1e-4)
  --bn-weight-decay     use weight_decay on batch normalization learnable
                        parameters, (default: false)
  --nesterov            use nesterov momentum, (default: false)
  --print-freq N, -p N  print frequency (default: 10)
  --resume PATH         path to latest checkpoint (default: none)
  --pretrained-from-file PATH
                        load weights from here
  --static-loss-scale STATIC_LOSS_SCALE
                        Static loss scale, positive power of 2 values can
                        improve amp convergence.
  --dynamic-loss-scale  Use dynamic loss scaling. If supplied, this argument
                        supersedes --static-loss-scale.
  --prof N              Run only N iterations
  --amp                 Run model AMP (automatic mixed precision) mode.
  --seed SEED           random seed used for numpy and pytorch
  --gather-checkpoints  Gather checkpoints throughout the training, without
                        this flag only best and last checkpoints will be
                        stored
  --raport-file RAPORT_FILE
                        file in which to store JSON experiment raport
  --evaluate            evaluate checkpoint/model
  --training-only       do not evaluate
  --no-checkpoints      do not store any checkpoints, useful for benchmarking
  --checkpoint-filename CHECKPOINT_FILENAME
  --workspace DIR       path to directory where checkpoints will be stored
  --memory-format {nchw,nhwc}
                        memory layout, nchw or nhwc

Dataset guidelines

To use your own dataset, divide it in directories as in the following scheme:

• Training images - train/<class id>/<image>
• Validation images - val/<class id>/<image>

If your dataset's has number of classes different than 1000, you need to pass --num-classes N flag to the training script.

Training process

All the results of the training will be stored in the directory specified with --workspace argument. Script will store:

• most recent checkpoint - checkpoint.pth.tar (unless --no-checkpoints flag is used).
• checkpoint with best validation accuracy - model_best.pth.tar (unless --no-checkpoints flag is used).
• JSON log - in the file specified with --raport-file flag.

Metrics gathered through training:

• train.loss - training loss
• train.total_ips - training speed measured in images/second
• train.compute_ips - training speed measured in images/second, not counting data loading
• train.data_time - time spent on waiting on data
• train.compute_time - time spent in forward/backward pass

To restart training from checkpoint use --resume option.

To start training from pretrained weights (e.g. downloaded from NGC) use --pretrained-from-file option.

The difference between those two is that the pretrained weights contain only model weights, and checkpoints, apart from model weights, contain optimizer state, LR scheduler state.

Checkpoints are suitable for dividing the training into parts, for example in order to divide the training job into shorter stages, or restart training after infrastructure fail.

Pretrained weights can be used as a base for finetuning the model to a different dataset, or as a backbone to detection models.

Inference process

Validation is done every epoch, and can be also run separately on a checkpointed model.

python ./main.py --arch resnext101-32x4d --evaluate --epochs 1 --resume <path to checkpoint> -b <batch size> <path to imagenet>

Metrics gathered through training:

• val.loss - validation loss
• val.top1 - validation top1 accuracy
• val.top5 - validation top5 accuracy
• val.total_ips - inference speed measured in images/second
• val.compute_ips - inference speed measured in images/second, not counting data loading
• val.data_time - time spent on waiting on data
• val.compute_time - time spent on inference

To run inference on JPEG image, you have to first extract the model weights from checkpoint:

python checkpoint2model.py --checkpoint-path <path to checkpoint> --weight-path <path where weights will be stored>

Then run classification script:

python classify.py --arch resnext101-32x4d --pretrained-from-file <path to weights from previous step> --precision AMP|FP32 --image <path to JPEG image>

You can also run ImageNet validation on pretrained weights:

python ./main.py --arch resnext101-32x4d --evaluate --epochs 1 --pretrained-from-file <path to pretrained weights> -b <batch size> <path to imagenet>

NGC Pretrained weights:

Pretrained weights can be downloaded from NGC:

wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/resnext101_32x4d_pyt_amp/versions/20.06.0/zip -O resnext101_32x4d_pyt_amp_20.06.0.zip

unzip resnext101_32x4d_pyt_amp_20.06.0.zip

To run inference on ImageNet, run:

python ./main.py --arch resnext101-32x4d --evaluate --epochs 1 --pretrained-from-file nvidia_resnext101-32x4d_200821.pth.tar -b <batch size> <path to imagenet>

To run inference on JPEG image using pretrained weights:

python classify.py --arch resnext101-32x4d --pretrained-from-file nvidia_resnext101-32x4d_200821.pth.tar --precision AMP|FP32 --image <path to JPEG image>