Technology Sharing

한어Русский языкEnglishFrançaisIndonesianSanskrit日本語DeutschPortuguêsΕλληνικάespañolItalianoSuomalainenLatina

Distillation

Model distillation is a method to optimize the performance of a small model by transferring the knowledge of a large model (teacher model) to a small model (student model). Distillation usually includes the following forms:

1. Soft Label Distillation

The student model is trained through the soft labels of the teacher model so that the student model can learn the output distribution of the teacher model.

import torch
import torch.nn as nn

# 定义教师模型和学生模型
teacher_model = ...
student_model = ...

# 定义损失函数
criterion = nn.KLDivLoss(reduction='batchmean')

# 教师模型生成软标签
teacher_model.eval()
with torch.no_grad():
    teacher_outputs = teacher_model(inputs)
soft_labels = torch.softmax(teacher_outputs / temperature, dim=1)

# 学生模型预测
student_outputs = student_model(inputs)
loss = criterion(torch.log_softmax(student_outputs / temperature, dim=1), soft_labels)

# 反向传播和优化
loss.backward()
optimizer.step()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

2. Feature Distillation

By letting the student model learn the teacher modelmiddle layerfeature representation to optimize the student model performance.

class FeatureExtractor(nn.Module):
    def __init__(self, model):
        super(FeatureExtractor, self).__init__()
        self.features = nn.Sequential(*list(model.children())[:-1])
    
    def forward(self, x):
        return self.features(x)

teacher_feature_extractor = FeatureExtractor(teacher_model)
student_feature_extractor = FeatureExtractor(student_model)

# 获取特征表示
teacher_features = teacher_feature_extractor(inputs)
student_features = student_feature_extractor(inputs)

# 定义特征蒸馏损失
feature_distillation_loss = nn.MSELoss()(student_features, teacher_features)

# 反向传播和优化
feature_distillation_loss.backward()
optimizer.step()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

3. Combined Distillation

Combining soft label distillation and feature distillation, we use the output distribution of the teacher model andFeature RepresentationTo train the student model.

# 定义损失函数
criterion = nn.KLDivLoss(reduction='batchmean')
mse_loss = nn.MSELoss()

# 教师模型生成软标签
teacher_model.eval()
with torch.no_grad():
    teacher_outputs = teacher_model(inputs)
soft_labels = torch.softmax(teacher_outputs / temperature, dim=1)

# 学生模型预测
student_outputs = student_model(inputs)
soft_label_loss = criterion(torch.log_softmax(student_outputs / temperature, dim=1), soft_labels)

# 获取特征表示
teacher_features = teacher_feature_extractor(inputs)
student_features = student_feature_extractor(inputs)
feature_loss = mse_loss(student_features, teacher_features)

# 组合损失
total_loss = soft_label_loss + alpha * feature_loss

# 反向传播和优化
total_loss.backward()
optimizer.step()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

Through the above distillation technology, it is possible to effectivelyOptimizing the modelThe structure reduces computational overhead and improves the model's reasoning speed and deployment efficiency while maintaining model performance.