使用PyTorch進行知識蒸餾的代碼示例

作者：Alessandro Lamberti 2022-12-19 15:16:46

隨著機器學習模型的復雜性和能力不斷增加。提高大型復雜模型在小數據集性能的一種有效技術是知識蒸餾，它包括訓練一個更小、更有效的模型來模仿一個更大的“教師”模型的行為。

在本文中，我們將探索知識蒸餾的概念，以及如何在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倍!也就是說，我們的小模型真的從大模型中學到了有用的東西。

責任編輯：華軒來源： DeepHub IMBA

機器學習模型