Enable AutoAugment and modernize DALI pipeline for ConvNets

Update DALI implementation to use modern "fn" API
instead of old class approach.

Add a codepath using AutoAugment in DALI training pipeline.
It can be easily extended to use other Automatic Augmentations.

The integration of DALI Pipeline with PyTorch additionally skips
the transposition when exposing NHWC data.

Extract the DALI implementation to separate file.
Update the readme and some configuration files for EfficientNet:
* dali-gpu is the default one, instead of PyTorch
* DALI supports AutoAugment (+ a mention of other Automatic Augmentations)

Fix a typo in the readme files:
--data-backends -> --data-backend

This PR is a backport of the changes made to this example, when it was
introduced into DALI codebase:
https://github.com/NVIDIA/DALI/tree/main/docs/examples/use_cases/pytorch/efficientnet

The changes were tested with the smallest EfficientNet only.

The usage od DALI GPU pipeline in the training can remove the CPU bottlneck
on both DGX-1V and DGX-A100 when running using AMP which was covered
in the blogpost:
https://developer.nvidia.com/blog/why-automatic-augmentation-matters/

Signed-off-by: Krzysztof Lecki <[email protected]>

Author: klecki

PyTorch/Classification/ConvNets/configs.yml

@@ -105,9 +105,9 @@ anchors:
         mixup: 0.2
         lr_schedule: cosine
         momentum: 0.9
-        warmup: 16 
+        warmup: 16
         epochs: 400
-        data_backend: pytorch
+        data_backend: dali-gpu
         augmentation: autoaugment
         num_classes: 1000
         interpolation: bicubic

PyTorch/Classification/ConvNets/efficientnet/README.md

@@ -92,7 +92,7 @@ This model uses SGD optimizer for B0 models and RMSPROP optimizer alpha=0.853  e
   * 0.0125 for 128 batch size for B0 models
   * 4.09e-06 for 32 batch size for B4 models
 scale the learning rate.
-* Learning rate schedule: 
+* Learning rate schedule:
   * cosine LR schedule for B0 models
   * linear LR schedule for B4 models
 * Weight decay (WD):
@@ -128,7 +128,7 @@ The following features are supported by this model:
 
 | Feature               | EfficientNet
 |-----------------------|--------------------------
-|[DALI](https://docs.nvidia.com/deeplearning/dali/release-notes/index.html)   |   Yes (without autoaugmentation)
+|[DALI](https://docs.nvidia.com/deeplearning/dali/release-notes/index.html)   |   Yes
 |[APEX AMP](https://nvidia.github.io/apex/amp.html) | Yes
 |[QAT](https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization)  |   Yes
 
@@ -143,10 +143,11 @@ We use [NVIDIA DALI](https://github.com/NVIDIA/DALI),
 which speeds up data loading when CPU becomes a bottleneck.
 DALI can use CPU or GPU, and outperforms the PyTorch native dataloader.
 
-Run training with `--data-backends dali-gpu` or `--data-backends dali-cpu` to enable DALI.
-For DGXA100 and DGX1 we recommend `--data-backends dali-cpu`.
+Run training with `--data-backend dali-gpu` or `--data-backend dali-cpu` to enable DALI.
+For DGXA100 and DGX1 we recommend `--data-backend dali-gpu`.
 
-DALI currently does not support Autoaugmentation, so for best accuracy it has to be disabled.
+[DALI supports AutoAugment](https://docs.nvidia.com/deeplearning/dali/user-guide/docs/auto_aug/auto_augment.html),
+as well as other [Automatic Augmentation schemes](https://docs.nvidia.com/deeplearning/dali/user-guide/docs/auto_aug/auto_aug.html).
 
 
 **[APEX](https://github.com/NVIDIA/apex)**
@@ -199,7 +200,7 @@ To enable mixed precision, you can:
 
 #### Enabling TF32
 
-TensorFloat-32 (TF32) is the new math mode in [NVIDIA A100](https://www.nvidia.com/en-us/data-center/a100/) GPUs for handling the matrix math also called tensor operations. TF32 running on Tensor Cores in A100 GPUs can provide up to 10x speedups compared to single-precision floating-point math (FP32) on Volta GPUs. 
+TensorFloat-32 (TF32) is the new math mode in [NVIDIA A100](https://www.nvidia.com/en-us/data-center/a100/) GPUs for handling the matrix math also called tensor operations. TF32 running on Tensor Cores in A100 GPUs can provide up to 10x speedups compared to single-precision floating-point math (FP32) on Volta GPUs.
 
 TF32 Tensor Cores can speed up networks using FP32, typically with no loss of accuracy. It is more robust than FP16 for models which require high dynamic range for weights or activations.
 
@@ -319,9 +320,9 @@ For example:
 You can download pre-trained weights from NGC:
 
 ```bash
-wget --content-disposition  -O 
+wget --content-disposition  -O
 
-unzip 
+unzip
 ```
 
 To run inference on ImageNet, run:
@@ -441,12 +442,12 @@ URL for each model can be found in the following table:
 
 | **Model** | **NGC weights URL** |
 |:---------:|:-------------------:|
-| efficientnet-b0 | https://api.ngc.nvidia.com/v2/models/nvidia/efficientnet_b0_pyt_amp/versions/20.12.0/files/nvidia_efficientnet-b0_210412.pth | 
-| efficientnet-b4 | https://api.ngc.nvidia.com/v2/models/nvidia/efficientnet_b4_pyt_amp/versions/20.12.0/files/nvidia_efficientnet-b4_210412.pth | 
-| efficientnet-widese-b0 | https://api.ngc.nvidia.com/v2/models/nvidia/efficientnet_widese_b0_pyt_amp/versions/20.12.0/files/nvidia_efficientnet-widese-b0_210412.pth | 
-| efficientnet-widese-b4 | https://api.ngc.nvidia.com/v2/models/nvidia/efficientnet_widese_b4_pyt_amp/versions/20.12.0/files/nvidia_efficientnet-widese-b4_210412.pth | 
-| efficientnet-quant-b0 | https://api.ngc.nvidia.com/v2/models/nvidia/efficientnet_b0_pyt_qat_ckpt_fp32/versions/21.03.0/files/nvidia-efficientnet-quant-b0-130421.pth | 
-| efficientnet-quant-b4 | https://api.ngc.nvidia.com/v2/models/nvidia/efficientnet_b4_pyt_qat_ckpt_fp32/versions/21.03.0/files/nvidia-efficientnet-quant-b4-130421.pth | 
+| efficientnet-b0 | https://api.ngc.nvidia.com/v2/models/nvidia/efficientnet_b0_pyt_amp/versions/20.12.0/files/nvidia_efficientnet-b0_210412.pth |
+| efficientnet-b4 | https://api.ngc.nvidia.com/v2/models/nvidia/efficientnet_b4_pyt_amp/versions/20.12.0/files/nvidia_efficientnet-b4_210412.pth |
+| efficientnet-widese-b0 | https://api.ngc.nvidia.com/v2/models/nvidia/efficientnet_widese_b0_pyt_amp/versions/20.12.0/files/nvidia_efficientnet-widese-b0_210412.pth |
+| efficientnet-widese-b4 | https://api.ngc.nvidia.com/v2/models/nvidia/efficientnet_widese_b4_pyt_amp/versions/20.12.0/files/nvidia_efficientnet-widese-b4_210412.pth |
+| efficientnet-quant-b0 | https://api.ngc.nvidia.com/v2/models/nvidia/efficientnet_b0_pyt_qat_ckpt_fp32/versions/21.03.0/files/nvidia-efficientnet-quant-b0-130421.pth |
+| efficientnet-quant-b4 | https://api.ngc.nvidia.com/v2/models/nvidia/efficientnet_b4_pyt_qat_ckpt_fp32/versions/21.03.0/files/nvidia-efficientnet-quant-b4-130421.pth |
 
 To run inference on ImageNet, run:
 
@@ -469,7 +470,7 @@ During the QAT process, evaluation is done in the same way as during standard tr
 
 or to evaluate a created checkpoint with the flag `--evaluate`:
 
-`python ./quant_main.py --arch efficientnet-quant-<version> --evaluate --epochs 1 --resume <path to checkpoint> -b <batch size> <path to imagenet>` 
+`python ./quant_main.py --arch efficientnet-quant-<version> --evaluate --epochs 1 --resume <path to checkpoint> -b <batch size> <path to imagenet>`
 
 It also can run on multi-GPU in an identical way as the standard `main.py` script:
 

PyTorch/Classification/ConvNets/image_classification/dali.py

@@ -0,0 +1,91 @@
+# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nvidia.dali import fn
+from nvidia.dali import types
+
+from nvidia.dali.pipeline.experimental import pipeline_def
+
+from nvidia.dali.auto_aug import auto_augment, trivial_augment
+
+
+@pipeline_def(enable_conditionals=True)
+def training_pipe(data_dir, interpolation, image_size, output_layout, automatic_augmentation,
+                  dali_device="gpu", rank=0, world_size=1):
+    rng = fn.random.coin_flip(probability=0.5)
+
+    jpegs, labels = fn.readers.file(name="Reader", file_root=data_dir, shard_id=rank,
+                                    num_shards=world_size, random_shuffle=True, pad_last_batch=True)
+
+    if dali_device == "gpu":
+        decoder_device = "mixed"
+        resize_device = "gpu"
+    else:
+        decoder_device = "cpu"
+        resize_device = "cpu"
+
+    # This padding sets the size of the internal nvJPEG buffers to be able to handle all images
+    # from full-sized ImageNet without additional reallocations
+    images = fn.decoders.image_random_crop(jpegs, device=decoder_device, output_type=types.RGB,
+                                           device_memory_padding=211025920,
+                                           host_memory_padding=140544512,
+                                           random_aspect_ratio=[0.75, 4.0 / 3.0],
+                                           random_area=[0.08, 1.0])
+
+    images = fn.resize(images, device=resize_device, size=[image_size, image_size],
+                       interp_type=interpolation, antialias=False)
+
+    # Make sure that from this point we are processing on GPU regardless of dali_device parameter
+    images = images.gpu()
+
+    images = fn.flip(images, horizontal=rng)
+
+    # Based on the specification, apply the automatic augmentation policy. Note, that from the point
+    # of Pipeline definition, this `if` statement relies on static scalar parameter, so it is
+    # evaluated exactly once during build - we either include automatic augmentations or not.
+    # We pass the shape of the image after the resize so the translate operations are done
+    # relative to the image size.
+    if automatic_augmentation is None:
+        output = images
+    elif automatic_augmentation == "autoaugment":
+        output = auto_augment.auto_augment_image_net(images, shape=[image_size, image_size])
+    else:
+        raise ValueError(f"Automatic augmentation: '{automatic_augmentation}'"
+                         f" is not supported for DALI")
+
+
+    output = fn.crop_mirror_normalize(output, dtype=types.FLOAT, output_layout=output_layout,
+                                      crop=(image_size, image_size),
+                                      mean=[0.485 * 255, 0.456 * 255, 0.406 * 255],
+                                      std=[0.229 * 255, 0.224 * 255, 0.225 * 255])
+
+    return output, labels
+
+
+@pipeline_def
+def validation_pipe(data_dir, interpolation, image_size, image_crop, output_layout, rank=0,
+                    world_size=1):
+    jpegs, label = fn.readers.file(name="Reader", file_root=data_dir, shard_id=rank,
+                                   num_shards=world_size, random_shuffle=False, pad_last_batch=True)
+
+    images = fn.decoders.image(jpegs, device="mixed", output_type=types.RGB)
+
+    images = fn.resize(images, resize_shorter=image_size, interp_type=interpolation,
+                       antialias=False)
+
+    output = fn.crop_mirror_normalize(images, dtype=types.FLOAT, output_layout=output_layout,
+                                      crop=(image_crop, image_crop),
+                                      mean=[0.485 * 255, 0.456 * 255, 0.406 * 255],
+                                      std=[0.229 * 255, 0.224 * 255, 0.225 * 255])
+    return output, label