# Source: content/notes/deep-learning/transfer-learning-and-finetuning.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

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

torch.manual_seed(22)
torch.set_num_threads(1)
x = torch.randn(256, 5)
teacher = nn.Linear(5, 3)
teacher.requires_grad_(False)
teacher.eval()
student = nn.Linear(5, 3)
temperature, mixture = 2.0, 0.7
with torch.no_grad():
    teacher_logits = teacher(x)
    soft_targets = (teacher_logits / temperature).softmax(-1)
    labels = teacher_logits.argmax(-1)
opt = torch.optim.Adam(student.parameters(), lr=0.04)

def objective():
    logits = student(x)
    soft = F.kl_div((logits / temperature).log_softmax(-1), soft_targets,
                    reduction="batchmean") * temperature**2
    return mixture * soft + (1 - mixture) * F.cross_entropy(logits, labels)

initial = objective().item()
for _ in range(100):
    opt.zero_grad(set_to_none=True)
    loss = objective()
    loss.backward()
    opt.step()
final = objective().item()
assert final < initial * 0.5
assert all(p.grad is None for p in teacher.parameters())
print({"initial_objective": initial, "final_objective": final,
       "teacher_agreement": (student(x).argmax(-1) == labels).float().mean().item()})