知识蒸馏——pytorch实现

轻量化网络

知识蒸馏可以理解为轻量化网络的一个tricks，轻量化网络是深度学习的一个大的发展趋势，尤其是在移动端，终端边缘计算这种对算力和运算时间有要求的场景中。
轻量化网络可以有以下四种方式实现：
1.压缩已经训练好的模型: 知识蒸馏，权值量化，剪枝，注意力迁移
2.直接训练轻量化网络：SqueezeNet，MobileNet等
3.加速卷积运算：低秩分解
4.硬件部署：Tensorrt，Jetson，Openvino等

知识蒸馏

知识蒸馏在轻量化网络方面具有很高的地位。下图是知识蒸馏的实现过程。

导入包

import torch
from torch import nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader
from torchinfo import summary
from tqdm import tqdm

#设置随机数种子，便于复现
torch.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 使用cudnn加速卷积运算
torch.backends.cudnn.benchmark = True

载入MNIST数据集

from torchvision.transforms.transforms import ToTensor
# 载入训练集
train_dataset = torchvision.datasets.MNIST(
    root = 'dataset/',
    train = True,
    transform = transforms.ToTensor(),
    download=True
)
# 生成测试集
test_dataset = torchvision.datasets.MNIST(
    root = 'dataset/',
    train = False,
    transform = transforms.ToTensor(),
    download=True
)
# 生成dataloader
train_dataloader = DataLoader(dataset=train_dataset,batch_size=32,shuffle=True)
test_dataloader = DataLoader(dataset=test_dataset,batch_size=32,shuffle=True)

构建教师模型

class TeacherModel(nn.Module):
  def __init__(self,in_channels=1,num_classes=10):
    super(TeacherModel,self).__init__()
    self.relu = nn.ReLU()
    self.fc1 = nn.Linear(784,1200)
    self.fc2 = nn.Linear(<