pytorch实现模型蒸馏

#首先导入模块、准备数据
import torch
from torch.utils.data import DataLoader
import torch.utils.data as Data
import torchvision.transforms as transforms
import numpy as np
import os
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

#根据自己情况加载自己的数据集
import torch
from torch.autograd import Variable
#对数据归一化处理
maxT = trainx.max()
print(maxT)
minT = trainx.min()
print(minT)
trainx = (trainx-minT)/(maxT-minT)
testx = (testx-minT)/(maxT-minT)

trainx = torch.tensor(trainx)
trainy = torch.tensor(trainy)
testx = torch.tensor(testx)
testy = torch.tensor(testy)
trainx =  Variable(torch.unsqueeze(trainx, dim=1).float(), requires_grad=False)
trainy = Variable(torch.unsqueeze(trainy, dim=1).float(), requires_grad=False)
testx = Variable(torch.unsqueeze(testx, dim=1).float(), requires_grad=False)
testy = Variable(torch.unsqueeze(testy, dim=1).float(), requires_grad=False)
print(trainx.shape)
print(testx.shape)

建立student模型并初始化，根据自己需求建立，如果studentnet模型效果太差可考虑加深。

class studentNet(nn.Module):
    def __init__(self):
        super(anNet,self).__init__()
        self.conv1 = nn.Conv2d(1,6,3)
        self.pool1 = nn.MaxPool2d(2,1)
        self.fc3 = nn.Linear(6*25*25,2)
    def forward(self,x):
        x = self.conv1(x)
        x = self.pool1(F.relu(x))
        x = x.view(x.size()[0],-1)
        x = self.fc3(x)
        return x
    def initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                torch.nn.init.xavier_normal_(m.weight.data)
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                torch.nn.init.normal_(m.weight.data, 0, 0.01)
                m.bias.data.zero_()