2021-11-05

# （1）加载必要的库
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

# （2）定义超参数
BATCH_SIZE = 16  # 每批处理的数据
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 用GPU还是CPU训练
EPOCHS = 10  # 训练数据集的轮次

# （3）构建pipeline，对图像做处理
pipeline = transforms.Compose([
    transforms.ToTensor(),  # /将图片转化成tensor
    transforms.Normalize((0.1307,), (0.3081,))  # 正则化：降低模型的复杂化
])

# （4）下载、加载数据
from torch.utils.data import DataLoader

# 下载数据集
train_set = datasets.MNIST('data', train=True, download=True, transform=pipeline)
test_set = datasets.MNIST('data', train=False, download=True, transform=pipeline)

# 加载数据
train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=True)

# （5） 构建网络模型
class Digit(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 10, 5)  # 卷积层 1：灰度图片的通道， 10：输出通道， 5： 卷积核
        self.conv2 = nn.Conv2d(10, 20, 3)  # 卷积层 10:输入通道，20 ：输出通道，3：卷积核大小
        self.fc1 = nn.Linear(20 * 10 * 10, 500)  # 全连接层 20*10*10：输入通道，500：输出通道
        self.fc2 = nn.Linear(500, 10)  # 全连接层  500：输入通道，10：输出通道

    def forward(self, x):  # 前向传播
        input_size = x.size(0)  # batch_size  BATCH_SIZE = 16  # 每批处理的数据
        x = self.conv1(x)  # 第一个卷积层 输入：batch*1*28*28 ， 输出：batch*10*24*24(28-5+1 = 24）
        x = F.relu(x)  # 激活函数 保持shape不变，输出：batch*10*24*24
        x = F.max_pool2d(x, 2, 2)  # 池化层 最大池化 输入：batch*10*24*24 输出：batch*10*12*12

        x = self.conv2(x)  # 第二个卷积层 输入：batch*10*12*12 输出：batch*20*10*10（12-3+1=10）
        x = F.relu(x)  # 激活函数

        x = x.view(input_size, -1)  # 拉平，把二维图片拉成一维度的，-1：自动计算维度：20*20*10 = 2000

        x = self.fc1(x)  # 全连接层 输入：batch*2000 输出：batch*500
        x = F.relu(x)  # 激活层，保持shpae不变

        x = self.fc2(x)  # 全连接层2 输入：batch*500  输出：batch*10

        output = F.log_softmax(x, dim=1)  # 计算分类就，每个数字的概率值
        return output  # 把概率值给返回回去

# （6） 定义优化器
model = Digit().to(DEVICE)  # 创建一个模型，并把模型放在设备上面去
optimizer = optim.Adam(model.parameters())  # 定义一个优化器，更新模型的参数。优化器有很多种，这只是其中的一个

# （7）定义训练方法
def train_model(model, device, train_loader, optimizer,epoch):  
    # model:模型， device：运行模型的设备（cpu还是gpu）， optimizer：优化器，epoch：训练集训练的轮数  EPOCHS = 10  # 训练数据集的轮次
    # 模型训练
    model.train()
    for batch_index, (data, target) in enumerate(train_loader):
        # 部署到DEVICE上去
        data, target = data.to(device), target.to(device)
        # 梯度初始化为0
        optimizer.zero_grad()
        # 训练后的结果
        output = model(data)
        # 计算损失
        loss = F.cross_entropy(output, target)  # 此处使用的是交叉熵损失，交叉熵损失适用于多分类的情况
        # 找到概率值最大的下标
        pred = output.max(1, keepdim=True)  # pred = output.argmax(dim=1)  1:表示维度，这里是指横轴
        # 反向传播
        loss.backward()
        # 参数优化 上一步的反向传播哦，不进行参数更新吗？
        optimizer.step()
        if batch_index % 3000 == 0:  # 每处理3000张图片打印一次损失函数
            print('Train Epoch : {} \t Loss : {:.6}'.format(epoch, loss.item()))  # 输出的数字是6位

# (8)定义测试方法
def test_model(model, device, test_loader):
    # 模型验证  验证什么？
    model.eval()
    # 正确率
    correct = 0.0
    # 测试损失  测试损失？这个是做什么用的？
    test_loss = 0.0
    with torch.no_grad():  # 不会计算梯度，也不会进行反向传播
        for data, target in test_loader:
            # 部署到device上
            data, target = data.to(device), target.to(device)
            # 测试数据
            output = model(data)
            # 计算测试损失
            test_loss = F.cross_entropy(output, target).item()
            # 找到概率值最大的下标
            pred = output.max(1, keepdim=True)[1]  # 值，索引
            # 累计正确的值
            correct += pred.eq(target.view_as(pred)).sum().item()
        test_loss /= len(test_loader.dataset)
        print('Test -- Average loss : {:.4f}, Accuracy : {:.3f}\n'.format(test_loss,
                                                                          100 * correct / len(test_loader.dataset)))

# （9）调用方法（7）（8）
for epoch in range(1, EPOCHS + 1):
    train_model(model, DEVICE, train_loader, optimizer, epoch)
    test_model(model, DEVICE, test_loader)