Package: darklight
Project DarkLightEfficient Knowledge Distillation in neural networks using TensorRT inference on teacher networkBackgroundKnowledge Distillation (KD) refers to the practice of using the outputs of a large teacher network train a (usually) smaller student network. This project leverages TensorRT to accelerate this process. It is common practice in KD, especiallydark knowledgetype techniques to pre-compute the logits from the teacher network and save them to disk. For training the student network, the pre-computed logits are used as is to teach the student. This saves GPU resources as one does not need to load the large teacher network to GPU memory during training.ProblemInA good teacher is patient and consistent, Beyer et. al. introduce the function matching approach for distilling the knowledge in a neural network. In this approach, rather than pre-computing the outputs from the teacher network, they are computed on the fly during training on the exact same input as seen by the student. However, this requires that the teacher model must share the GPU memory and compute resources and leads to the following question:How to achieve the best teacher inference and student training performance on a GPU?SolutionUse TensorRT to set up an inference engine and perform blazing fast inferenceUse logits from TensorRT inference to train the student network.Note that TensorRT works only on NVIDIA GPUs. AMD, Intel GPUs or TPUs are not supported.EnvironmentInstall withpipDarkLight can now be installed via PyPi (pip) withpipinstalldarklightHowever, this will not install TensorRT, which can be installed with either of these methods.Recommended method
This project uses pytorch CUDA, tensorrt>=8.0, opencv and pycuda. The recommended way to get all these is to use an NGC docker container with a recent version of PyTorch.sudodockerrun-it--ipc=host--net=host--gpusallnvcr.io/nvidia/pytorch:22.08-py3/bin/bash#if you want to load an external disk to the container, use the --volume switch#Once the container is up and running, install pycudapipinstallpycudadarklightCustom envIf you want to use your own environment with PyTorch, you need to get TensorRT and pycuda.Follow theofficial guideto download TensorRT deb file and install it with the script provided in this repo. Finally install pycudagitclonehttps://github.com/dataplayer12/darklight.gitcddarklight
bashinstall_trt.sh# if needed modify the version of deb file in the script before running.# This script will also install pycuda# this might fail for a number of reasons which is why NGC container is recommendedHow to useHere is a minimal example demonstrating the use ofresnet152to trainresnet18.importdarklightasdlimporttorchfromtorchvisionimportmodelsteacher=models.resnet152(pretrained=True)#substitue these with any other models you writestudent=models.resnet18(pretrained=False)dl.exportonnx(teacher,'rn152.onnx',bsize=1,hw=[224,224])delteacher#free up CPU or GPU memory used by teacher pytorch modeldm=dl.ImageNetManager('/sfnvme/imagenet/',size=[224,224],bsize=128)opt_params={'optimizer':torch.optim.AdamW,'okwargs':{'lr':1e-4,'weight_decay':0.05},'scheduler':torch.optim.lr_scheduler.CosineAnnealingWarmRestarts,'skwargs':{'T_0':10,'T_mult':2},'amplevel':None}#TensorRT inference engine is constructed for the teacher from onnx file#CUDA stream scope ensures interoperability between pycuda, TensorRT and pytorchtrainer=dl.StudentTrainer(student,dm,'rn152.onnx',opt_params=opt_params)trainer.train(epochs=50,save='dltest_{}.pth')Release Notesv0.1: Alpha of alphav0.2: Working libraryv0.3: Documentation addedv0.4: Support for multi-GPU training and stream creation is not requiredStatusCurrently supports image classification KD.The core TensorRT functionality works well (can also be used for pure inference)TensorRT accelerated training is verified (accelerate inference on teacher network with TRT)Implemented Soft Target Loss byHinton et. al.Implemented Hard Label Distillation byTouvron et. al.Immediate ToDosImprove TRT inference and training by transfering input only once.Benchmark dynamic shapes on TRTBenchmark PyTorch v/s TensorRT inference speed/memoryBetter documentationBetter unit testsMake PyPi packageRoadmapThis project will support KD on semantic segmentationKD support for object detection is planned.Teacher inference with other backends (OpenVINO,MIVisionX, OpenCV DNN module) planned but are not a high priority.