《PyTorch深度学习实践》09 多分类问题

1. 说明

本系列博客记录B站课程《PyTorch深度学习实践》的实践代码课程链接请点我

2. 代码

加载了mnist数据库

# ---------------------------
# @Time     : 2022/4/15 20:25
# @Author   : lcq
# @File     : 09_MultiClassification.py
# @Function : 多分类问题
# ---------------------------

import torch
from torchvision import transforms
from torchvision import datasets            # 分割成小批量数据
from torch.utils.data import DataLoader     # 分割成小批量数据
import torch.nn.functional as F
import torch.optim as optim                 # 优化器

# Step1：数据准备
batch_size = 64

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])  # 归一化,均值和方差
train_dataset = datasets.MNIST(root='../dataset/mnist/',    # 表示获取mnist数据库
                               train=True,                  # 表示获取训练集
                               download=True,               # 表示本地没有则下载
                               transform=transform)
train_loader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size)
test_dataset = datasets.MNIST(root='../dataset/mnist/', train=False, download=True, transform=transform)
test_loader = DataLoader(test_dataset, shuffle=False, batch_size=batch_size)


x_data = train_dataset.data.shape
column = x_data[1] * x_data[2]    # 获取样本的特征数量


# Step2： 定义网络模型
class Net(torch.nn.Module):
    def __init__(self, input_column):
        super(Net, self).__init__()
        self.l1 = torch.nn.Linear(input_column, 512)
        self.l2 = torch.nn.Linear(512, 256)
        self.l3 = torch.nn.Linear(256, 128)
        self.l4 = torch.nn.Linear(128, 64)
        self.l5 = torch.nn.Linear(64, 10)

    def forward(self, x, input_column):
        x = x.view(-1, input_column)  # -1其实就是自动获取mini_batch
        x = F.relu(self.l1(x))
        x = F.relu(self.l2(x))
        x = F.relu(self.l3(x))
        x = F.relu(self.l4(x))
        return self.l5(x)  # 最后一层不做激活，不进行非线性变换


model = Net(column)

# Step3： 定义损失函数和优化函数
Loss = torch.nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)


# Step4：训练
def train(epoch):
    running_loss = 0.0
    for batch_index, data in enumerate(train_loader, 0):
        X_input, Y_label = data
        optimizer.zero_grad()                       # 梯度归零
        y_pred = model.forward(X_input, column)     # 预测
        loss = Loss(y_pred, Y_label)

        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if batch_index % 10 == 9:
            print('[%d, %d] loss: %.3f' % (epoch, batch_index, running_loss/10))
            running_loss = 0.0


# Step5: 测试
def test():
    correct = 0
    total = 0
    with torch.no_grad():
        for data in test_loader:
            X_test_input, Y_test_label = data
            y_test_pred = model.forward(X_test_input, column)
            _, predicted = torch.max(y_test_pred.data, dim=1)   # dim=1, 表示求出列的最大值； 返回两个数，第一个为该列最大值，第二个为最大值的行索引
            total += Y_test_label.size(0)
            correct += (predicted == Y_test_label).sum().item()
    print("测试集上的正确率： ", correct/total)


for epoch in range(1):
    train(epoch)
    test()