『聚合』 机器学习从入门到放弃：卷积神经网络CNN（二）

import torch
from torch import nn
from torch.nn import functional as F
from torch import optim
import torchvision
from matplotlib import pyplot as plt
from utils import plot_curve, plot_image, one_hot, predict_plot_image

# step 1 : load dataset
batch_size = 512
# https://blog.csdn.net/weixin_44211968/article/details/123739994
# DataLoader 和 dataset 数据集的应用
train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('./data', train=True, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize(
                                       (0.1307,), (0.3081,)
                                   )
                               ])),
    batch_size=batch_size, shuffle=True
)

test_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('./data', train=True, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize(
                                       (0.1307,), (0.3081,)
                                   )
                               ])),
    batch_size=batch_size, shuffle=False
)


# step 2 : CNN网络
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 图片是灰度图片，只有一个通道
        # 第一层
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16,
                               kernel_size=5, stride=1, padding=2)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)

        # 第二层
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32,
                               kernel_size=5, stride=1, padding=2)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        # 全连接层
        self.fc1 = nn.Linear(in_features=7 * 7 * 32, out_features=256)
        self.relufc = nn.ReLU()
        self.fc2 = nn.Linear(in_features=256, out_features=10)

    # 定义前向传播过程的计算函数
    def forward(self, x):
        # 第一层卷积、激活函数和池化
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        # 第二层卷积、激活函数和池化
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        # 将数据平展成一维
        x = x.view(-1, 7 * 7 * 32)
        # 第一层全连接层
        x = self.fc1(x)
        x = self.relufc(x)
        # 第二层全连接层
        x = self.fc2(x)
        return x


# step3: 定义损失函数和优化函数
# 学习率
learning_rate = 0.005
# 定义损失函数，计算模型的输出与目标标签之间的交叉熵损失
criterion = nn.CrossEntropyLoss()
model = CNN()
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9)

# step4: 模型训练
train_loss = []
# 定义迭代次数
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 将神经网络模型 net 移动到指定的设备上。
model = model.to(device)
# 为了和全连接网路的训练次数一致，我们采用相同的迭代次数
total_step = len(train_loader)
num_epochs = 3
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()  # 清空上一个batch的梯度信息
        # 将输入数据 inputs 喂入神经网络模型 net 中进行前向计算，得到模型的输出结果 outputs。
        outputs = model(images)
        # 使用交叉熵损失函数
        loss = criterion(outputs, labels)
        # 使用反向传播算法计算模型参数的梯度信息
        loss.backward()
        # 更新梯度
        optimizer.step()

        train_loss.append(loss.item())
        # 输出训练结果
        if (i + 1) % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, num_epochs, i + 1, total_step,
                                                                     loss.item()))

print('Finished Training')
# 打印 loss 损失图
plot_curve(train_loss)


# step 5 : 准确度测试
# 测试CNN模型
with torch.no_grad(): # 进行评测的时候网络不更新梯度
    correct = 0
    total = 0
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    print('Accuracy of the network on the 10000 test images: {} %'.format(100 * correct / total))

Epoch [1/3], Step [100/118], Loss: 0.4683
Epoch [2/3], Step [100/118], Loss: 0.2781
Epoch [3/3], Step [100/118], Loss: 0.1155
Finished Training
Accuracy of the network on the 10000 test images: 96.47 %