使用PyTorch進行知識蒸餾的程式碼範例-人工智慧-PHP中文網

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

王林

發布： 2023-04-11 22:31:13

轉載

957 人瀏覽過

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

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

在本文中，我們將探討知識蒸餾的概念，以及如何在PyTorch中實現它。我們將看到如何使用它將一個龐大、笨重的模型壓縮成一個更小、更有效率的模型，並且仍然保留原始模型的準確性和性能。

我們先定義知識蒸餾要解決的問題。

我們訓練了一個大型深度神經網路來執行複雜的任務，例如影像分類或機器翻譯。這個模型可能有數千層和數百萬個參數，這使得它很難部署在現實應用程式、邊緣設備等。而這個超大的模型還需要大量的運算資源來運行，這使得它在一些資源受限的平台上無法運作。

解決這個問題的一種方法是使用知識蒸餾將大模型壓縮成較小的模型。這個過程包括訓練一個較小的模型來模仿給定任務中大型模型的行為。

我們將使用來自Kaggle的胸部x光資料集進行肺炎分類來進行知識蒸餾的範例。我們使用的資料集被組織成3個資料夾(train, test, val)，並包含每個影像類別的子資料夾(Pneumonia/Normal)。共有5,863張x射線影像(JPEG)和2個類別(肺炎/正常)。

比較這兩個類別的圖片:

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

資料的載入和預處理與我們是否使用知識蒸餾或特定模型無關，程式碼片段可能如下所示:

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

登入後複製

損失函數是下面兩個東西的加權和: