import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.nn import functional as F

# 教師模型（較大網(wǎng)絡(luò)）
class TeacherModel(nn.Module):
    def __init__(self):
        super(TeacherModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(7*7*64, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = x.view(-1, 7*7*64)  # Flatten the tensor
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 學(xué)生模型（較小網(wǎng)絡(luò)）
class StudentModel(nn.Module):
    def __init__(self):
        super(StudentModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(7*7*32, 64)
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = x.view(-1, 7*7*32)  # Flatten the tensor
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

def distillation_loss(y_student, y_teacher, T=2.0, alpha=0.7):
    # 計(jì)算軟標(biāo)簽的交叉熵?fù)p失
    soft_loss = nn.KLDivLoss(reduction='batchmean')(
        F.log_softmax(y_student / T, dim=1),
        F.softmax(y_teacher / T, dim=1)
    )

    # 計(jì)算真實(shí)標(biāo)簽的交叉熵?fù)p失
    hard_loss = F.cross_entropy(y_student, torch.argmax(y_teacher, dim=1))

    # 綜合蒸餾損失
    return alpha * soft_loss + (1 - alpha) * hard_loss

def train(model, device, train_loader, optimizer, epoch, teacher_model=None, T=2.0, alpha=0.7):
    model.train()
    running_loss = 0.0

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad()

        # 教師模型和學(xué)生模型的輸出
        output = model(data)
        with torch.no_grad():  # 教師模型在蒸餾時不更新參數(shù)
            teacher_output = teacher_model(data)

        # 計(jì)算蒸餾損失
        loss = distillation_loss(output, teacher_output, T, alpha)

        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Train Epoch: {epoch} \tLoss: {running_loss / len(train_loader):.6f}")
    
    
batch_size = 64
epochs = 10
lr = 0.001
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 數(shù)據(jù)加載
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = datasets.MNIST('.', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('.', train=False, download=True, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 初始化教師模型和學(xué)生模型
teacher_model = TeacherModel().to(device)
student_model = StudentModel().to(device)

# 教師模型訓(xùn)練（簡單訓(xùn)練）
optimizer_teacher = optim.Adam(teacher_model.parameters(), lr=lr)
teacher_model.train()
for epoch in range(1, epochs + 1):
    train_teacher(teacher_model, device, train_loader, optimizer_teacher, epoch)

# 學(xué)生模型訓(xùn)練（蒸餾）
optimizer_student = optim.Adam(student_model.parameters(), lr=lr)
student_model.train()
for epoch in range(1, epochs + 1):
     train(student_model, device, train_loader, optimizer_student, epoch, teacher_mod