mnist分类 卷积方法 详细注释版2022-04-15

#导入所用库
import os
import torch
import torchvision
from torch import optim, nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.utils.tensorboard import SummaryWriter
import torch.nn.functional as F


#获得数据
def get_data():
    transf = transforms.Compose([transforms.ToTensor(),transforms.Normalize([0.5],0.5)]) #定义对数据的处理，totensor将数据转换成pytorch计算的tensor，normalize对数据进行归一化
    train_dataset = datasets.MNIST(root='./data',train= True,transform=transf,download= False)#定义train_dataset，download，在第一次运行时候设置为True。
    test_dataset = datasets.MNIST(root='./data', train=False, transform=transf, download=False)#定义test_dataset，download，在第一次运行时候设置为True。

    train_dataloader = DataLoader(train_dataset,batch_size=64,shuffle=True,drop_last=True)#定义获取数据方式：dataloader
    test_dataloder = DataLoader(test_dataset,batch_size = 1,shuffle = False,drop_last=True)

    return train_dataloader,test_dataloder

#定义网络模型
#网络模型为三层卷积两层全连接网络，每层卷积之后跟着最大池化和激活
class CONV_NET(nn.Module):       #定义卷积模型类，需要继承nn.Module
    def __init__(self):          #构造函数
        super(CONV_NET, self).__init__() #先调用父类构造函数
        self.conv1 = nn.Conv2d(1, 10, kernel_size = 5)  # 定义第一类卷积层，输入通道为1，输出通道为10，卷积核大小为5
        self.conv2 = nn.Conv2d(10, 20, kernel_size=3)   # 定义第二类卷积层，输入通道为10，输出通道为20，卷积核大小为3
        self.conv3 = nn.Conv2d(20, 50, kernel_size=2)   # 定义第三类卷积层，输入通道为20，输出通道为50，卷积核大小为3
        self.conv2_drop = nn.Dropout2d()                # 二维dropout
        self.conv3_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(200, 50)                   # 第一个全连接层 输入维度200，输出50
        self.fc2 = nn.Linear(50, 10)                    # 第二个全连接层 输入维度50，输出10

    def forward(self, x):          #定义调用模型时数据流向，x为输入的图像
        x = F.relu(F.max_pool2d(self.conv1(x), 2))      #第一个卷积，池化，激活   #输入batch 1 28 28，卷积输出batch 10 24 24 池化输出batch 10 12 12
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))#第二个卷积，池化，激活#输入batch 10 12 12，卷积输出 输出batch 20 10 10，池化输出batch 20 5 5
        # drop不改变维度随机删去一些数据
        x = F.relu(F.max_pool2d(self.conv3_drop(self.conv3(x)), 2))#第三个卷积，池化，激活 #输入 batch 20 5 5 卷积输出 batch 50 4 4  池化 batch 50 2 2
        x = x.view(-1, 200)  # view将数据变为200维的数据   batch 50 2 2变为batch 200
        x = F.relu(self.fc1(x))  # 第一层全连接层 输入维度 batch 200 输出batch 50
        x = F.dropout(x, training=self.training)  #dropout层，防止过拟合
        x = self.fc2(x)            # 第二层全连接层 输入维度 batch 50 输出batch 10
        return x     # 输出

if __name__ =="__main__":   #main 函数，程序入口
    # 指定超参数
    batch_size = 64  #每次处理batch 大小
    lr = 1e-2        #学习率
    epochs = 3       #训练轮数
    log_dir = './logs/' #日志目录
    writer = SummaryWriter(log_dir) #定义日志记录对象
    momentum = 0.5  # 随机梯度下降优化的动量系数

    train_loader , test_loader = get_data()      #获取训练数据和测试数据

    model = CONV_NET()          #构建模型对象
    if os.path.isfile('./model.pth'): #判断当前目录下model.pth是否存在，model.pth为本次实验保存的模型
        model_dict = torch.load('./model.pth') #加载模型，读取为字典
        model.load_state_dict(model_dict)       #将字典的数据赋值为网络模型的参数
        print(model)                        #打印模型
        print('load success')               #输出提示

    loss_func = nn.CrossEntropyLoss()       #定义损失函数
    optimizer = optim.SGD(model.parameters(),lr= lr,momentum=momentum) #定义优化器

    model.train()  #将模型设置为训练模式
    batch = 0       #batch为计数变量
    for i in range(epochs): #循环多少轮
        for data in train_loader: #利用train_dataloader拿数据
            img ,label = data     # 拿到的数据分别是img和label
            image = torchvision.utils.make_grid(img)  #将一个batch的img展开成一张图像 为了后面写入tensorboard
            out = model(img)            #将img输入到模型中得到输出

            loss = loss_func(out,label) #计算输出和真实值之间的误差
            optimizer.zero_grad()       #将当前网络的梯度置为o
            loss.backward()             #误差反向传播，计算网络每层梯度
            optimizer.step()            #按照梯度方向，更新网络参数
            batch += 1                  #计数变量加1
            print('epoch = {},batch = {}, loss = {}'.format(i,batch, loss)) #打印当前轮数，训练次数，损失
            writer.add_scalar('loss_value',loss,global_step=batch)      #使用tensorboard 记录损失
            # writer.add_image('batch_image',image,global_step=batch)   # 使用tensorboard 记录当前参与计算的图像
        torch.save(model.state_dict(),'./model.pth')                    #每训练一轮，保存一次模型
    writer.add_graph(model, input_to_model=img, )                       #tensorboard 记录当前模型计算图
    writer.close()          #关闭tensorboard记录

    model.eval()            #模型设置为写模式
    count = 0               #记录分类正确的个数
    for img ,label in test_loader: #每次拿出一个图像进行测试
        out = model(img)           #单个图像得到分类置信度
        _,predict = torch.max(out,1)    #根据置信度使用max函数将置信度最大的类别作为分类结果
        if predict == label:            #判断预测分类结果和真实的是否一致
            count += 1                  #如果一致则分类正确的计数加1

    print("acc = {}".format(count/len(test_loader)))    #所有测试集数据测试完成后计算准确率