GPUNet is a new family of Convolutional Neural Networks crafted by NVIDIA AI.
This repo allows a user to easily train GPUNet, reproduce our results, test the accuracy of pre-trained checkpoints, and benchmark GPUNet latency. For customizing GPUNet, refer to Model customization.
- To get started, clone the repo:
git clone https://github.com/NVIDIA/DeepLearningExamples
cd DeepLearningExamples/PyTorch/Classification/GPUNet
-
Download ImageNet from the offical website. Recursively unzip the dataset, and locate the train and val folders. Refer to Prepare the dataset for more details.
-
Build and run the GPUNet PyTorch container, assuming you have installed the docker.
docker build -t gpunet .
docker run --gpus all -it --rm --network=host --shm-size 600G --ipc=host -v /path/to/imagenet:/root/data/imagenet/ gpunet
Prepare the Dataset
-
Extract the training data:
mkdir train && mv ILSVRC2012_img_train.tar train/ && cd train
tar -xvf ILSVRC2012_img_train.tar && rm -f ILSVRC2012_img_train.tar
find . -name "*.tar" | while read NAME ; do mkdir -p "${NAME%.tar}"; tar -xvf "${NAME}" -C "${NAME%.tar}"; rm -f "${NAME}"; done
cd ..
- Extract the validation data and move the images to subfolders:
mkdir val && mv ILSVRC2012_img_val.tar val/ && cd val && tar -xvf ILSVRC2012_img_val.tar
wget -qO- https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh | bash
The directory where the train/ and val/ directories are placed is referred to as /path/to/imagenet/ in this document.
Training
We have provided the training launch scripts for you to reproduce the GPUNet accuracy by training from scratch. For example, a user can copy the launch script in GPUNet-0.train.params or the training hyper-parameters below to reproduce the accuracy.
GPUNet training hyperparameters:
- GPUNet-0
./train.sh 8 /root/data/imagenet/ --model gpunet_0 --sched step --decay-epochs 2.4 --decay-rate .97 --opt rmsproptf -b 192 --epochs 450 --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.3 --drop-connect 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --lr .06 --num-classes 1000 --enable-distill False --crop-pct 1.0 --img-size 320 --amp
- GPUNet-1
./train.sh 8 /root/data/imagenet/ --model gpunet_1 --sched step --decay-epochs 2.4 --decay-rate .97 --opt rmsproptf -b 192 --epochs 450 --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.3 --drop-connect 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --lr .06 --num-classes 1000 --enable-distill False --crop-pct 1.0 --img-size 288 --amp
- GPUNet-2
./train.sh 8 /root/data/imagenet/ --model gpunet_2 --sched step --decay-epochs 2.4 --decay-rate .97 --opt rmsproptf -b 192 --epochs 450 --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.3 --drop-connect 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --lr .06 --num-classes 1000 --enable-distill False --crop-pct 1.0 --img-size 384 --amp
- GPUNet-D1 with distillation
./train.sh 8 /root/data/imagenet/ --model gpunet_d1 --sched step --decay-epochs 2.4 --decay-rate .97 --opt rmsproptf -b 192 --epochs 450 --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.3 --drop-connect 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --lr .06 --num-classes 1000 --enable-distill True --crop-pct 1.0 --img-size 456 --amp --test-teacher False --teacher tf_efficientnet_b5_ns --teacher-img-size 456
- GPUNet-D2 with distillation
./train.sh 8 /root/data/imagenet/ --model gpunet_d2 --sched step --decay-epochs 2.4 --decay-rate .97 --opt rmsproptf -b 128 --epochs 450 --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.3 --drop-connect 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --lr .06 --num-classes 1000 --enable-distill True --crop-pct 1.0 --img-size 528 --amp --test-teacher False --teacher tf_efficientnet_b6_ns --teacher-img-size 528
- GPUNet-P0 with distillation
./train.sh 8 /root/data/imagenet/ --model gpunet_p0 --sched step --decay-epochs 2.4 --decay-rate 0.97 --opt rmsproptf -b 256 --epochs 450 --opt-eps 0.001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.3 --drop-connect 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --lr 0.08 --num-classes 1000 --enable-distill True --crop-pct 0.875 --img-size 224 --amp --test-teacher False --teacher tf_efficientnet_b2 --teacher-img-size 260
- GPUNet-P1 with distillation
./train.sh 8 /root/data/imagenet/ --model gpunet_p1 --sched step --decay-epochs 2.4 --decay-rate 0.97 --opt rmsproptf -b 256 --epochs 450 --opt-eps 0.001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.3 --drop-connect 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --lr 0.08 --num-classes 1000 --enable-distill True --crop-pct 0.875 --img-size 224 --amp --test-teacher False --teacher tf_efficientnet_b2 --teacher-img-size 260
You need to call train.sh to start the training, and here is an example of arguments to train.sh.
./train.sh 8 >>launch with 8 GPUs.
/root/data/imagenet/ >>path to the imagenet.
--model gpunet_d1 >>name of GPUNet.
--sched step >>stepwise learning rate scheduler.
--decay-epochs 2.4 >>epoch interval to decay LR.
--decay-rate .97 >>LR decay rate (default: 0.1).
--opt rmsproptf >>optimizer.
-b 192 >>batch size.
--epochs 450 >>total training epochs.
--opt-eps .001 >>optimizer epsilon.
-j 8 >>the number of threads for data loader.
--lr .06 >>learning rate.
--warmup-lr 1e-6 >>warmup learning rate.
--weight-decay 1e-5 >>weight-decay rate.
--drop 0.3 >>dropout rate.
--drop-connect 0.2 >>drop connect rate.
--model-ema >>enable tracking moving average of model weights.
--model-ema-decay 0.9999 >>decay factor for model weights moving average (default: 0.9998).
--aa rand-m9-mstd0.5 >>using the random augmentation.
--remode pixel >>random erase mode.
--reprob 0.2 >>random erase prob.
--num-classes 1000 >>the number of output classes.
--amp >>enable the amp training.
--crop-pct 1.0 >>input image center crop percent.
--output ./output/ >>path to output folder.
--img-size 456 >>image size for the student model, i.e., gpunet_d1.
--enable-distill True >>to turn on/off the distillation.
--test-teacher False >>to test the accuracy of teacher model
--teacher tf_efficientnet_b5 >>the name of teacher model
--teacher-img-size 456 >>the image size to the teacher model. Note the student and teacher may have different image resolutions.
Training with distillation
We recommend running the distillation on a GPU with large DRAM; for example, 80G A100, since it needs to fit another teacher network.
- The following describes the usage of distillation.
- --enable-distill [Boolean]
- Enable or disable distillation.
- --teacher [String]
- Specify the name of the teacher model from Timm. You can find the full list of teacher models here, and you can pick any models from the model column. We expect the teacher model to be larger and better than the selected GPUNet.
- --teacher-img-size [Int]
- Specify the image resolution to the teacher model. The teacher model may use a larger resolution, and you can find the teacher's resolution from the column img_size to the selected teacher. Internally we use one data loader for both the teacher and student models, and we downsample an image batch with the teacher's resolution to the student's resolution using Interpolation.
- --enable-distill [Boolean]
Inference
We also allow a user to evaluate the accuracy of pre-trained GPUNet checkpoints and benchmark the model's TensorRT latency. For evaluating GPUNet on a custom dataset, refer to Train on your data.
Evaluate the pre-trained GPUNet checkpoints
In the eval.py, we have listed seven configurations to the released GPUNet models in the table below.
| batch | Distillation | GPU | Latency |
|---|---|---|---|
| 1 | No | GV100 | 0.65ms |
| 1 | No | GV100 | 0.85ms |
| 1 | No | GV100 | 1.75ms |
| 1 | Yes | GV100 | 0.5ms-D |
| 1 | Yes | GV100 | 0.8ms-D |
| 1 | Yes | GV100 | 1.25ms-D |
| 1 | Yes | GV100 | 2.25ms-D |
A user can easily evaluate the accuracy of a pre-trained checkpoint using the following code:
from configs.model_hub import get_configs, get_model_list
from models.gpunet_builder import GPUNet_Builder
modelJSON, cpkPath = get_configs(batch=1, latency="0.65ms", gpuType="GV100") >>Get the model configurations and checkpoints.
builder = GPUNet_Builder() >>Build an instance of GPUNet constructor.
model = builder.get_model(modelJSON) >>Build the GPUNet based on the model json.
builder.export_onnx(model) >>Export Pytorch model to ONNX for benchmarking the latency.
builder.test_model( >>Test the checkpoint accuracy.
model,
testBatch=200,
checkpoint=cpkPath,
imgRes=(3, model.imgRes, model.imgRes),
dtype="fp16",
crop_pct=1,
val_path="/root/data/imagenet/val",
)
Benchmark the GPUNet latency
We will need the ONNX file of the GPUNet model to reproduce the latency. builder.export_onnx(model) will export an ONNX file named gpunet.onnx. You can get the FP16 latency with the following command:
trtexec --onnx=gpunet.onnx --fp16 --workspace=10240
Here gpunet.onnx is configured to benchmark the latency at the batch = 1 to be consistent with the GPUNet paper. You can also look at the torch.onnx API to benchmark different settings, such as batch sizes. Finally, we report the median GPUNet compute time; here is an example output of a network with batch=1, latency=0.65ms, gpuType=GV100.
[04/07/2022-19:40:17] [I] GPU Compute Time: min = 0.554077 ms, max = 0.572388 ms, mean = 0.564606 ms, median = 0.564209 ms, percentile(99%) = 0.570312 ms