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

    知识蒸馏属于模型的压缩一种方法，但其实这种方法又属于一种伪压缩，是将一个性能较好的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:训练迭代数

    1. def teacher_train(teacher_model, train_loader, test_loader, loss_func, epochs):
    2.  
    3.     teacher_model.train()
    4.     teacher_model.cuda()
    5.     # train
    6.     for i in range(epochs):
    7.         for data, label in train_loader:
    8.             data = data.to(device)
    9.             label = label.to(device)
    10.             output = teacher_model(data)
    11.             loss = loss_func(output, label)
    12.             optimizer_teacher.zero_grad()
    13.             loss.backward()
    14.             optimizer_teacher.step()
    15.         print("loss: ", loss)
    16.         # eval
    17.         correct = 0
    18.         teacher_model.eval()
    19.         teacher_model.cuda()
    20.         for test_data, test_label in test_loader:
    21.             test_data = test_data.to(device)
    22.             test_label = test_label.to(device)
    23.             with torch.no_grad():
    24.  
    25.                 output = teacher_model(test_data)
    26.                 # acc = torch.mean((torch.argmax(F.softmax(output, dim=-1), dim=-1) == test_label).type(torch.FloatTensor))
    27.                 # print("teacher acc: ", acc)
    28.                 _, pred = torch.max(output, dim=1)
    29.                 correct += float(torch.sum(pred == test_label))
    30.         print('test_acc:{}'.format(correct / len(test_dataset)))
    31.     return teacher_model

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

    1. teacher model train
    2. loss:  tensor(0.0891, device='cuda:0', grad_fn=)
    3. test_acc:0.9845
    4. loss:  tensor(0.0132, device='cuda:0', grad_fn=)
    5. test_acc:0.9865
    6. loss:  tensor(0.0019, device='cuda:0', grad_fn=)
    7. test_acc:0.9909
    8. loss:  tensor(0.0042, device='cuda:0', grad_fn=)
    9. test_acc:0.9909
    10. loss:  tensor(0.0034, device='cuda:0', grad_fn=)
    11. test_acc:0.9917
    12. teacher model trained finished!

    student未KD 训练 

    参数说明：

    student_model:选用的student网络

    train_loader:训练集

    test_loader:测试集

    loss_func:损失函数

    epochs:训练迭代数

    1. def student_train(student_model, train_loader, test_loader, loss_func, epochs):
    2.  
    3.     student_model.train()
    4.     student_model.cuda()
    5.     # train
    6.     for i in range(epochs):
    7.         for data, label in train_loader:
    8.             data = data.to(device)
    9.             label = label.to(device)
    10.             output = student_model(data)
    11.             loss = loss_func(output, label)
    12.             optimizer_student.zero_grad()
    13.             loss.backward()
    14.             optimizer_student.step()
    15.  
    16.         print("student loss: ", loss)
    17.         # eval
    18.         correct = 0
    19.         student_model.eval()
    20.         student_model.cuda()
    21.         for test_data, test_label in test_loader:
    22.             test_data = test_data.to(device)
    23.             test_label = test_label.to(device)
    24.             with torch.no_grad():
    25.                 output = student_model(test_data)
    26.                 # acc = torch.mean((torch.argmax(F.softmax(output, dim=-1), dim=-1) == test_label).type(torch.FloatTensor))
    27.                 # print("teacher acc: ", acc)
    28.                 _, pred = torch.max(output, dim=1)
    29.                 correct += float(torch.sum(pred == test_label))
    30.         print('student test_acc:{}'.format(correct / len(test_dataset)))

    没有KD train的效果如下： 

    1. student model ready train
    2. student loss:  tensor(0.1876, device='cuda:0', grad_fn=)
    3. student test_acc:0.9588
    4. student loss:  tensor(0.0219, device='cuda:0', grad_fn=)
    5. student test_acc:0.9737
    6. student loss:  tensor(0.0588, device='cuda:0', grad_fn=)
    7. student test_acc:0.9812
    8. student loss:  tensor(0.0024, device='cuda:0', grad_fn=)
    9. student test_acc:0.9853
    10. student loss:  tensor(0.0022, device='cuda:0', grad_fn=)
    11. student test_acc:0.9814
    12.  
    13.  student model trained finished!

    KD train代码

    参数说明：

    teacher_model:为已经训练好的teacher

    student_model:待KD的student网络

    train_loader:训练集

    test_loader:测试集

    1. def KD_train(teacher_model, student_model, train_loader, test_loader,loss_func, epochs):
    2.     teacher_model.eval()
    3.     student_model.train()
    4.     student_model.cuda()
    5.     HL = nn.CrossEntropyLoss()
    6.     for i in range(epochs):
    7.         for data, labels in train_loader:
    8.             data = data.to(device)
    9.             labels = labels.to(device)
    10.             teacher_output = teacher_model(data)
    11.             student_output = student_model(data)
    12.             soft_loss = KD_loss(teacher_output, student_output)
    13.             hard_loss = HL(student_output, labels)
    14.             loss = hard_loss + alpha*soft_loss
    15.             optimizer_student.zero_grad()
    16.             loss.backward()
    17.             optimizer_student.step()
    18.         print("KD loss: ", loss)
    19.         student_model.eval()
    20.         ACC = 0
    21.         for data, labels in test_loader:
    22.             with torch.no_grad():
    23.                 data = data.to(device)
    24.                 labels = labels.to(device)
    25.                 output = student_model(data)
    26.                 _, pred = torch.max(output, dim=1)
    27.                 ACC += float(torch.sum(pred == labels))
    28.         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，你也可以改为交叉熵。

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

    KD tran后student结果： 

    1. KD loss:  tensor(0.2580, device='cuda:0', grad_fn=)
    2. KD test_acc:0.9753
    3. KD loss:  tensor(0.1686, device='cuda:0', grad_fn=)
    4. KD test_acc:0.9748
    5. KD loss:  tensor(0.0827, device='cuda:0', grad_fn=)
    6. KD test_acc:0.9849
    7. KD loss:  tensor(0.0098, device='cuda:0', grad_fn=)
    8. KD test_acc:0.9865
    9. KD loss:  tensor(0.0114, device='cuda:0', grad_fn=)
    10. KD test_acc:0.988

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

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

    完整代码

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

    1. import torch
    2.  
    3. from torch.optim import Adam, SGD
    4. import torch.nn.functional as F
    5. import torch.nn as nn
    6. from torchvision.models import resnet50, resnet34, resnet18, MobileNetV2
    7. import torchvision
    8. import torchvision.transforms as transforms
    9.  
    10.  
    11. Temp = 2.  # 温度常数
    12. alpha = 0.5
    13. def KD_loss(p, q):  # p指的老师老师的预测(经过softmax)，q是学生的预测
    14.     pt = F.softmax(p / Temp, dim=1)
    15.     ps = F.log_softmax(q / Temp, dim=1)
    16.     return nn.KLDivLoss(reduction='mean')(ps, pt) * (Temp**2)
    17. def teacher_train(teacher_model, train_loader, test_loader, loss_func, epochs):
    18.  
    19.     teacher_model.train()
    20.     teacher_model.cuda()
    21.     # train
    22.     for i in range(epochs):
    23.         for data, label in train_loader:
    24.             data = data.to(device)
    25.             label = label.to(device)
    26.             output = teacher_model(data)
    27.             loss = loss_func(output, label)
    28.             optimizer_teacher.zero_grad()
    29.             loss.backward()
    30.             optimizer_teacher.step()
    31.         print("loss: ", loss)
    32.         # eval
    33.         correct = 0
    34.         teacher_model.eval()
    35.         teacher_model.cuda()
    36.         for test_data, test_label in test_loader:
    37.             test_data = test_data.to(device)
    38.             test_label = test_label.to(device)
    39.             with torch.no_grad():
    40.  
    41.                 output = teacher_model(test_data)
    42.                 # acc = torch.mean((torch.argmax(F.softmax(output, dim=-1), dim=-1) == test_label).type(torch.FloatTensor))
    43.                 # print("teacher acc: ", acc)
    44.                 _, pred = torch.max(output, dim=1)
    45.                 correct += float(torch.sum(pred == test_label))
    46.         print('test_acc:{}'.format(correct / len(test_dataset)))
    47.     return teacher_model
    48.  
    49. def student_train(student_model, train_loader, test_loader, loss_func, epochs):
    50.  
    51.     student_model.train()
    52.     student_model.cuda()
    53.     # train
    54.     for i in range(epochs):
    55.         for data, label in train_loader:
    56.             data = data.to(device)
    57.             label = label.to(device)
    58.             output = student_model(data)
    59.             loss = loss_func(output, label)
    60.             optimizer_student.zero_grad()
    61.             loss.backward()
    62.             optimizer_student.step()
    63.  
    64.         print("student loss: ", loss)
    65.         # eval
    66.         correct = 0
    67.         student_model.eval()
    68.         student_model.cuda()
    69.         for test_data, test_label in test_loader:
    70.             test_data = test_data.to(device)
    71.             test_label = test_label.to(device)
    72.             with torch.no_grad():
    73.                 output = student_model(test_data)
    74.                 # acc = torch.mean((torch.argmax(F.softmax(output, dim=-1), dim=-1) == test_label).type(torch.FloatTensor))
    75.                 # print("teacher acc: ", acc)
    76.                 _, pred = torch.max(output, dim=1)
    77.                 correct += float(torch.sum(pred == test_label))
    78.         print('student test_acc:{}'.format(correct / len(test_dataset)))
    79.  
    80. def KD_train(teacher_model, student_model, train_loader, test_loader,loss_func, epochs):
    81.     teacher_model.eval()
    82.     student_model.train()
    83.     student_model.cuda()
    84.     HL = nn.CrossEntropyLoss()
    85.     for i in range(epochs):
    86.         for data, labels in train_loader:
    87.             data = data.to(device)
    88.             labels = labels.to(device)
    89.             teacher_output = teacher_model(data)
    90.             student_output = student_model(data)
    91.             soft_loss = KD_loss(teacher_output, student_output)
    92.             hard_loss = HL(student_output, labels)
    93.             loss = hard_loss + alpha*soft_loss
    94.             optimizer_student.zero_grad()
    95.             loss.backward()
    96.             optimizer_student.step()
    97.         print("KD loss: ", loss)
    98.         student_model.eval()
    99.         ACC = 0
    100.         for data, labels in test_loader:
    101.             with torch.no_grad():
    102.                 data = data.to(device)
    103.                 labels = labels.to(device)
    104.                 output = student_model(data)
    105.                 _, pred = torch.max(output, dim=1)
    106.                 ACC += float(torch.sum(pred == labels))
    107.         print('KD test_acc:{}'.format(ACC / len(test_dataset)))
    108.  
    109.  
    110. def do_train(teacher_model, student_model, train_loader, test_loader, loss_func, epochs):
    111.     #教师训练
    112.     teacher_model.train()
    113.     teacher_model.to(device)
    114.     print("teacher model train")
    115.     Teacher = teacher_train(teacher_model, train_loader, test_loader, loss_func, epochs)
    116.     print("teacher model trained finished!")
    117.  
    118.     # print("\n student model ready train")
    119.     # student_train(student_model, train_loader, test_loader, loss_func, epochs)
    120.     # print("\n student model trained finished!")
    121.  
    122.     print("\n KD model ready train")
    123.     KD_train(Teacher, student_model, train_loader, test_loader, loss_func, epochs)
    124.  
    125.  
    126. if __name__=="__main__":
    127.     # 准备数据集
    128.     batch_size = 64
    129.     train_dataset = torchvision.datasets.MNIST('./data/', train=True, download=True,
    130.                                                transform=transforms.Compose([
    131.                                                    transforms.Resize(28),
    132.                                                    transforms.ToTensor(),
    133.                                                    transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
    134.                                                    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
    135.                                                    transforms.Grayscale(num_output_channels=3)
    136.                                                ])
    137.                                                )
    138.     test_dataset = torchvision.datasets.MNIST('./data/', train=False, download=True,
    139.                                               transform=transforms.Compose([
    140.                                                   transforms.Resize(28),  # resnet默认图片输入大小224*224
    141.                                                   transforms.ToTensor(),
    142.                                                   transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
    143.                                                   transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
    144.                                                   transforms.Grayscale(num_output_channels=3)
    145.                                               ])
    146.                                               )
    147.     train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    148.     test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
    149.  
    150.     sample, label = next(iter(train_loader))
    151.     print(sample.shape)
    152.     print("当前类： ", label)
    153.  
    154.  
    155.  
    156.     num_classes = 10
    157.     lr = 0.01
    158.     epochs = 5
    159.     device = torch.device('cuda:0')
    160.     teacher_model = resnet18(num_classes=num_classes)
    161.     student_model = resnet50(num_classes=num_classes)
    162.     optimizer_teacher = SGD(teacher_model.parameters(), lr=lr, momentum=0.9)
    163.     optimizer_student = SGD(student_model.parameters(), lr=lr, momentum=0.9)
    164.     loss_function = nn.CrossEntropyLoss()
    165.     do_train(teacher_model, student_model, train_loader, test_loader, loss_function, epochs)
    166.  
    167.  
    168.

  • 相关阅读:
    阿里云-系统盘-磁盘扩容
    元数据治理利器 - Apache Atlas
    java中log使用总结
    Mybiosource抗 DGKE 抗体丨Mybiosource DGKE兔多克隆抗体
    浏览器本地存储webStroage
    浅谈中国程序员的四个层次，你在第几层？
    push，pop指令
    独立成分分析ICA/FastICA
    TodoList案例
    内网资料传外网速度太慢怎么办，一招教你解决
  • 原文地址：https://blog.csdn.net/z240626191s/article/details/128154154