Implementation of Teacher-Student Model in PyTorch

With a pre-trained “teacher” network, teacher-student training is a method for accelerating training and enhancing the convergence of a neural network. It is widely used to train smaller, less expensive networks from more expensive, larger ones since it is both popular and effective. In a previous post, we discussed the concept of Knowlege distillation as the idea behind the Teacher-Student model. In this post, we’ll discuss the fundamentals of teacher-student training, demonstrate how to do it in PyTorch, and examine the results of using this approach. If you’re not familiar with softmax cross entropy, our introduction to it might be a helpful pre-read. This is a part of our series on training targets.

Main Concept

The concept is basic. Start by training a sizable neural network (the teacher) with training data as per normal. Then, build a second, smaller network (the student), and train it to replicate the teacher’s outcomes. For instance, teacher preparation might look like this:

for (batch_idx, batch) in enumerate(train_ldr):
    X = batch[0]  # the predictors / inputs
    Y = batch[1]  # the targets 
    out = teacher(X) 
. . .

But training the student looks like:

for (batch_idx, batch) in enumerate(train_ldr):
    X = batch[0]    # the…