分类网络知识蒸馏【附代码】

时间:2022-12-06 19:17:48

知识蒸馏属于模型的压缩一种方法,但其实这种方法又属于一种伪压缩,是将一个性能较好的teacher网络“压缩”进一个性能较差的student网络中,或者是可类似于在teacher的指导下让student进行学习进而提高性能。

知识蒸馏是一种思想,并不像其他压缩方法有现成的库,因此对于实际需求与场景需要自己去实现。蒸馏也分为“离线”蒸馏与“在线”蒸馏。前者是建立T-S进行KD训练,而后者可以说是一种自学习,让student自己做自己的teacher。

同时蒸馏还分为逻辑蒸馏和特征蒸馏,前者是在两个网络最终输出部分建立loss关系,而后者是在网络中间的某些特征部分建立loss进行蒸馏。

本文是以手写数字为例,teacher选用的resnet18,student选用的resnet50【大家可能会想resnet50比resnet18强啊,为啥resnet50是student,这是因为我在实际测试的时候发现在手写数字这个数据上resnet18的准确率比resnet50高,猜测是因为在低分辨率下resnet50虽然loss在下降,但由于网络较深,特征丢失也明显,网络退化较明显】。当然这里你也可以尝试resnet做teacher,mobilnet做student【我这样训练了一下发现对mobilnet提升变化不大】


注:这里不做模型和蒸馏改进,仅仅是给大家展示一下效果,至于更细化的蒸馏有兴趣的可以自己去研究。【有关目标检测方面的KD 训练,我将会在明年以后推出】


目录

teacher train代码

student未KD 训练 

KD train代码

KD_loss代码:

完整代码


 

teacher train代码

参数说明:

teacher_model:选用的teacher网络

train_loader:训练集

test_loader:测试集

loss_func:损失函数

epochs:训练迭代数

def teacher_train(teacher_model, train_loader, test_loader, loss_func, epochs):

    teacher_model.train()
    teacher_model.cuda()
    # train
    for i in range(epochs):
        for data, label in train_loader:
            data = data.to(device)
            label = label.to(device)
            output = teacher_model(data)
            loss = loss_func(output, label)
            optimizer_teacher.zero_grad()
            loss.backward()
            optimizer_teacher.step()
        print("loss: ", loss)
        # eval
        correct = 0
        teacher_model.eval()
        teacher_model.cuda()
        for test_data, test_label in test_loader:
            test_data = test_data.to(device)
            test_label = test_label.to(device)
            with torch.no_grad():

                output = teacher_model(test_data)
                # acc = torch.mean((torch.argmax(F.softmax(output, dim=-1), dim=-1) == test_label).type(torch.FloatTensor))
                # print("teacher acc: ", acc)
                _, pred = torch.max(output, dim=1)
                correct += float(torch.sum(pred == test_label))
        print('test_acc:{}'.format(correct / len(test_dataset)))
    return teacher_model

训练结果(我只训练了5轮): 

teacher model train
loss:  tensor(0.0891, device='cuda:0', grad_fn=<NllLossBackward>)
test_acc:0.9845
loss:  tensor(0.0132, device='cuda:0', grad_fn=<NllLossBackward>)
test_acc:0.9865
loss:  tensor(0.0019, device='cuda:0', grad_fn=<NllLossBackward>)
test_acc:0.9909
loss:  tensor(0.0042, device='cuda:0', grad_fn=<NllLossBackward>)
test_acc:0.9909
loss:  tensor(0.0034, device='cuda:0', grad_fn=<NllLossBackward>)
test_acc:0.9917
teacher model trained finished!

student未KD 训练 

参数说明:

student_model:选用的student网络

train_loader:训练集

test_loader:测试集

loss_func:损失函数

epochs:训练迭代数

def student_train(student_model, train_loader, test_loader, loss_func, epochs):

    student_model.train()
    student_model.cuda()
    # train
    for i in range(epochs):
        for data, label in train_loader:
            data = data.to(device)
            label = label.to(device)
            output = student_model(data)
            loss = loss_func(output, label)
            optimizer_student.zero_grad()
            loss.backward()
            optimizer_student.step()

        print("student loss: ", loss)
        # eval
        correct = 0
        student_model.eval()
        student_model.cuda()
        for test_data, test_label in test_loader:
            test_data = test_data.to(device)
            test_label = test_label.to(device)
            with torch.no_grad():
                output = student_model(test_data)
                # acc = torch.mean((torch.argmax(F.softmax(output, dim=-1), dim=-1) == test_label).type(torch.FloatTensor))
                # print("teacher acc: ", acc)
                _, pred = torch.max(output, dim=1)
                correct += float(torch.sum(pred == test_label))
        print('student test_acc:{}'.format(correct / len(test_dataset)))

没有KD train的效果如下: 

student model ready train
student loss:  tensor(0.1876, device='cuda:0', grad_fn=<NllLossBackward>)
student test_acc:0.9588
student loss:  tensor(0.0219, device='cuda:0', grad_fn=<NllLossBackward>)
student test_acc:0.9737
student loss:  tensor(0.0588, device='cuda:0', grad_fn=<NllLossBackward>)
student test_acc:0.9812
student loss:  tensor(0.0024, device='cuda:0', grad_fn=<NllLossBackward>)
student test_acc:0.9853
student loss:  tensor(0.0022, device='cuda:0', grad_fn=<NllLossBackward>)
student test_acc:0.9814

 student model trained finished!

KD train代码

参数说明:

teacher_model:为已经训练好的teacher

student_model:待KD的student网络

train_loader:训练集

test_loader:测试集

def KD_train(teacher_model, student_model, train_loader, test_loader,loss_func, epochs):
    teacher_model.eval()
    student_model.train()
    student_model.cuda()
    HL = nn.CrossEntropyLoss()
    for i in range(epochs):
        for data, labels in train_loader:
            data = data.to(device)
            labels = labels.to(device)
            teacher_output = teacher_model(data)
            student_output = student_model(data)
            soft_loss = KD_loss(teacher_output, student_output)
            hard_loss = HL(student_output, labels)
            loss = hard_loss + alpha*soft_loss
            optimizer_student.zero_grad()
            loss.backward()
            optimizer_student.step()
        print("KD loss: ", loss)
        student_model.eval()
        ACC = 0
        for data, labels in test_loader:
            with torch.no_grad():
                data = data.to(device)
                labels = labels.to(device)
                output = student_model(data)
                _, pred = torch.max(output, dim=1)
                ACC += float(torch.sum(pred == labels))
        print('KD test_acc:{}'.format(ACC / len(test_dataset)))

代码中的teacher_output是teacher网络的输出,student_output是student的输出,两者之间设计的KD_loss代码如下:

KD_loss代码:

Temp为温度系数,默认为2【可以根据自己的数据集去尝试】

alpha是hard与soft的平衡系数【默认0.5,也是根据自己的实际情况调整】

损失函数采用的KL,你也可以改为交叉熵。

Temp = 2.  # 温度常数
alpha = 0.5
def KD_loss(p, q):  # p指的老师老师的预测(经过softmax),q是学生的预测
    pt = F.softmax(p / Temp, dim=1)
    ps = F.log_softmax(q / Temp, dim=1)
    return nn.KLDivLoss(reduction='mean')(ps, pt) * (Temp**2)

KD tran后student结果: 

KD loss:  tensor(0.2580, device='cuda:0', grad_fn=<AddBackward0>)
KD test_acc:0.9753
KD loss:  tensor(0.1686, device='cuda:0', grad_fn=<AddBackward0>)
KD test_acc:0.9748
KD loss:  tensor(0.0827, device='cuda:0', grad_fn=<AddBackward0>)
KD test_acc:0.9849
KD loss:  tensor(0.0098, device='cuda:0', grad_fn=<AddBackward0>)
KD test_acc:0.9865
KD loss:  tensor(0.0114, device='cuda:0', grad_fn=<AddBackward0>)
KD test_acc:0.988

 


 可以看出经过KD训练后student略有提升【主要手写数字这个太容易训练,稍微一训练就可以有较高的准确率】,如果换成别的数据集【比如猫狗数据集可能会明显点,可以自己试试】。

如果要换teacher和student网络,只需要在代码中将teacher_model和student_model网络进行替换即可。

完整代码

目标检测方面的KD比较麻烦,这个以后再讲。

import torch

from torch.optim import Adam, SGD
import torch.nn.functional as F
import torch.nn as nn
from torchvision.models import resnet50, resnet34, resnet18, MobileNetV2
import torchvision
import torchvision.transforms as transforms


Temp = 2.  # 温度常数
alpha = 0.5
def KD_loss(p, q):  # p指的老师老师的预测(经过softmax),q是学生的预测
    pt = F.softmax(p / Temp, dim=1)
    ps = F.log_softmax(q / Temp, dim=1)
    return nn.KLDivLoss(reduction='mean')(ps, pt) * (Temp**2)
def teacher_train(teacher_model, train_loader, test_loader, loss_func, epochs):

    teacher_model.train()
    teacher_model.cuda()
    # train
    for i in range(epochs):
        for data, label in train_loader:
            data = data.to(device)
            label = label.to(device)
            output = teacher_model(data)
            loss = loss_func(output, label)
            optimizer_teacher.zero_grad()
            loss.backward()
            optimizer_teacher.step()
        print("loss: ", loss)
        # eval
        correct = 0
        teacher_model.eval()
        teacher_model.cuda()
        for test_data, test_label in test_loader:
            test_data = test_data.to(device)
            test_label = test_label.to(device)
            with torch.no_grad():

                output = teacher_model(test_data)
                # acc = torch.mean((torch.argmax(F.softmax(output, dim=-1), dim=-1) == test_label).type(torch.FloatTensor))
                # print("teacher acc: ", acc)
                _, pred = torch.max(output, dim=1)
                correct += float(torch.sum(pred == test_label))
        print('test_acc:{}'.format(correct / len(test_dataset)))
    return teacher_model

def student_train(student_model, train_loader, test_loader, loss_func, epochs):

    student_model.train()
    student_model.cuda()
    # train
    for i in range(epochs):
        for data, label in train_loader:
            data = data.to(device)
            label = label.to(device)
            output = student_model(data)
            loss = loss_func(output, label)
            optimizer_student.zero_grad()
            loss.backward()
            optimizer_student.step()

        print("student loss: ", loss)
        # eval
        correct = 0
        student_model.eval()
        student_model.cuda()
        for test_data, test_label in test_loader:
            test_data = test_data.to(device)
            test_label = test_label.to(device)
            with torch.no_grad():
                output = student_model(test_data)
                # acc = torch.mean((torch.argmax(F.softmax(output, dim=-1), dim=-1) == test_label).type(torch.FloatTensor))
                # print("teacher acc: ", acc)
                _, pred = torch.max(output, dim=1)
                correct += float(torch.sum(pred == test_label))
        print('student test_acc:{}'.format(correct / len(test_dataset)))

def KD_train(teacher_model, student_model, train_loader, test_loader,loss_func, epochs):
    teacher_model.eval()
    student_model.train()
    student_model.cuda()
    HL = nn.CrossEntropyLoss()
    for i in range(epochs):
        for data, labels in train_loader:
            data = data.to(device)
            labels = labels.to(device)
            teacher_output = teacher_model(data)
            student_output = student_model(data)
            soft_loss = KD_loss(teacher_output, student_output)
            hard_loss = HL(student_output, labels)
            loss = hard_loss + alpha*soft_loss
            optimizer_student.zero_grad()
            loss.backward()
            optimizer_student.step()
        print("KD loss: ", loss)
        student_model.eval()
        ACC = 0
        for data, labels in test_loader:
            with torch.no_grad():
                data = data.to(device)
                labels = labels.to(device)
                output = student_model(data)
                _, pred = torch.max(output, dim=1)
                ACC += float(torch.sum(pred == labels))
        print('KD test_acc:{}'.format(ACC / len(test_dataset)))


def do_train(teacher_model, student_model, train_loader, test_loader, loss_func, epochs):
    #教师训练
    teacher_model.train()
    teacher_model.to(device)
    print("teacher model train")
    Teacher = teacher_train(teacher_model, train_loader, test_loader, loss_func, epochs)
    print("teacher model trained finished!")

    # print("\n student model ready train")
    # student_train(student_model, train_loader, test_loader, loss_func, epochs)
    # print("\n student model trained finished!")

    print("\n KD model ready train")
    KD_train(Teacher, student_model, train_loader, test_loader, loss_func, epochs)


if __name__=="__main__":
    # 准备数据集
    batch_size = 64
    train_dataset = torchvision.datasets.MNIST('./data/', train=True, download=True,
                                               transform=transforms.Compose([
                                                   transforms.Resize(28),
                                                   transforms.ToTensor(),
                                                   transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
                                                   transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
                                                   transforms.Grayscale(num_output_channels=3)
                                               ])
                                               )
    test_dataset = torchvision.datasets.MNIST('./data/', train=False, download=True,
                                              transform=transforms.Compose([
                                                  transforms.Resize(28),  # resnet默认图片输入大小224*224
                                                  transforms.ToTensor(),
                                                  transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
                                                  transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
                                                  transforms.Grayscale(num_output_channels=3)
                                              ])
                                              )
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    sample, label = next(iter(train_loader))
    print(sample.shape)
    print("当前类: ", label)



    num_classes = 10
    lr = 0.01
    epochs = 5
    device = torch.device('cuda:0')
    teacher_model = resnet18(num_classes=num_classes)
    student_model = resnet50(num_classes=num_classes)
    optimizer_teacher = SGD(teacher_model.parameters(), lr=lr, momentum=0.9)
    optimizer_student = SGD(student_model.parameters(), lr=lr, momentum=0.9)
    loss_function = nn.CrossEntropyLoss()
    do_train(teacher_model, student_model, train_loader, test_loader, loss_function, epochs)