終于把神經(jīng)網(wǎng)絡中的知識蒸餾搞懂了！！！
作者：程序員小寒
在深度學習中，尤其是在計算機視覺和自然語言處理等任務中，深度神經(jīng)網(wǎng)絡（DNN）常常有非常龐大的參數(shù)量。
大家好，我是小寒
今天給大家分享神經(jīng)網(wǎng)絡中的一個關鍵知識點，知識蒸餾
知識蒸餾是一種模型壓縮方法，用于將大型神經(jīng)網(wǎng)絡（教師模型）中的知識轉移到較小的神經(jīng)網(wǎng)絡（學生模型）中。
這一技術能夠在保持或接近原始模型性能的情況下，顯著減小模型的體積，從而提升推理效率。
知識蒸餾在很多場景中非常有用，尤其是在計算資源有限或需要部署到邊緣設備的應用中。
知識蒸餾的背景和動機
在深度學習中，尤其是在計算機視覺和自然語言處理等任務中，深度神經(jīng)網(wǎng)絡（DNN）常常有非常龐大的參數(shù)量。盡管這些大型模型（如BERT、ResNet等）能夠取得非常好的性能，但它們也面臨著存儲、計算和延遲等挑戰(zhàn)。為了克服這一問題，知識蒸餾被提出作為一種方法，通過訓練較小的學生模型來模擬大型教師模型的行為。
知識蒸餾的基本概念
- 教師模型（Teacher Model）
通常是一個預訓練的、復雜的深度神經(jīng)網(wǎng)絡，具有較高的精度，但計算和存儲開銷較大。 - 學生模型（Student Model）
學生模型相對簡單，參數(shù)較少，推理速度更快，目標是通過知識蒸餾從教師模型中獲取知識，提升其性能。 - 軟標簽（Soft Labels）
軟標簽是教師模型輸出的概率分布，而非簡單的類別標簽。
教師模型通常使用 softmax 層生成的概率分布作為軟標簽，這些分布包含了類別間的相對關系。 - 溫度（Temperature）
在蒸餾過程中，通常使用一個溫度參數(shù)來調節(jié)教師模型輸出的概率分布的“平滑程度”。較高的溫度會使得輸出分布更加平滑，從而讓學生模型學習到更多的類間關系。
知識蒸餾的流程
- 訓練教師模型
首先訓練一個大型的、高性能的教師模型。
該模型在給定的訓練數(shù)據(jù)集上表現(xiàn)非常好，具有高精度，但計算開銷較大。 - 生成軟標簽
用教師模型對訓練數(shù)據(jù)進行預測，得到每個樣本的類別概率分布（即軟標簽）。
可以使用 softmax 函數(shù)將教師模型的原始輸出轉換為概率分布，并通過調節(jié)溫度參數(shù)來控制這些概率分布的平滑度。 - 訓練學生模型
使用教師模型生成的軟標簽來訓練一個較小的學生模型。
學生模型的目標是模仿教師模型的輸出，從而盡可能地學習到教師模型的知識。
訓練過程中，學生模型同時會使用真實標簽（硬標簽）和軟標簽進行監(jiān)督學習。 - 損失函數(shù)設計知識蒸餾的損失函數(shù)通常由兩個部分組成。傳統(tǒng)的監(jiān)督損失：計算學生模型輸出與真實標簽之間的交叉熵。蒸餾損失：計算學生模型輸出與教師模型輸出之間的差異，通常使用 KL 散度度量兩個概率分布之間的差異。因此，知識蒸餾的損失函數(shù)通常是這兩個損失的加權和：
溫度的作用
在知識蒸餾中，溫度 T 控制了教師模型輸出的“軟標簽”分布的平滑程度。
較高的溫度會使得輸出的概率分布更加平滑，減少類間的差異，使學生模型能夠學習到更多的類之間的相似性。
- 在高溫度下，教師模型的輸出概率分布更加平滑，類之間的概率差異較小。
- 在低溫度下，輸出概率分布變得更加尖銳，教師模型的預測結果接近于硬標簽。
通過調節(jié)溫度，可以讓學生模型更好地學習到教師模型的知識。
知識蒸餾的優(yōu)點
- 模型壓縮
通過蒸餾，學生模型通常比教師模型更小，參數(shù)數(shù)量更少，可以大幅度降低計算和存儲開銷。 - 提高推理速度
由于學生模型體積較小，推理速度較快，適合部署到移動設備或資源有限的邊緣設備上。
案例分享
以下是一個基于 PyTorch 實現(xiàn)的簡單示例代碼，展示了如何進行神經(jīng)網(wǎng)絡中的知識蒸餾。
首先，定義教師模型和學生模型。
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.nn import functional as F
# 教師模型（較大網(wǎng)絡）
class TeacherModel(nn.Module):
def __init__(self):
super(TeacherModel, self).__init__()
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
self.fc1 = nn.Linear(7*7*64, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = x.view(-1, 7*7*64) # Flatten the tensor
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
# 學生模型（較小網(wǎng)絡）
class StudentModel(nn.Module):
def __init__(self):
super(StudentModel, self).__init__()
self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1)
self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
self.fc1 = nn.Linear(7*7*32, 64)
self.fc2 = nn.Linear(64, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = x.view(-1, 7*7*32) # Flatten the tensor
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x接下來，定義蒸餾損失函數(shù)。
def distillation_loss(y_student, y_teacher, T=2.0, alpha=0.7):
# 計算軟標簽的交叉熵損失
soft_loss = nn.KLDivLoss(reduction='batchmean')(
F.log_softmax(y_student / T, dim=1),
F.softmax(y_teacher / T, dim=1)
)
# 計算真實標簽的交叉熵損失
hard_loss = F.cross_entropy(y_student, torch.argmax(y_teacher, dim=1))
# 綜合蒸餾損失
return alpha * soft_loss + (1 - alpha) * hard_loss接下來定義一個訓練函數(shù)，其中教師模型先訓練好，然后使用蒸餾損失訓練學生模型。
def train(model, device, train_loader, optimizer, epoch, teacher_model=None, T=2.0, alpha=0.7):
model.train()
running_loss = 0.0
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
# 教師模型和學生模型的輸出
output = model(data)
with torch.no_grad(): # 教師模型在蒸餾時不更新參數(shù)
teacher_output = teacher_model(data)
# 計算蒸餾損失
loss = distillation_loss(output, teacher_output, T, alpha)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f"Train Epoch: {epoch} \tLoss: {running_loss / len(train_loader):.6f}")
batch_size = 64
epochs = 10
lr = 0.001
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 數(shù)據(jù)加載
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = datasets.MNIST('.', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('.', train=False, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# 初始化教師模型和學生模型
teacher_model = TeacherModel().to(device)
student_model = StudentModel().to(device)
# 教師模型訓練（簡單訓練）
optimizer_teacher = optim.Adam(teacher_model.parameters(), lr=lr)
teacher_model.train()
for epoch in range(1, epochs + 1):
train_teacher(teacher_model, device, train_loader, optimizer_teacher, epoch)
# 學生模型訓練（蒸餾）
optimizer_student = optim.Adam(student_model.parameters(), lr=lr)
student_model.train()
for epoch in range(1, epochs + 1):
train(student_model, device, train_loader, optimizer_student, epoch, teacher_mod責任編輯：武曉燕
來源：
程序員學長
