pytorch入门 - LetNet5 神经网络

import torch
import torch.nn as nn
from torchsummary import summary
class LeNet5(nn.Module):
    def __init__(self, num_classes=10):
        super(LeNet5, self).__init__()
        # 第一卷积层: 输入通道1(灰度图), 输出通道6, 5x5卷积核, padding=2保持尺寸
        self.conv1 = nn.Conv2d(
            in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=2
        )
        self.sig = nn.Sigmoid()  # Sigmoid激活函数
        self.pool = nn.AvgPool2d(kernel_size=2, stride=2)  # 平均池化层
        # 第二卷积层: 输入通道6, 输出通道16, 5x5卷积核, 无padding
        self.conv2 = nn.Conv2d(
            in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0
        )
        self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2)  # 第二个平均池化层
        self.flatten = nn.Flatten()  # 展平层，将多维输入一维化
        # 第一个全连接层: 输入16 * 5 * 5=400, 输出120
        self.f5 = nn.Linear(16 * 5 * 5, 120)
        # 第二个全连接层: 输入120, 输出84
        self.f6 = nn.Linear(120, 84)
        # 输出层: 输入84, 输出类别数
        self.f7 = nn.Linear(84, num_classes)
        self.softmax = nn.Softmax(dim=1)  # Softmax激活函数
    def forward(self, x):
        # 第一卷积块
        x = self.conv1(x)  # 卷积
        x = self.sig(x)     # 激活
        x = self.pool(x)    # 池化
        # 第二卷积块
        x = self.conv2(x)   # 卷积
        x = self.sig(x)     # 激活
        x = self.pool2(x)   # 池化
        # 全连接部分
        x = self.flatten(x)  # 展平
        x = self.f5(x)       # 全连接
        x = self.sig(x)      # 激活
        x = self.f6(x)       # 全连接
        x = self.sig(x)      # 激活
        x = self.f7(x)       # 输出层
        # 注意: 训练时通常不在这里使用softmax，因为CrossEntropyLoss已经包含了softmax
        return x
if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = LeNet5(num_classes=10).to(device)
    summary(model, (1, 28, 28))  # 打印模型结构摘要

import os
import sys
sys.path.append(os.getcwd())  # 添加当前目录到系统路径，以便导入自定义模块
import time
from torchvision.datasets import FashionMNIST  # 导入FashionMNIST数据集
from torchvision import transforms  # 图像预处理
from torch.utils.data import DataLoader, random_split  # 数据加载和划分
import numpy as np
import matplotlib.pyplot as plt  # 绘图
import torch
from torch import nn, optim  # 神经网络和优化器
import copy  # 用于模型参数深拷贝
import pandas as pd  # 数据处理
from LetNet5_model.main import LeNet5  # 导入我们的LeNet5模型
def train_val_date_load():
    # 加载FashionMNIST训练集
    train_dataset = FashionMNIST(
        root="./data",  # 数据存储路径
        train=True,     # 加载训练集
        download=True,  # 自动下载
        transform=transforms.Compose([
            transforms.Resize(size=28),  # 调整大小到28x28
            transforms.ToTensor(),       # 转为Tensor并归一化到[0,1]
        ]),
    )
    # 按8:2划分训练集和验证集
    train_date, val_data = random_split(
        train_dataset,
        [
            int(len(train_dataset) * 0.8),  # 80%训练
            len(train_dataset) - int(len(train_dataset) * 0.8),  # 20%验证
        ],
    )
    # 创建数据加载器
    train_loader = DataLoader(
        dataset=train_date, batch_size=128, shuffle=True, num_workers=1
    )
    val_loader = DataLoader(
        dataset=val_data, batch_size=128, shuffle=True, num_workers=1
    )
    return train_loader, val_loader
def train_model_process(model, train_loader, val_loader, epochs=10):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    optimizer = optim.Adam(model.parameters(), lr=0.001)  # Adam优化器
    criterion = nn.CrossEntropyLoss()  # 交叉熵损失
    model.to(device)  # 模型移到设备
    # 初始化变量记录最佳模型和训练过程
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    train_loss_all = []
    val_loss_all = []
    train_acc_all = []
    val_acc_all = []
    since = time.time()  # 计时开始
    for epoch in range(epochs):
        print(f"Epoch {epoch + 1}/{epochs}")
        # 初始化统计变量
        train_loss = 0.0
        train_correct = 0
        val_loss = 0.0
        val_correct = 0
        train_num = 0
        val_num = 0
        # 训练阶段
        for step, (images, labels) in enumerate(train_loader):
            images, labels = images.to(device), labels.to(device)
            model.train()  # 训练模式
            outputs = model(images)
            pre_lab = torch.argmax(outputs, dim=1)  # 预测标签
            loss = criterion(outputs, labels)  # 计算损失
            optimizer.zero_grad()  # 梯度清零
            loss.backward()        # 反向传播
            optimizer.step()       # 参数更新
            # 统计信息
            train_loss += loss.item() * images.size(0)
            train_correct += torch.sum(pre_lab == labels.data)
            train_num += labels.size(0)
        # 验证阶段
        for step, (images, labels) in enumerate(val_loader):
            images, labels = images.to(device), labels.to(device)
            model.eval()  # 评估模式
            with torch.no_grad():  # 不计算梯度
                outputs = model(images)
                pre_lab = torch.argmax(outputs, dim=1)
                loss = criterion(outputs, labels)
                val_loss += loss.item() * images.size(0)
                val_correct += torch.sum(pre_lab == labels.data)
                val_num += labels.size(0)
        # 记录本轮结果
        train_loss_all.append(train_loss  train_num)
        val_loss_all.append(val_loss / val_num)
        train_acc = train_correct.double() / train_num
        val_acc = val_correct.double() / val_num
        train_acc_all.append(train_acc.item())
        val_acc_all.append(val_acc.item())
        print(f"Train Loss: {train_loss / train_num:.4f}, Train Acc: {train_acc:.4f}, "
              f"Val Loss: {val_loss / val_num:.4f}, Val Acc: {val_acc:.4f}")
        # 更新最佳模型
        if val_acc_all[-1] > best_acc:
            best_acc = val_acc_all[-1]
            best_model_wts = copy.deepcopy(model.state_dict())
    # 训练结束
    time_elapsed = time.time() - since
    print(f"Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s\n"
          f"Best val Acc: {best_acc:.4f}")
    # 保存模型和训练过程
    torch.save(model.state_dict(), "./models/le_net5_best_model.pth")
    train_process = pd.DataFrame({
        "epoch": range(1, epochs + 1),
        "train_loss_all": train_loss_all,
        "val_loss_all": val_loss_all,
        "train_acc_all": train_acc_all,
        "val_acc_all": val_acc_all,
    })
    return train_process
def matplot_acc_loss(train_process):
    # 绘制训练曲线
    plt.figure(figsize=(12, 5))
    # 损失曲线
    plt.subplot(1, 2, 1)
    plt.plot(train_process["epoch"], train_process["train_loss_all"], label="Train Loss")
    plt.plot(train_process["epoch"], train_process["val_loss_all"], label="Val Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.title("Loss vs Epoch")
    plt.legend()
    # 准确率曲线
    plt.subplot(1, 2, 2)
    plt.plot(train_process["epoch"], train_process["train_acc_all"], label="Train Acc")
    plt.plot(train_process["epoch"], train_process["val_acc_all"], label="Val Acc")
    plt.xlabel("Epoch")
    plt.ylabel("Accuracy")
    plt.title("Accuracy vs Epoch")
    plt.legend()
    plt.tight_layout()
    plt.ion()
    plt.show()
    plt.savefig("./models/le_net5_output.png")
if __name__ == "__main__":
    traindatam, valdata = train_val_date_load()  # 加载数据
    result = train_model_process(LeNet5(), traindatam, valdata, 10)  # 训练模型
    matplot_acc_loss(result)  # 绘制曲线 

import os
import sys
sys.path.append(os.getcwd())  # 添加当前目录到系统路径
import torch
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import FashionMNIST
from LetNet5_model.main import LeNet5
def test_data_load():
    # 加载测试集
    test_dataset = FashionMNIST(
        root="./data",
        train=False,  # 测试集
        download=True,
        transform=transforms.Compose([
            transforms.Resize(size=28),
            transforms.ToTensor(),
        ]),
    )
    # 创建测试数据加载器
    test_loader = DataLoader(
        dataset=test_dataset, batch_size=128, shuffle=True, num_workers=1
    )
    return test_loader
def test_model_process(model, test_loader):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    model.eval()  # 评估模式
    correct = 0
    total = 0
    with torch.no_grad():  # 不计算梯度
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)  # 获取预测类别
            total += labels.size(0)
            correct += torch.sum(predicted == labels.data)  # 统计正确数
    accuracy = correct / total * 100
    print(f"Test Accuracy: {accuracy:.2f}%")  # 打印测试准确率
if __name__ == "__main__":
    test_loader = test_data_load()  # 加载测试数据
    model = LeNet5()  # 实例化模型
    model.load_state_dict(torch.load("./models/le_net5_best_model.pth"))  # 加载训练好的权重
    test_model_process(model, test_loader)  # 测试模型