使用PyTorch進行知識蒸餾的程式碼範例

隨著機器學習模型的複雜性和能力而不斷增加。提高大型複雜模型在小資料集表現的一種有效技術是知識蒸餾，它包括訓練一個更小、更有效的模型來模仿一個更大的「教師」模型的行為。

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

以上是使用PyTorch進行知識蒸餾的程式碼範例的詳細內容。更多資訊請關注PHP中文網其他相關文章！