随着机器学习模型的复杂性和能力一直减少。进步大型简单模型在小数据集性能的一种无效技术是常识蒸馏，它包含训练一个更小、更无效的模型来模拟一个更大的“老师”模型的行为。
在本文中，咱们将摸索常识蒸馏的概念，以及如何在PyTorch中实现它。咱们将看到如何应用它将一个宏大、轻便的模型压缩成一个更小、更高效的模型，并且依然保留原始模型的准确性和性能。
咱们首先定义常识蒸馏要解决的问题。
咱们训练了一个大型深度神经网络来执行简单的工作，比方图像分类或机器翻译。这个模型可能有数千层和数百万个参数，这使得它很难部署在事实应用程序、边缘设施等中。并且这个超大的模型还须要大量的计算资源来运行，这使得它在一些资源受限的平台上无奈工作。
解决这个问题的一种办法是应用常识蒸馏将大模型压缩成较小的模型。这个过程包含训练一个较小的模型来模拟给定工作中大型模型的行为。
咱们将应用来自Kaggle的胸部x光数据集进行肺炎分类来进行常识蒸馏的示例。咱们应用的数据集被组织成3个文件夹(train, test, val)，并蕴含每个图像类别的子文件夹(Pneumonia/Normal)。共有5,863张x射线图像(JPEG)和2个类别(肺炎/失常)。
比拟一下这两个类的图片:
数据的加载和预处理与咱们是否应用常识蒸馏或特定模型无关，代码片段可能如下所示:
transforms_train = transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]) transforms_test = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]) train_data = ImageFolder(root=train_dir, transform=transforms_train) test_data = ImageFolder(root=test_dir, transform=transforms_test) train_loader = DataLoader(train_data, batch_size=32, shuffle=True) test_loader = DataLoader(test_data, batch_size=32, shuffle=True)老师模型
在这个背景中老师模型咱们应用Resnet-18并且在这个数据集上进行了微调。
import torch import torch.nn as nn import torchvision class TeacherNet(nn.Module): def __init__(self): super().__init__() self.model = torchvision.models.resnet18(pretrained=True) for params in self.model.parameters(): params.requires_grad_ = False n_filters = self.model.fc.in_features self.model.fc = nn.Linear(n_filters, 2) def forward(self, x): x = self.model(x) return x微调训练的代码如下
def train(model, train_loader, test_loader, optimizer, criterion, device): dataloaders = {'train': train_loader, 'val': test_loader} for epoch in range(30): print('Epoch {}/{}'.format(epoch, num_epochs - 1)) print('-' * 10) for phase in ['train', 'val']: if phase == 'train': model.train() else: model.eval() running_loss = 0.0 running_corrects = 0 for inputs, labels in tqdm.tqdm(dataloaders[phase]): inputs = inputs.to(device) labels = labels.to(device) optimizer.zero_grad() with torch.set_grad_enabled(phase == 'train'): outputs = model(inputs) loss = criterion(outputs, labels) _, preds = torch.max(outputs, 1) if phase == 'train': loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) epoch_loss = running_loss / len(dataloaders[phase].dataset) epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset) print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))这是一个规范的微调训练步骤，训练后咱们能够看到该模型在测试集上达到了91%的准确性，这也就是咱们没有抉择更大模型的起因，因为作为测试91的准确率曾经足够作为基类模型来应用了。
咱们晓得模型有1170万个参数，因而不肯定可能适应边缘设施或其余特定场景。
学生模型
咱们的学生是一个更浅的CNN，只有几层和大概100k个参数。
class StudentNet(nn.Module): def __init__(self): super().__init__() self.layer1 = nn.Sequential( nn.Conv2d(3, 4, kernel_size=3, padding=1), nn.BatchNorm2d(4), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2) ) self.fc = nn.Linear(4 * 112 * 112, 2) def forward(self, x): out = self.layer1(x) out = out.view(out.size(0), -1) out = self.fc(out) return out看代码就十分的简略，对吧。
如果我能够简略地训练这个更小的神经网络，我为什么还要费神进行常识蒸馏呢?咱们最初会附上咱们通过超参数调整等伎俩从头训练这个网络的后果最为比照。
然而当初咱们持续咱们的常识蒸馏的步骤
常识蒸馏训练
训练的根本步骤是不变的，然而区别是如何计算最终的训练损失，咱们将应用老师模型损失，学生模型的损失和蒸馏损失一起来计算最终的损失。
class DistillationLoss: def __init__(self): self.student_loss = nn.CrossEntropyLoss() self.distillation_loss = nn.KLDivLoss() self.temperature = 1 self.alpha = 0.25 def __call__(self, student_logits, student_target_loss, teacher_logits): distillation_loss = self.distillation_loss(F.log_softmax(student_logits / self.temperature, dim=1), F.softmax(teacher_logits / self.temperature, dim=1)) loss = (1 - self.alpha) * student_target_loss + self.alpha * distillation_loss return loss损失函数是上面两个货色的加权和:
- 分类损失，称为student_target_loss
- 蒸馏损失，学生对数和老师对数之间的穿插熵损失
简略的讲，咱们的老师模型须要教诲学生如何“思考”的，这就是指的是它的不确定性;例如，如果老师模型的最终输入概率是[0.53,0.47]，咱们心愿学生也失去同样相似后果，这些预测之间的差别就是蒸馏损失。
为了管制损失，还有有两个主要参数:
- 蒸馏损失的权重：0意味着咱们只思考蒸馏损失，反之亦然。
- 温度：掂量老师预测的不确定性。
在下面的要点中，alpha和temperature的值都是依据咱们尝试过一些组合失去的最佳后果的值。
后果比照
这是这个试验的表格摘要。
咱们能够分明地看到应用更小(99.14%)，更浅的CNN所取得的微小益处:与无蒸馏训练相比，准确率晋升了10点，并且比Resnet-18快11倍!也就是说，咱们的小模型真的从大模型中学到了有用的货色。
https://avoid.overfit.cn/post/482f5c111e7344179e2aba57865427ea
作者：Alessandro Lamberti