With modified architecture and initialization this ResNet50 version gives ~0.5% better accuracy than original.
The following sections provide greater details of the dataset, running training and inference, and the training results.
Scripts and sample code
In the root directory, the most important files are:
main.py: the script that controls the logic of training and validation of the ResNet-like modelsDockerfile: Instructions for Docker to build a container with the basic set of dependencies to run ResNet like models for image classificationrequirements.txt: a set of extra Python requirements for running ResNet-like models
The model/ directory contains the following modules used to define ResNet family models:
resnet.py: the definition of ResNet, ResNext, and SE-ResNext modelblocks/conv2d_block.py: the definition of 2D convolution blockblocks/resnet_bottleneck_block.py: the definition of ResNet-like bottleneck blocklayers/*.py: definitions of specific layers used in the ResNet-like model
The utils/ directory contains the following utility modules:
cmdline_helper.py: helper module for command line processingdata_utils.py: module defining input data pipelinesdali_utils.py: helper module for DALIimage_processing.py: image processing and data augmentation functionslearning_rate.py: definition of used learning rate scheduleoptimizers.py: definition of used custom optimizershooks/*.py: definitions of specific hooks allowing logging of training and inference process
The runtime/ directory contains the following module that define the mechanics of the training process:
runner.py: module encapsulating the training, inference and evaluation
Parameters
The main.py script
The script for training and evaluating the ResNet-50 v1.5 model has a variety of parameters that control these processes.
usage: main.py [-h] [--arch {resnet50,resnext101-32x4d,se-resnext101-32x4d}]
[--mode {train,train_and_evaluate,evaluate,predict,training_benchmark,inference_benchmark}]
[--export_dir EXPORT_DIR] [--to_predict TO_PREDICT]
--batch_size BATCH_SIZE [--num_iter NUM_ITER]
[--run_iter RUN_ITER] [--iter_unit {epoch,batch}]
[--warmup_steps WARMUP_STEPS] [--model_dir MODEL_DIR]
[--results_dir RESULTS_DIR] [--log_filename LOG_FILENAME]
[--display_every DISPLAY_EVERY] [--seed SEED]
[--gpu_memory_fraction GPU_MEMORY_FRACTION] [--gpu_id GPU_ID]
[--finetune_checkpoint FINETUNE_CHECKPOINT] [--use_final_conv]
[--quant_delay QUANT_DELAY] [--quantize] [--use_qdq]
[--symmetric] [--data_dir DATA_DIR]
[--data_idx_dir DATA_IDX_DIR] [--dali]
[--synthetic_data_size SYNTHETIC_DATA_SIZE] [--lr_init LR_INIT]
[--lr_warmup_epochs LR_WARMUP_EPOCHS]
[--weight_decay WEIGHT_DECAY] [--weight_init {fan_in,fan_out}]
[--momentum MOMENTUM] [--label_smoothing LABEL_SMOOTHING]
[--mixup MIXUP] [--cosine_lr] [--xla]
[--data_format {NHWC,NCHW}] [--amp]
[--static_loss_scale STATIC_LOSS_SCALE]
JoC-RN50v1.5-TF
optional arguments:
-h, --help show this help message and exit.
--arch {resnet50,resnext101-32x4d,se-resnext101-32x4d}
Architecture of model to run.
--mode {train,train_and_evaluate,evaluate,predict,training_benchmark,inference_benchmark}
The execution mode of the script.
--export_dir EXPORT_DIR
Directory in which to write exported SavedModel.
--to_predict TO_PREDICT
Path to file or directory of files to run prediction
on.
--batch_size BATCH_SIZE
Size of each minibatch per GPU.
--num_iter NUM_ITER Number of iterations to run.
--run_iter RUN_ITER Number of training iterations to run on single run.
--iter_unit {epoch,batch}
Unit of iterations.
--warmup_steps WARMUP_STEPS
Number of steps considered as warmup and not taken
into account for performance measurements.
--model_dir MODEL_DIR
Directory in which to write model. If undefined,
results dir will be used.
--results_dir RESULTS_DIR
Directory in which to write training logs, summaries
and checkpoints.
--log_filename LOG_FILENAME
Name of the JSON file to which write the training log.
--display_every DISPLAY_EVERY
How often (in batches) to print out running
information.
--seed SEED Random seed.
--gpu_memory_fraction GPU_MEMORY_FRACTION
Limit memory fraction used by training script for DALI.
--gpu_id GPU_ID Specify ID of the target GPU on multi-device platform.
Effective only for single-GPU mode.
--finetune_checkpoint FINETUNE_CHECKPOINT
Path to pre-trained checkpoint which will be used for
fine-tuning.
--use_final_conv Use convolution operator instead of MLP as last layer.
--quant_delay QUANT_DELAY
Number of steps to be run before quantization starts
to happen.
--quantize Quantize weights and activations during training.
(Defaults to Assymmetric quantization)
--use_qdq Use QDQV3 op instead of FakeQuantWithMinMaxVars op for
quantization. QDQv3 does only scaling.
--symmetric Quantize weights and activations during training using
symmetric quantization.
Dataset arguments:
--data_dir DATA_DIR Path to dataset in TFRecord format. Files should be
named 'train-*' and 'validation-*'.
--data_idx_dir DATA_IDX_DIR
Path to index files for DALI. Files should be named
'train-*' and 'validation-*'.
--dali Enable DALI data input.
--synthetic_data_size SYNTHETIC_DATA_SIZE
Dimension of image for synthetic dataset.
Training arguments:
--lr_init LR_INIT Initial value for the learning rate.
--lr_warmup_epochs LR_WARMUP_EPOCHS
Number of warmup epochs for learning rate schedule.
--weight_decay WEIGHT_DECAY
Weight Decay scale factor.
--weight_init {fan_in,fan_out}
Model weight initialization method.
--momentum MOMENTUM SGD momentum value for the Momentum optimizer.
--label_smoothing LABEL_SMOOTHING
The value of label smoothing.
--mixup MIXUP The alpha parameter for mixup (if 0 then mixup is not
applied).
--cosine_lr Use cosine learning rate schedule.
Generic optimization arguments:
--xla Enable XLA (Accelerated Linear Algebra) computation
for improved performance.
--data_format {NHWC,NCHW}
Data format used to do calculations.
--amp Enable Automatic Mixed Precision to speedup
computation using tensor cores.
Automatic Mixed Precision arguments:
--static_loss_scale STATIC_LOSS_SCALE
Use static loss scaling in FP32 AMP.
Quantization Aware Training
Quantization Aware training (QAT) simulates quantization during training by quantizing weights and activation layers. This will help reduce the loss in accuracy when we convert the network trained in FP32 to INT8 for faster inference. QAT introduces additional nodes in the graph which will be used to learn the dynamic ranges of weights and activation layers. Tensorflow provides a quantization tool which automatically adds these nodes in-place. Typical workflow for training QAT networks is to train a model until convergence and then finetune with the quantization layers. It is recommended that QAT is performed on a single GPU.
- For 1 GPU
- Command:
sh resnet50v1.5/training/GPU1_RN50_QAT.sh <path to pre-trained ckpt dir> <path to dataset directory> <result_directory>
- Command:
It is recommended to finetune a model with quantization nodes rather than train a QAT model from scratch. The latter can also be performed by setting quant_delay parameter.
quant_delay is the number of steps after which quantization nodes are added for QAT. If we are fine-tuning, quant_delay is set to 0.
For QAT network, we use tf.quantization.quantize_and_dequantize operation.
These operations are automatically added at weights and activation layers in the RN50 by using tf.contrib.quantize.experimental_create_training_graph utility. Support for using tf.quantization.quantize_and_dequantize
operations for tf.contrib.quantize.experimental_create_training_graph has been added in TensorFlow 20.01-py3 NGC container and later versions, which is required for this task.
Post process checkpoint
postprocess_ckpt.py is a utility to convert the final classification FC layer into a 1x1 convolution layer using the same weights. This is required to ensure TensorRT can parse QAT models successfully.
This script should be used after performing QAT to reshape the FC layer weights in the final checkpoint.
Arguments:
* --input : Path to the trained checkpoint of RN50.
* --output : Name of the new checkpoint file which has the FC layer weights reshaped into 1x1 conv layer weights.
* --dense_layer : Name of the FC layer
Exporting Frozen graphs
To export frozen graphs (which can be used for inference with TensorRT), use:
python export_frozen_graph.py --checkpoint <path_to_checkpoint> --quantize --use_final_conv --use_qdq --symmetric --input_format NCHW --compute_format NCHW --output_file=<output_file_name>
Arguments:
--checkpoint: Optional argument to export the model with checkpoint weights.--quantize: Optional flag to export quantized graphs.--use_qdq: Use quantize_and_dequantize (QDQ) op instead of FakeQuantWithMinMaxVars op for quantization. QDQ does only scaling.--input_format: Data format of input tensor (Default: NCHW). Use NCHW format to optimize the graph with TensorRT.--compute_format: Data format of the operations in the network (Default: NCHW). Use NCHW format to optimize the graph with TensorRT.
Inference process
To run inference on a single example with a checkpoint and a model script, use:
python main.py --mode predict --model_dir <path to model> --to_predict <path to image> --results_dir <path to results>
The optional --xla and --amp flags control XLA and AMP during inference. To run inference using TF-TRT, please use the following command:
python inference.py --model <path to model> --tf-trt --batch-size <inference_batch_size> --data-dir <path to data>
The optional --amp flag controls AMP during inference.