PyTorch Distilling on Synthetic Data

The teacher scores the training rows once, under no_grad. Temperature 3 flattens each row before the student sees it.

teacher.eval()
with torch.no_grad():
    soft = F.softmax(teacher(x) / temp, dim=1)

KL loss compares the student's log-probabilities with the soft rows. The temperature squared factor restores the scale that the division shrank.

logp = F.log_softmax(model(x) / temp, dim=1)
loss = F.kl_div(logp, soft, reduction="batchmean") * (temp * temp)

import torch
import torch.nn.functional as F
from torch import nn


def draw_shape(label, g):
    rows = torch.arange(16).view(16, 1).float()
    cols = torch.arange(16).view(1, 16).float()
    if label == 0:
        centre = 8 + int(torch.randint(-2, 3, (1,), generator=g))
        mask = (rows >= centre - 1) & (rows <= centre + 1) & (cols >= 0)
    elif label == 1:
        cr = 7.5 + int(torch.randint(-1, 2, (1,), generator=g))
        cc = 7.5 + int(torch.randint(-1, 2, (1,), generator=g))
        dist = torch.sqrt((rows - cr) ** 2 + (cols - cc) ** 2)
        mask = (dist >= 4.5) & (dist <= 6.5)
    else:
        top = int(torch.randint(2, 7, (1,), generator=g))
        left = int(torch.randint(2, 7, (1,), generator=g))
        mask = (rows >= top) & (rows <= top + 7) & (cols >= left) & (cols <= left + 7)
    return mask.float()


def make_shapes(count, seed, noise):
    g = torch.Generator().manual_seed(seed)
    labels = torch.arange(count) % 3
    labels = labels[torch.randperm(count, generator=g)]
    images = [draw_shape(int(label), g) for label in labels]
    x = torch.stack(images).unsqueeze(1)
    x = x + noise * torch.randn(x.shape, generator=g)
    return x.flatten(1), labels


class Teacher(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(256, 64), nn.ReLU(), nn.Linear(64, 3))

    def forward(self, x):
        return self.net(x)


class Student(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(256, 16), nn.ReLU(), nn.Linear(16, 3))

    def forward(self, x):
        return self.net(x)


def accuracy(model, x, y):
    model.eval()
    with torch.no_grad():
        return float((model(x).argmax(dim=1) == y).float().mean())


def train_teacher(x, y, epochs=40):
    torch.manual_seed(1)
    model = Teacher()
    opt = torch.optim.Adam(model.parameters(), lr=0.01)
    for epoch in range(epochs):
        model.train()
        loss = F.cross_entropy(model(x), y)
        opt.zero_grad()
        loss.backward()
        opt.step()
    return model


def train_hard(x, y, epochs=15):
    torch.manual_seed(2)
    model = Student()
    opt = torch.optim.Adam(model.parameters(), lr=0.01)
    for epoch in range(epochs):
        model.train()
        loss = F.cross_entropy(model(x), y)
        opt.zero_grad()
        loss.backward()
        opt.step()
    return model


def train_soft(x, teacher, epochs=15, temp=3.0):
    torch.manual_seed(2)
    model = Student()
    opt = torch.optim.Adam(model.parameters(), lr=0.01)
    teacher.eval()
    with torch.no_grad():
        soft = F.softmax(teacher(x) / temp, dim=1)
    for epoch in range(epochs):
        model.train()
        logp = F.log_softmax(model(x) / temp, dim=1)
        loss = F.kl_div(logp, soft, reduction="batchmean") * (temp * temp)
        opt.zero_grad()
        loss.backward()
        opt.step()
    return model


train_x, train_y = make_shapes(60, seed=42, noise=0.4)
val_x, val_y = make_shapes(120, seed=43, noise=0.4)

teacher = train_teacher(train_x, train_y)
hard = train_hard(train_x, train_y, epochs=10)
soft = train_soft(train_x, teacher, epochs=10)

print(f"Teacher parameters: {sum(p.numel() for p in teacher.parameters())}")
print(f"Student parameters: {sum(p.numel() for p in Student().parameters())}")
print(f"Teacher validation accuracy: {accuracy(teacher, val_x, val_y):.3f}")
print(f"Student on hard labels accuracy: {accuracy(hard, val_x, val_y):.3f}")
print(f"Student on soft targets accuracy: {accuracy(soft, val_x, val_y):.3f}")