dataloader_tools.py

import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader

def load_data():
    # 载入MNIST训练集
    train_dataset = torchvision.datasets.MNIST(
        root = "../datasets/",
        train=True,
        transform=transforms.ToTensor(),
        download=True
    )

    # 载入MNIST测试集
    test_dataset = torchvision.datasets.MNIST(
        root = "../datasets/",
        train=False,
        transform=transforms.ToTensor(),
        download=True
    )

    # 生成训练集和测试集的dataloader
    train_dataloader = DataLoader(dataset=train_dataset,batch_size=12,shuffle=True)
    test_dataloader = DataLoader(dataset=test_dataset,batch_size=12,shuffle=False)
    return train_dataloader,test_dataloader

models.py

import torch
from torch import nn
# 教师模型
class TeacherModel(nn.Module):
    def __init__(self,in_channels=1,num_classes=10):
        super(TeacherModel,self).__init__()
        self.relu = nn.ReLU()
        self.fc1 = nn.Linear(784,1200)
        self.fc2 = nn.Linear(1200,1200)
        self.fc3 = nn.Linear(1200,num_classes)
        self.dropout = nn.Dropout(p=0.5) #p=0.5是丢弃该层一半的神经元.
    def forward(self,x):
        x = x.view(-1,784)
        x = self.fc1(x)
        x = self.dropout(x)
        x = self.relu(x)

        x = self.fc2(x)
        x = self.dropout(x)
        x = self.relu(x)

        x = self.fc3(x)
        return x

class StudentModel(nn.Module):
    def __init__(self,in_channels=1,num_classes=10):
        super(StudentModel,self).__init__()
        self.relu = nn.ReLU()
        self.fc1 = nn.Linear(784,20)
        self.fc2 = nn.Linear(20,20)
        self.fc3 = nn.Linear(20,num_classes)
    def forward(self,x):
        x = x.view(-1,784)
        x = self.fc1(x)
        x = self.relu(x)

        x = self.fc2(x)
        x = self.relu(x)

        x = self.fc3(x)
        return x

train_tools.py

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

def train(epochs, model, model_name, lr,train_dataloader,test_dataloader,device):
    # ----------------------开始计时-----------------------------------
    start_time = time.time()

    # 设置参数开始训练
    best_acc, best_epoch = 0, 0
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    for epoch in range(epochs):
        model.train()
        # 训练集上训练模型权重
        for data, targets in tqdm.tqdm(train_dataloader):
            # 把数据加载到GPU上
            data = data.to(device)
            targets = targets.to(device)

            # 前向传播
            preds = model(data)
            loss = criterion(preds, targets)

            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # 测试集上评估模型性能
        model.eval()
        num_correct = 0
        num_samples = 0
        with torch.no_grad():
            for x, y in test_dataloader:
                x = x.to(device)
                y = y.to(device)
                preds = model(x)
                predictions = preds.max(1).indices  # 返回每一行的最大值和该最大值在该行的列索引
                num_correct += (predictions == y).sum()
                num_samples += predictions.size(0)
            acc = (num_correct / num_samples).item()
            if acc > best_acc:
                best_acc = acc
                best_epoch = epoch
                # 保存模型最优准确率的参数
                torch.save(model.state_dict(), f"../weights/{model_name}_best_acc_params.pth")
        model.train()
        print('Epoch:{}\t Accuracy:{:.4f}'.format(epoch + 1, acc),f'loss={loss}')
    print(f'最优准确率的epoch为{best_epoch},值为:{best_acc},最优参数已经保存到:weights/{model_name}_best_acc_params.pth')

    # -------------------------结束计时------------------------------------
    end_time = time.time()
    run_time = end_time - start_time
    # 将输出的秒数保留两位小数
    if int(run_time) < 60:
        print(f'训练用时为:{round(run_time, 2)}s')
    else:
        print(f'训练用时为:{round(run_time / 60, 2)}minutes')

def distill_train(epochs,teacher_model,student_model,model_name,train_dataloader,test_dataloader,alpha,lr,temp,device):
    # -------------------------------------开始计时--------------------------------
    start_time = time.time()

    # 定以损失函数
    hard_loss = nn.CrossEntropyLoss()
    soft_loss = nn.KLDivLoss(reduction="batchmean")
    # 定义优化器
    optimizer = torch.optim.Adam(student_model.parameters(), lr=lr)

    best_acc,best_epoch = 0,0
    for epoch in range(epochs):
        student_model.train()
        # 训练集上训练模型权重
        for data,targets in tqdm.tqdm(train_dataloader):
            # 把数据加载到GPU上
            data = data.to(device)
            targets = targets.to(device)

            # 教师模型预测
            with torch.no_grad():
                teacher_preds = teacher_model(data)
            # 学生模型预测
            student_preds = student_model(data)
            # 计算hard_loss
            student_hard_loss = hard_loss(student_preds,targets)

            # 计算蒸馏后的预测结果及soft_loss
            ditillation_loss = soft_loss(
                F.softmax(student_preds/temp,dim=1),
                F.softmax(teacher_preds/temp,dim=1)
            )

            # 将hard_loss和soft_loss加权求和
            loss = temp * temp * alpha * student_hard_loss + (1-alpha)*ditillation_loss

            # 反向传播,优化权重
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        #测试集上评估模型性能
        student_model.eval()
        num_correct = 0
        num_samples = 0
        with torch.no_grad():
            for x,y in test_dataloader:
                x = x.to(device)
                y = y.to(device)
                preds = student_model(x)
                predictions = preds.max(1).indices #返回每一行的最大值和该最大值在该行的列索引
                num_correct += (predictions ==y).sum()
                num_samples += predictions.size(0)
            acc = (num_correct/num_samples).item()
            if acc>best_acc:
                best_acc = acc
                best_epoch = epoch
                # 保存模型最优准确率的参数
                torch.save(student_model.state_dict(),f"../weights/{model_name}_best_acc_params.pth")
        student_model.train()
        print('Epoch:{}\t Accuracy:{:.4f}'.format(epoch+1,acc))
        print(f'student_hard_loss={student_hard_loss},ditillation_loss={ditillation_loss},loss={loss}')
    print(f'最优准确率的epoch为{best_epoch},值为:{best_acc},')

    # --------------------------------结束计时----------------------------------
    end_time = time.time()
    run_time = end_time - start_time
    # 将输出的秒数保留两位小数
    if int(run_time) < 60:
        print(f'训练用时为:{round(run_time, 2)}s')
    else:
        print(f'训练用时为:{round(run_time / 60, 2)}minutes')

训练教师网络

import torch
from torchinfo import summary #用来可视化的
import models
import dataloader_tools
import train_tools

# 设置随机数种子
torch.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 使用cuDNN加速卷积运算
torch.backends.cudnn.benchmark = True

# 载入MNIST训练集和测试集
train_dataloader,test_dataloader = dataloader_tools.load_data()

# 定义教师模型
model = models.TeacherModel()
model = model.to(device)
# 打印模型的参数
summary(model)

# 定义参数并开始训练
epochs = 10
lr = 1e-4
model_name = 'teacher'
train_tools.train(epochs,model,model_name,lr,train_dataloader,test_dataloader,device)

最优准确率的epoch为9,值为:0.9868999719619751

用非蒸馏的方法训练学生网络

import torch
from torchinfo import summary #用来可视化的
import dataloader_tools
import models
import train_tools

# 设置随机数种子
torch.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 使用cuDNN加速卷积运算
torch.backends.cudnn.benchmark = True

# 生成训练集和测试集的dataloader
train_dataloader,test_dataloader = dataloader_tools.load_data()

# 从头训练学生模型
model = models.StudentModel()
model = model.to(device)
# 查看模型参数
print(summary(model))

# 定义参数并开始训练
epochs = 10
lr = 1e-4
model_name = 'student'
train_tools.train(epochs, model, model_name, lr,train_dataloader,test_dataloader,device)

最优准确率的epoch为9,准确率为:0.9382999539375305,最优参数已经保存到:weights/student_best_acc_params.pth
训练用时为:1.74minutes

用知识蒸馏的方法训练student model

import torch
import train_tools
import models
import dataloader_tools

# 设置随机数种子
torch.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 使用cuDNN加速卷积运算
torch.backends.cudnn.benchmark = True

# 加载数据
train_dataloader,test_dataloader = dataloader_tools.load_data()


# 加载训练好的teacher model
teacher_model = models.TeacherModel()
teacher_model = teacher_model.to(device)
teacher_model.load_state_dict(torch.load('../weights/teacher_best_acc_params.pth'))
teacher_model.eval()

# 准备新的学生模型
student_model = models.StudentModel()
student_model = student_model.to(device)
student_model.train()

# 开始训练
lr = 0.0001
epochs = 20
alpha = 0.3 # hard_loss权重
temp = 7 # 蒸馏温度
model_name = 'distill_student_loss'
# 调用train_tools中的
train_tools.distill_train(epochs,teacher_model,student_model,model_name,train_dataloader,test_dataloader,alpha,lr,temp,device)

最优准确率的epoch为9,值为:0.9204999804496765,
训练用时为:2.14minutes

loss改为：

# temp的平方乘在student_hard_loss
loss = temp * temp * alpha * student_hard_loss + (1 - alpha) * ditillation_loss

最优准确率的epoch为9,值为:0.9336999654769897,
训练用时为:2.12minutes

loss改为：

# temp的平方乘ditillation_loss
loss = alpha * student_hard_loss + temp * temp * (1 - alpha) * ditillation_loss

最优准确率的epoch为9,值为:0.9176999926567078,
训练用时为:2.09minutes

上面的几种loss，蒸馏损失都出现了负数的情况。不太对劲。

其它开源的知识蒸馏算法如下：

open-mmlab开源的工具箱包含知识蒸馏算法

mmrazor

github.com/open-mmlab/mmrazor

NAS：神经架构搜索
剪枝：Pruning
KD: 知识蒸馏
Quantization: 量化

自定义知识蒸馏算法：

mmdeploy

可以把算法部署到一些厂商支持的中间格式，如ONNX，tensorRT等。

HobbitLong的RepDistiller

github.com/HobbitLong/RepDistiller

里面有12种最新的知识蒸馏算法。

蒸馏网络可以应用于同一种模型，将大的学习的知识蒸馏到小的上面。
如下将resnet100做教师网络，resnet32做学生网络。

将一种模型迁移到另一种模型上。如vgg13做教师网络，mobilNetv2做学生网络：