KnowledgeDistillation Layer (Caffe implementation)

This is a CPU implementation of knowledge distillation in Caffe.
This code is heavily based on softmax_loss_layer.hpp and softmax_loss_layer.cpp.

Please refer to the paper

Hinton, G. Vinyals, O. and Dean, J. Distilling knowledge in a neural network. 2015.

Installation

1. Install Caffe in your directory CAFFE
   
2. Clone this repository in your directory ROOT
   

   cd $ROOT
   git clone https://github.com/wentianli/knowledge_distillation_caffe.git

3. Move files to your Caffe folder
   

   cp $ROOT/knowledge_distillation_layer.hpp $CAFFE/include/caffe/layers
   cp $ROOT/knowledge_distillation_layer.cpp $CAFFE/src/caffe/layers

4. Modify $CAFFE/src/caffe/proto/caffe.proto
   add optional KnowledgeDistillationParameter in LayerParameter

   message LayerParameter {
     ...
   
   //next available layer-specific ID
   optional KnowledgeDistillationParameter knowledge_distillation_param = 147;
   }

   
   add message KnowledgeDistillationParameter
   

   message KnowledgeDistillationParameter {
   optional float temperature = 1 [default = 1];
   }

5. Build Caffe
   

Usage

KnowledgeDistillation Layer has one specific parameter temperature.

The layer takes 2 or 3 input blobs:
bottom[0]: the logits of the student
bottom[1]: the logits of the teacher
bottom[2](optional): label inputs
The logits are first divided by temperatrue T, then mapped to probability distributions over classes using the softmax function. The layer computes KL divergence instead of cross entropy. The gradients are multiplied by T^2, as suggested in the paper.

1. Common setting in prototxt (2 input blobs are given)

   layer {
   name: "KD"
   type: "KnowledgeDistillation"
   bottom: "student_logits"
   bottom: "taecher_logits"
   top: "KL_div"
   include { phase: TRAIN }
   knowledge_distillation_param { temperature: 4 } #usually larger than 1
   loss_weight: 1
   }

2. If you have ignore_label, 3 input blobs should be given

   
   layer {
   name: "KD"
   type: "KnowledgeDistillation"
   bottom: "student_logits"
   bottom: "taecher_logits"
   bottom: "label"
   top: "KL_div"
   include { phase: TRAIN }
   knowledge_distillation_param { temperature: 4 }
   loss_param {ignore_label: 2}
   loss_weight: 1
   }