【Pytorch】使用Pytorch进行知识蒸馏

1. 导入各种包

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import torch
import torch.nn.functional as F
import torchvision
from torch import nn
from torchvision import transforms
from torch.utils.data import DataLoader
# from torchinfo import summary
from tqdm import tqdm

2. 设置随机种子

#设置随机种子
torch.manual_seed(0)

# device = torch.device("cuda" if torch.cuda.is_available() else "pcu") # 使用云GPU

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

3. 加载 MNIST 数据集

执行后，MNIST数据集会下载到"dataset/"文件夹下

# 载入训练集
train_dataset = torchvision.datasets.MNIST(
    root="dataset/", # MNIST数据集存放目录
    train=True, #为train=True 时，加载训练集
    transform=transforms.ToTensor(),  # 图像处理、转不同格式显示
    download=True
)
# 载入测试集
test_dataset = torchvision.datasets.MNIST(
    root="dataset/",
    train=False, #为train=False 时，加载测试集
    transform=transforms.ToTensor(), # 图像处理、转不同格式显示
    download=True
)
train_loder = DataLoader(dataset=train_dataset,batch_size=32,shuffle=True)
test_loder  = DataLoader(dataset=test_dataset, batch_size=32,shuffle=False)

#设置随机种子
torch.manual_seed(0)

# device = torch.device("cuda" if torch.cuda.is_available() else "pcu") # 使用云GPU

附录

1. 关于 import torch.nn as nn

torch.nn是用于设置网络中的全连接层的，需要注意在二维图像处理的任务中，全连接层的输入与输出一般都设置为二维张量，形状通常为[batch_size, size]，不同于卷积层要求输入输出是四维张量。

in_features指的是输入的二维张量的大小，即输入的[batch_size, size]中的size。
out_features指的是输出的二维张量的大小，即输出的二维张量的形状为[batch_size，output_size]，当然，它也代表了该全连接层的神经元个数。从输入输出的张量的shape角度来理解，相当于一个输入为[batch_size, in_features]的张量变换成了[batch_size, out_features]的输出张量。

import torch as t
from torch import nn

# in_features由输入张量的形状决定，out_features则决定了输出张量的形状 
connected_layer = nn.Linear(in_features = 64*64*3, out_features = 1)

# 假定输入的图像形状为[64,64,3]
input = t.randn(1,64,64,3)

# 将四维张量转换为二维张量之后，才能作为全连接层的输入
input = input.view(1,64*64*3)
print(input.shape)
output = connected_layer(input) # 调用全连接层
print(output.shape)

# 运行结果:
# input shape is %s torch.Size([1, 12288])
# output shape is %s torch.Size([1, 1])

2. 关于 nn.functional

import torch.nn.functional as F
包含 torch.nn 库中所有函数
同时包含大量 loss 和 activation function

import torch.nn.functional as F

loss_func = F.cross_entropy
loss = loss_func(model(x), y)

loss.backward()

其中 loss.backward() 更新模型的梯度，包括 weights 和 bias

3. 关于from torch.utils.data import DataLoader

数据加载器，结合了数据集和取样器，并且可以提供多个线程处理数据集。
在训练模型时使用到此函数，用来把训练数据分成多个小组，此函数每次抛出一组数据。直至把所有的数据都抛出。就是做一个数据的初始化。