Knowledge Distillation - Tensorflow

March 12, 2018 ยท View on GitHub

This is an implementation for the basic idea behind Hinton's Knowledge Distillation Paper. We do not reproduce the exact results but rather show that the idea works.

While a few other implementations are available, the code flow is not very intuitive. Here we generate the soft targets from the teacher in an on-line manner while training the student network.

The big and small models (with some modification - We currently have a simple softmax regression as in TF's tutorial) have been taken from here.

While this may not (or may) be a good way to implement the distillation architecture, it leads to a good improvement in the (small) student model. In case you find any bug or have any suggestions feel free to create an issue or even send in a pull request.

Requirements

Tensorflow 1.3 or above

Running the code

Train the Teacher Model

 python main.py --model_type teacher --checkpoint_dir teachercpt --num_steps 5000 --temperature 5
 

Train the Student Model (in a standalone manner for comparison)

 python main.py --model_type student --checkpoint_dir studentcpt --num_steps 5000
 

Train the Student Model (Using Soft Targets from the teacher model)

 python main.py --model_type student --checkpoint_dir studentcpt --load_teacher_from_checkpoint true --load_teacher_checkpoint_dir teachercpt --num_steps 5000 --temperature 5
 

Results (For different temperature values)

ModelAccuracy - 2Accuracy - 5
Teacher Only97.998.12
Distillation89.1490.77
Student Only88.8488.84

The small model when trained without the soft labels always use temperature=1.

References

Distilling the Knowledge in a Neural Network