用pytorch实现CNN完成手写数据集分类

初学pytorch，做一个小项目练习一下

一.首先，创建一个python 项目，项目结构如下

data做数据处理

model定义模型

train训练模型代码

Utils一些辅助性的函数

val验证，也就是测试

main.py主函数

（simlpe_cnn.pth是训练之后产生的，是训练好的模型）

二。

在model中定义模型结构

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # 第一个卷积层, 输入通道1（对于灰度图像）, 输出通道32, 卷积核大小3x3
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        # 第二个卷积层, 输入通道32, 输出通道64, 卷积核大小3x3
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        # 第一个池化层, 使用2x2的最大池化
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        # 全连接层, 将64 * 7 * 7维的数据（因为两次池化，28x28的图像变为7x7）连接到100个神经元
        self.fc1 = nn.Linear(64 * 7 * 7, 100)
        # 输出层, 从100个神经元到10个输出类别
        self.fc2 = nn.Linear(100, 10)

    def forward(self, x):
        # 通过第一个卷积层 + ReLU激活 + 池化
        x = self.pool1(F.relu(self.conv1(x)))
        # 通过第二个卷积层 + ReLU激活 + 池化
        x = self.pool1(F.relu(self.conv2(x)))
        # 展平所有除批量外的维度
        x = x.view(-1, 64 * 7 * 7)
        # 通过第一个全连接层 + ReLU激活
        x = F.relu(self.fc1(x))
        # 通过第二个全连接层得到最终输出
        x = self.fc2(x)
        return x

三。在data中定义mnist数据集的下载与预处理

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

def load_data(batch_size=64):
    # 定义数据转换
    # 将数据转换为tensor，并进行标准化（mean和std为0.1307和0.3081是MNIST数据特有的）
    transform = transforms.Compose([
        transforms.ToTensor(),#将numpy转换成tensor
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    # 下载并加载训练集
    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

    # 下载并加载测试集
    test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    return train_loader, test_loader

四。utils中定义一些辅助函数

import torch

def save_model(model, path):
    """
    保存PyTorch模型到指定路径。
    """
    torch.save(model.state_dict(), path)
    print(f"Model saved to {path}")

def load_model(model, path):
    """
    从指定路径加载PyTorch模型。
    """
    model.load_state_dict(torch.load(path))
    model.eval()  # 设置为评估模式
    print(f"Model loaded from {path}")
    return model

def set_device():
    """
    配置和返回适用的设备：GPU如果可用，否则CPU。
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using {device} device")
    return device

五。训练模块

import torch
import torch.nn as nn
import torch.optim as optim

def train_model(model, device, train_loader, epochs=10, lr=0.001):
    """
    训练模型的函数。
    
    参数:
    model - PyTorch模型。
    device - 设备类型，可以是'cuda'或'cpu'。
    train_loader - 训练数据加载器。
    epochs - 训练的总周期数。
    lr - 学习率。
    """
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)

    # 将模型移到正确的设备
    model.to(device)
    
    # 开始训练
    model.train()  # 将模型设置为训练模式
    for epoch in range(epochs):
        total_loss = 0
        for batch_idx, (data, target) in enumerate(train_loader): #批索引，对应索引中所有图片的数据与对应标签值
            # 将数据送到设备
            data, target = data.to(device), target.to(device)
            
            # 前向传播
            output = model(data)
            
            # 计算损失
            loss = criterion(output, target)
            
            # 反向传播和优化
            optimizer.zero_grad()  # 清除之前的梯度
            loss.backward()  # 反向传播
            optimizer.step()  # 更新参数
            
            # 累加损失
            total_loss += loss.item()
        
        # 打印平均损失
        print(f'Epoch {epoch + 1}/{epochs} \t Loss: {total_loss / len(train_loader):.6f}')

六。验证模块

import torch
import torch.nn as nn

def test_model(model, device, test_loader):
    """
    测试模型的函数。

    参数:
    model - PyTorch模型。
    device - 设备类型，可以是'cuda'或'cpu'。
    test_loader - 测试数据加载器。
    """
    model.eval()  # 设置模型为评估模式
    criterion = nn.CrossEntropyLoss()  # 使用交叉熵损失计算损失
    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)  # 计算平均损失
    accuracy = 100. * correct / len(test_loader.dataset)  # 计算准确率
    print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')

    return test_loss, accuracy

七。main.py

import argparse
import torch
from model.SimpleCNN import SimpleCNN
from train.train import train_model
from val.test import test_model
from data.data import load_data
from Utils.utils import set_device, save_model, load_model

def main(args):
    # 设置设备
    device = set_device()
    
    # 加载数据
    train_loader, test_loader = load_data(batch_size=args.batch_size)
    
    # 初始化或加载模型
    model = SimpleCNN()
    if args.model_path:
        model = load_model(model, args.model_path)
    model.to(device)
    
    if args.mode == 'train':
        # 训练模型
        train_model(model, device, train_loader, epochs=args.epochs, lr=args.lr)
        # 保存模型
        save_model(model, 'simple_cnn.pth')
        print(f"Model saved as simple_cnn.pth")
        
    if args.mode == 'test':
        # 测试模型
        test_loss, accuracy = test_model(model, device, test_loader)
        # 输出测试结果
        print(f"Test Loss: {test_loss:.4f}, Test Accuracy: {accuracy:.2f}%")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Train or Test CNN model on MNIST dataset')
    #创建argparse对象，实例化一个ArgumentParser对象，并赋值给parser

    parser.add_argument('--mode', type=str, choices=['train', 'test'], required=True, help='Mode to run the script in training or testing')
    parser.add_argument('--batch_size', type=int, default=64, help='Input batch size for training (default: 64)')
    parser.add_argument('--epochs', type=int, default=10, help='Number of epochs to train (default: 10)')
    parser.add_argument('--lr', type=float, default=0.001, help='Learning rate (default: 0.001)')
    parser.add_argument('--model_path', type=str, default='', help='Path to an existing model to use for testing or continue training')
    args = parser.parse_args()
    #这一行解析命令行输入的参数，并将解析后的参数作为属性保存在 args 对象中。
    main(args)

tips：使用命令行调用main.py，实现不同模式的程序运行

#训练：python main.py --mode train --epochs 15 --lr 0.0005

#测试：python main.py --mode test --model_path simple_cnn.pth