PyTorch Distillation Results, Honestly

Agreement counts matching picks, right or wrong. Four of these five picks match.

teacher_pred = torch.tensor([0, 1, 2, 2, 1])
student_pred = torch.tensor([0, 1, 2, 1, 1])
print("Agreements:", int((teacher_pred == student_pred).sum()))
print("Agreement rate:", round((teacher_pred == student_pred).float().mean().item(), 2))

The closing program slices the validation rows by true class. Each slice gets its own accuracy pair and agreement value.

for c in range(3):
    mask = val_y == c
    t_acc = float((t_pred[mask] == val_y[mask]).float().mean())
    agree = float((t_pred[mask] == s_pred[mask]).float().mean())

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)
soft = train_soft(train_x, teacher, epochs=10)

names = ["bar", "ring", "square"]
teacher.eval()
soft.eval()
with torch.no_grad():
    t_pred = teacher(val_x).argmax(dim=1)
    s_pred = soft(val_x).argmax(dim=1)

print(f"Teacher validation accuracy: {accuracy(teacher, val_x, val_y):.3f}")
print(f"Student validation accuracy: {accuracy(soft, val_x, val_y):.3f}")
print("Class | Teacher acc | Student acc | Agreement")
for c in range(3):
    mask = val_y == c
    t_acc = float((t_pred[mask] == val_y[mask]).float().mean())
    s_acc = float((s_pred[mask] == val_y[mask]).float().mean())
    agree = float((t_pred[mask] == s_pred[mask]).float().mean())
    print(f"{names[c]:6} | {t_acc:.3f}       | {s_acc:.3f}       | {agree:.3f}")
print(f"Overall agreement: {float((t_pred == s_pred).float().mean()):.3f}")