深度学习入门——手写数字识别MINIST数据集

# 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 定义超参数hyperparameter
BATCH_SIZE = 64 # 每批处理的数据
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 是否用GPU还是CPU训练
EPOCHS = 30 # 训练数据集的轮次数

# 3 构建pipeline，对图像做处理transforms
pipeline = transforms.Compose([
    transforms.ToTensor(), # 将图片转换成tensor
    transforms.Normalize((0.1307,), (0.3081,)) # 正则化（均值，标准差）-->当模型过拟合overfitting时可以降低模型复杂度
])

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

# 下载数据集
train_set = datasets.MNIST("dataset", train=True, download=True, transform=pipeline)

test_set = datasets.MNIST("dataset", train=True, 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：卷积层Kernel）
        self.conv2 = nn.Conv2d(10, 20, 3) #第二个卷积层（10：输入通道，20：输出通道，3：卷积层Kernel）
        self.fc1 = nn.Linear(20*10*10, 500) #第一个全连接层（20*10*10：输入通道，500：输出通道）
        self.fc2 = nn.Linear(500, 10) #第二个全连接层（500：输入通道，10：输出通道【0~9】）
        
    def forward(self, x):
        input_size = x.size(0) # batch_size
        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*10*10=2000
        
        x = self.fc1(x) # 输入：batch*2000 输出：batch*500
        x = F.relu(x) # 激活， 保持shape不变
        
        x = self.fc2(x) # 输入：batch*500，输出：batch*10
        
        output = F.log_softmax(x, dim=1) #计算分类后，每个数字0~9的概率
        
        return output

# 6 定义优化器
model = Digit().to(DEVICE)

optimizer = optim.Adam(model.parameters())

# 7 定义训练函数
def train_model(model, device, train_loader, optimizer, epoch):
    # 模型训练
    model.train()
    for batch_index, (img, target) in enumerate(train_loader):
        # 将数据部署到DEVICE上去
        img, target = img.to(device), target.to(device)
        # 梯度初始化为0
        optimizer.zero_grad()
        # 训练后的结果
        output = model(img)
        # 计算loss
        loss = F.cross_entropy(output, target) #cross_entropy适合多分类问题，将计算结果与真实值对比
        # 反向传播
        loss.backward()
        # 参数优化
        optimizer.step() # 用step方法更新参数
        # 每隔3000张图片打印一次loss
        if batch_index % 3000 == 0:
            print("Train Epoch : {} \t Loss : {:.6f}".format(epoch, loss.item()))

# 8 定义测试方法
def test_model(model, device, test_loader):
    # 模型验证
    model.eval()
    # 初始化正确率
    correct = 0.0
    # 初始化测试loss
    test_loss = 0.0
    with torch.no_grad(): # 测试时不会计算梯度，也不会进行反向传播
        for img, target in test_loader:
            # 部署到DEVICE上
            img, target = img.to(device), target.to(device)
            # 测试数据
            output = model(img)
            # 计算测试损失
            test_loss += F.cross_entropy(output, target).item()
            # 找到概率最大下标
            pred = output.max(1, keepdim=True)[1] #值，索引 
                # pred = output.argmax(dim=1)
                # pred = torch.max(output, dim=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)

运行后即可打印出训练信息

具体搭建视频可参考03-2 轻松学 PyTorch 手写字体识别 MNIST ( 实战 - 上 )_哔哩哔哩_bilibilihttps://www.bilibili.com/video/BV1WT4y177SA