PyTorch实现CNN用于MNIST手写数字识别任务的python完整代码

        这段代码使用PyTorch实现了一个简化版的卷积神经网络（CNN），用于MNIST手写数字识别任务。

        MNIST数据集是机器学习领域最经典的手写数字图像分类数据集之一，由美国国家标准与技术研究院（NIST）整理，包含7万张28×28像素的灰度图像，其中6万张用于训练、1万张用于测试，每张图像对应一个0-9的手写数字标签。数据集中的图像经过标准化处理，背景为黑色，数字前景为白色，像素值范围为0-255（通常会归一化至0-1），格式简洁且易于处理，常被用作深度学习入门的“Hello World”级示例（如CNN模型的基础测试），也是评估分类算法性能的基准数据集之一，在PyTorch、TensorFlow等框架中可直接加载调用，适合用于演示数据预处理、模型训练和推理的全流程。

一、代码即运行结果

（一）python代码示例

import torch

import torch.nn as nn

import torch.optim as optim

from torchvision import datasets, transforms

from torch.utils.data import DataLoader

# 超参数设置

batch_size = 64

learning_rate = 0.001

epochs = 10

# 数据预处理与加载

transform = transforms.Compose([

    transforms.ToTensor(),

transforms.Normalize((0.1307,), (0.3081,))  # MNIST均值与标准差

])

train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)

test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 定义CNN模型（简化版）

class Net(nn.Module):

    def __init__(self):

        super(Net, self).__init__()

        self.conv_layers = nn.Sequential(

            nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1),

            nn.ReLU(),

            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),

            nn.ReLU(),

            nn.MaxPool2d(kernel_size=2, stride=2)

        )

        self.fc_layers = nn.Sequential(

            nn.Linear(32 * 7 * 7, 128),

            nn.ReLU(),

            nn.Dropout(0.5),

            nn.Linear(128, 10)

        )

    def forward(self, x):

        x = self.conv_layers(x)

        x = x.view(x.size(0), -1)  # 展平特征图

        x = self.fc_layers(x)

        return x

model = Net()

criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 训练循环

def train(model, device, train_loader, optimizer, epoch):

    model.train()

    for batch_idx, (data, target) in enumerate(train_loader):

        data, target = data.to(device), target.to(device)

        optimizer.zero_grad()

        output = model(data)

        loss = criterion(output, target)

        loss.backward()

        optimizer.step()

        if batch_idx % 100 == 0:

            print(f"Epoch: {epoch}, Batch: {batch_idx}/{len(train_loader)}, Loss: {loss.item():.4f}")

# 测试评估

def test(model, device, test_loader):

    model.eval()

    test_loss = 0

    correct = 0

    with torch.no_grad():

        for data, target in test_loader:

            data, target = data.to(device), target.to(device)

            output = model(data)

            test_loss += criterion(output, target).item()  # 累加损失

            pred = output.argmax(dim=1, keepdim=True)  # 取概率最大的类别

            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    accuracy = correct / len(test_loader.dataset)

    print(f"\nTest Set: Average Loss: {test_loss:.4f}, Accuracy: {accuracy:.2f} ({correct}/{len(test_loader.dataset)})")

# 执行训练与测试

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model.to(device)

for epoch in range(1, epochs+1):

    train(model, device, train_loader, optimizer, epoch)

    test(model, device, test_loader)

下面从数据处理、模型架构、训练流程、评估方法四个方面详细解析。

二、数据处理流程

（一）数据预处理

1.python代码示例：

加载数据集并对数据预处理，如果本地没有加载过，需联网下载。

transform = transforms.Compose([

    transforms.ToTensor(),  # 将图像转换为Tensor，范围[0,1]

transforms.Normalize((0.1307,), (0.3081,))  # 标准化：减去均值0.1307，除以标准差0.3081

])

2.Tensor（张量）

        在这里就要额外拓展说一下，Tensor是什么，之后才对它标准化。

        Tensor（张量）是深度学习框架（如PyTorch、TensorFlow）中用于表示和处理数据的核心数据结构，本质上是一个多维数组，可以理解为向量（一维）、矩阵（二维）在更高维度上的扩展。在图像领域，将图像转换为Tensor是深度学习流程中的关键预处理步骤，其核心作用是将图像的像素值（通常为Python列表、NumPy数组或PIL图像对象）转换为适合神经网络计算的格式：