KD-Lib Documentation

Route Constrained Optimization (RCO) is a knowledge distillation algorithm based on curriculum learning. Instead of supervising a student model with a fully converged teacher, RCO supervises the student using anchor points selected from the trajectory (route) the teacher model traveled in parameter space. This approach is designed to reduce the lower bound of congruence loss during distillation, hint, and mimicking learning.

To implement RCO, you use the RCO class from KD_Lib.KD. You must provide a teacher model, a student model, their respective optimizers, and the data loaders. You can control how frequently the student mimics the teacher's trajectory using the epoch_interval parameter.

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from KD_Lib.KD import RCO

# 1. Prepare Data
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST("mnist_data", train=True, download=True, transform=transforms.Compose([
        transforms.ToTensor(), 
        transforms.Normalize((0.1307,), (0.3081,))
    ])),
    batch_size=32, shuffle=True
)

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST("mnist_data", train=False, transform=transforms.Compose([
        transforms.ToTensor(), 
        transforms.Normalize((0.1307,), (0.3081,))
    ])),
    batch_size=32, shuffle=True
)

# 2. Setup Device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 3. Define Models
teacher_model = <your_model_instance>
student_model = <your_model_instance>

# 4. Define Optimizers
teacher_optimizer = optim.SGD(teacher_model.parameters(), lr=0.01)
student_optimizer = optim.SGD(student_model.parameters(), lr=0.01)

# 5. Initialize RCO Distiller
distiller = RCO(
    teacher_model, 
    student_model, 
    train_loader, 
    test_loader, 
    teacher_optimizer, 
    student_optimizer, 
    epoch_interval=5, 
    device=device
)

# 6. Execute Training Workflow
distiller.train_teacher(epochs=20)      # Train the teacher model
distiller.train_students(epochs=20)    # Train the student model
distiller.evaluate(teacher=True)      # Evaluate the teacher model
distiller.evaluate()                   # Evaluate the student model

Self-training is a process where a student model is first trained normally to obtain a pre-trained model, which is then used as a teacher to train itself by transferring soft targets.

To implement this using KD_Lib, you need to:

1. Prepare your train_loader and test_loader using PyTorch.
2. Define your student_model (the architecture you wish to train).
3. Define a student_optimizer (e.g., torch.optim.SGD).
4. Initialize the SelfTraining class from KD_Lib.KD by passing the student model, loaders, optimizer, and the target device.
5. Call .train_student(epochs=N) to perform the self-training loop.
6. Call .evaluate() to assess the model's performance.

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from KD_Lib.KD import SelfTraining

# 1. Prepare datasets and dataloaders
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "mnist_data",
        train=True,
        download=True,
        transform=transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        ),
    ),
    batch_size=32,
    shuffle=True,
)

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "mnist_data",
        train=False,
        transform=transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        ),
    ),
    batch_size=32,
    shuffle=True,
)

# 2. Set device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 3. Define student model and optimizer
student_model = <your model>
student_optimizer = optim.SGD(student_model.parameters(), lr=0.01)

# 4. Initialize KD_Lib SelfTraining and run
distiller = SelfTraining(student_model, train_loader, test_loader, student_optimizer, 
                         device=device)  

distiller.train_student(epochs=5)  # Train the student model
distiller.evaluate()            # Evaluate the student model

To implement Class-wise Self-knowledge Distillation (CSKD) for regularizing class-wise predictions, you can use the CSKD class from KD_Lib.KD.vision. This method penalizes the predictive distribution between different samples of the same label to mitigate overconfident predictions and reduce intra-class variations.

Key Requirements:

1. Pairwise Sampling: The dataloader must use pairwise sampling. Use the load_dataset utility from KD_Lib.KD.vision.CSKD.sampler with the sampling_mode='pair' argument to ensure compatibility.
2. Distiller Setup: Initialize the CSKD object by providing the teacher model (set to None for self-distillation), the student model, the training loader, the test loader, a teacher optimizer (set to None for self-distillation), and the student optimizer.
3. Training: Use train_student() to execute the training loop. You can enable loss plotting and model saving via parameters.

import torch
import torch.optim as optim
from torchvision import datasets, transforms
from KD_Lib.KD.vision import CSKD
from KD_Lib.KD.vision.CSKD.sampler import load_dataset

# 1. Prepare datasets with pairwise sampling
train_loader, val_loader = load_dataset('cifar100', '~/data/', 'pair', batch_size=128)

# 2. Define your student model and optimizer
student_model = <your model>
student_optimizer = optim.SGD(student_model.parameters(), 0.01)

# 3. Initialize the CSKD distiller
# For self-distillation, teacher=None and teacher_optimizer=None
distiller = CSKD(None, student_model, train_loader, val_loader, 
                  None, student_optimizer)  

# 4. Train, Evaluate, and Inspect
distiller.train_student(epochs=5, plot_losses=True, save_model=True)
distiller.evaluate(teacher=False)
print(distiller.get_parameters())

The ProbShift algorithm is used to handle incorrect soft targets by swapping the value of the ground truth (theoretical maximum) and the value of the predicted class (predicted maximum). This ensures that maximum confidence is reached at the ground truth label.

To implement ProbShift, you need to initialize the ProbShift class from KD_Lib.KD with the following components:

• teacher_model: The teacher network.
• student_model: The student network.
• train_loader: The training DataLoader.
• test_loader: The testing DataLoader.
• teacher_optimizer: The optimizer for the teacher model.
• student_optimizer: The optimizer for the student model.
• device: The torch device (e.g., cuda or cpu).

Once initialized, you can use the following workflow:

1. distiller.train_teacher(epochs=N): Trains the teacher model.
2. distiller.train_students(epochs=N): Trains the student model using the probability shift distillation.
3. distiller.evaluate(teacher=True): Evaluates the teacher model.
4. distiller.evaluate(): Evaluates the student model.

import torch
from KD_Lib.KD import ProbShift

# ... setup models, loaders, and optimizers ...

distiller = ProbShift(
    teacher_model, 
    student_model, 
    train_loader, 
    test_loader, 
    teacher_optimizer, 
    student_optimizer, 
    device=device
)

distiller.train_teacher(epochs=5)

distiller.train_students(epochs=5)

distiller.evaluate(teacher=True)

distiller.evaluate()

Deep Mutual Learning (DML) is an online algorithm where an ensemble of students (a cohort) learns collaboratively. Instead of a one-way transfer from a pre-trained teacher, DML uses a pool of untrained students that learn simultaneously. Each student is trained using two losses: a conventional supervised learning loss and a mimicry loss that aligns the student's class posterior with the class probabilities of other students in the cohort.

To implement DML using KD_Lib:

1. Define your student models (they can have different architectures).
2. Create a student_cohort (a tuple of the models).
3. Create a list of optimizers for these models, ensuring the order of optimizers matches the order of models in the cohort.
4. Initialize the DML class with the cohort, dataloaders, optimizers, and the target device.
5. Call .train_students() to begin the collaborative training process.

import torch
from torchvision import datasets, transforms
from KD_Lib.KD import DML

# 1. Setup Data
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST("mnist_data", train=True, download=True, transform=transforms.Compose([
        transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))
    ])),
    batch_size=32, shuffle=True
)
test_loader = torch.utils.data.DataLoader(
    datasets.MNIST("mnist_data", train=False, transform=transforms.Compose([
        transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))
    ])),
    batch_size=32, shuffle=True
)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 2. Define Student Cohort (models can have different architectures)
student_model_1 = <your_model_1>
student_model_2 = <your_model_2>
student_cohort = (student_model_1, student_model_2)

# 3. Define Optimizers (order MUST match student_cohort)
optimizer_1 = torch.optim.SGD(student_model_1.parameters(), 0.01)
optimizer_2 = torch.optim.SGD(student_model_2.parameters(), 0.01)
optimizers = [optimizer_1, optimizer_2]

# 4. Train and Evaluate
distiller = DML(student_cohort, train_loader, test_loader, optimizers, device=device)
distiller.train_students(epochs=5, plot_losses=True, save_model=True)
distiller.evaluate()

The Virtual Teacher algorithm allows you to train a student model using a teacher model that is designed with 100% accuracy via label smoothing regularization. The teacher outputs a distribution where the correct class is assigned a specific probability (e.g., 0.9), and the remaining probability is distributed across other classes.

To implement this, follow these steps:

1. Prepare Data: Set up your torch.utils.data.DataLoader for both training and testing.
2. Define Models: Initialize your student_model and your teacher model (or use the VirtualTeacher abstraction which handles the distribution logic).
3. Configure Optimizer: Define an optimizer (e.g., torch.optim.SGD) for the student model.
4. Initialize VirtualTeacher: Use the VirtualTeacher class from KD_Lib.KD. You must provide the student_model, train_loader, test_loader, student_optimizer, the correct_prob (the probability assigned to the correct label), and the target device.
5. Execute Training: Call .train_student(epochs=N) to begin training and .evaluate() to check performance.

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from KD_Lib.KD import VirtualTeacher

# 1. Define datasets and dataloaders
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "mnist_data",
        train=True,
        download=True,
        transform=transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        ),
    ),
    batch_size=32,
    shuffle=True,
)

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "mnist_data",
        train=False,
        transform=transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        ),
    ),
    batch_size=32,
    shuffle=True,
)

# 2. Set device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 3. Define student model and optimizer
student_model = <your model>
student_optimizer = optim.SGD(student_model.parameters(), lr=0.01)

# 4. Train using KD_Lib
distiller = VirtualTeacher(
    student_model, 
    train_loader, 
    test_loader, 
    student_optimizer, 
    correct_prob=0.9, 
    device=device
)

distiller.train_student(epochs=5)  # Train the student model
distiller.evaluate()              # Evaluate the student model

The LabelSmoothReg class in KD_Lib.KD implements Label Smoothing Regularization (LSR). This technique modifies the ground truth label distribution to prevent the model from becoming overconfident. In this implementation, incorrect teacher predictions are replaced with labels where the correct classes are assigned a specific probability (controlled by correct_prob).

To use it, you must provide a teacher model, a student model, data loaders for training and testing, and optimizers for both models.

import torch
from KD_Lib.KD import LabelSmoothReg

# ... setup models, loaders, and optimizers ...

distiller = LabelSmoothReg(
    teacher_model,
    student_model,
    train_loader,
    test_loader,
    teacher_optimizer,
    student_optimizer,
    correct_prob=0.9,
    device=device
)

# Training and evaluation workflow
distiller.train_teacher(epochs=5)      # Train the teacher model
distiller.train_students(epochs=5)     # Train the student model
distiller.evaluate(teacher=True)       # Evaluate the teacher model
distiller.evaluate()                 # Evaluate the student model