A library for accelerating knowledge distillation training
Project description
Project DarkLight
Efficient Knowledge Distillation in neural networks using TensorRT inference on teacher network
Background
Knowledge 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, especially dark knowledge type 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.
Problem
In A 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?
Solution
- Use TensorRT to set up an inference engine and perform blazing fast inference
- Use 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.
Environment
- Install with
pip
DarkLight can now be installed via PyPi (pip) with
pip install darklight
However, 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.
sudo docker run -it --ipc=host --net=host --gpus all nvcr.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 pycuda
pip install pycuda darklight
- Custom env
If you want to use your own environment with PyTorch, you need to get TensorRT and pycuda.
Follow the official guide to download TensorRT deb file and install it with the script provided in this repo. Finally install pycuda
git clone https://github.com/dataplayer12/darklight.git
cd darklight
bash install_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 recommended
How to use
Here is a minimal example demonstrating the use of resnet152
to train resnet18
.
import darklight as dl
import torch
from torchvision import models
teacher=models.resnet152(pretrained=True) #substitue these with any other models you write
student= models.resnet18(pretrained=False)
dl.exportonnx(teacher, 'rn152.onnx', bsize=1, hw=[224,224])
del teacher #free up CPU or GPU memory used by teacher pytorch model
dm=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
}
stream=torch.cuda.Stream()
with torch.cuda.stream(stream):
#TensorRT inference engine is constructed for the teacher from onnx file
#CUDA stream scope ensures interoperability between pycuda, TensorRT and pytorch
trainer=dl.StudentTrainer(student, dm, 'rn152.onnx', opt_params=opt_params)
trainer.train(epochs=50, save='dltest_{}.pth')
Status
- Currently 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 by Hinton et. al.
- Implemented Hard Label Distillation by Touvron et. al.
Immediate ToDos
- Improve TRT inference and training by transfering input only once.
- Benchmark dynamic shapes on TRT
- Benchmark PyTorch v/s TensorRT inference speed/memory
- Better documentation
- Better unit tests
- Make PyPi package
Roadmap
- This project will support KD on semantic segmentation
- KD support for object detection is planned.
- Teacher inference with other backends (OpenVINO, MIVisionX, OpenCV DNN module) planned but are not a high priority.
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distributions
Built Distribution
File details
Details for the file darklight-0.3-py3-none-any.whl
.
File metadata
- Download URL: darklight-0.3-py3-none-any.whl
- Upload date:
- Size: 37.8 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/4.0.1 CPython/3.8.10
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | 4b3b7adfb4d8e86bb42b9685fa82650ad1320671fa63412ea0486709c596f1fa |
|
MD5 | a804b03be58d4536be5c36c6c8297fa1 |
|
BLAKE2b-256 | f074c39207ba2278aa4c68d48706273f483387285fbed4ce1b0a38c508034246 |