隨著機(jī)器學(xué)習(xí)模型的復(fù)雜性和能力不斷增加。提高大型復(fù)雜模型在小數(shù)據(jù)集性能的一種有效技術(shù)是知識(shí)蒸餾，它包括訓(xùn)練一個(gè)更小、更有效的模型來模仿一個(gè)更大的“教師”模型的行為。
在本文中，我們將探索知識(shí)蒸餾的概念，以及如何在PyTorch中實(shí)現(xiàn)它。我們將看到如何使用它將一個(gè)龐大、笨重的模型壓縮成一個(gè)更小、更高效的模型，并且仍然保留原始模型的準(zhǔn)確性和性能。
我們首先定義知識(shí)蒸餾要解決的問題。
我們訓(xùn)練了一個(gè)大型深度神經(jīng)網(wǎng)絡(luò)來執(zhí)行復(fù)雜的任務(wù)，比如圖像分類或機(jī)器翻譯。這個(gè)模型可能有數(shù)千層和數(shù)百萬個(gè)參數(shù)，這使得它很難部署在現(xiàn)實(shí)應(yīng)用程序、邊緣設(shè)備等中。并且這個(gè)超大的模型還需要大量的計(jì)算資源來運(yùn)行，這使得它在一些資源受限的平臺(tái)上無法工作。
解決這個(gè)問題的一種方法是使用知識(shí)蒸餾將大模型壓縮成較小的模型。這個(gè)過程包括訓(xùn)練一個(gè)較小的模型來模仿給定任務(wù)中大型模型的行為。
我們將使用來自Kaggle的胸部x光數(shù)據(jù)集進(jìn)行肺炎分類來進(jìn)行知識(shí)蒸餾的示例。我們使用的數(shù)據(jù)集被組織成3個(gè)文件夾(train, test, val)，并包含每個(gè)圖像類別的子文件夾(Pneumonia/Normal)。共有5,863張x射線圖像(JPEG)和2個(gè)類別(肺炎/正常)。
比較一下這兩個(gè)類的圖片:
數(shù)據(jù)的加載和預(yù)處理與我們是否使用知識(shí)蒸餾或特定模型無關(guān)，代碼片段可能如下所示:
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)
教師模型
在這個(gè)背景中教師模型我們使用Resnet-18并且在這個(gè)數(shù)據(jù)集上進(jìn)行了微調(diào)。
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
微調(diào)訓(xùn)練的代碼如下
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))
這是一個(gè)標(biāo)準(zhǔn)的微調(diào)訓(xùn)練步驟，訓(xùn)練后我們可以看到該模型在測(cè)試集上達(dá)到了91%的準(zhǔn)確性，這也就是我們沒有選擇更大模型的原因，因?yàn)樽鳛闇y(cè)試91的準(zhǔn)確率已經(jīng)足夠作為基類模型來使用了。
我們知道模型有1170萬個(gè)參數(shù)，因此不一定能夠適應(yīng)邊緣設(shè)備或其他特定場(chǎng)景。
學(xué)生模型
我們的學(xué)生是一個(gè)更淺的CNN，只有幾層和大約100k個(gè)參數(shù)。
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
看代碼就非常的簡(jiǎn)單，對(duì)吧。
如果我可以簡(jiǎn)單地訓(xùn)練這個(gè)更小的神經(jīng)網(wǎng)絡(luò)，我為什么還要費(fèi)心進(jìn)行知識(shí)蒸餾呢?我們最后會(huì)附上我們通過超參數(shù)調(diào)整等手段從頭訓(xùn)練這個(gè)網(wǎng)絡(luò)的結(jié)果最為對(duì)比。
但是現(xiàn)在我們繼續(xù)我們的知識(shí)蒸餾的步驟
知識(shí)蒸餾訓(xùn)練
訓(xùn)練的基本步驟是不變的，但是區(qū)別是如何計(jì)算最終的訓(xùn)練損失，我們將使用教師模型損失，學(xué)生模型的損失和蒸餾損失一起來計(jì)算最終的損失。
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
損失函數(shù)是下面兩個(gè)東西的加權(quán)和:
- 分類損失，稱為student_target_loss
- 蒸餾損失，學(xué)生對(duì)數(shù)和教師對(duì)數(shù)之間的交叉熵?fù)p失
簡(jiǎn)單的講，我們的教師模型需要教導(dǎo)學(xué)生如何“思考”的，這就是指的是它的不確定性;例如，如果教師模型的最終輸出概率是[0.53,0.47]，我們希望學(xué)生也得到同樣類似結(jié)果，這些預(yù)測(cè)之間的差異就是蒸餾損失。
為了控制損失，還有有兩個(gè)主要參數(shù):
- 蒸餾損失的權(quán)重：0意味著我們只考慮蒸餾損失，反之亦然。
- 溫度：衡量教師預(yù)測(cè)的不確定性。
在上面的要點(diǎn)中，alpha和temperature的值都是根據(jù)我們嘗試過一些組合得到的最佳結(jié)果的值。
結(jié)果對(duì)比
這是這個(gè)實(shí)驗(yàn)的表格摘要。
我們可以清楚地看到使用更小(99.14%)，更淺的CNN所獲得的巨大好處:與無蒸餾訓(xùn)練相比，準(zhǔn)確率提升了10點(diǎn)，并且比Resnet-18快11倍!也就是說，我們的小模型真的從大模型中學(xué)到了有用的東西。
