LeNet（后续）和AlexNet、ResNet

优化MNIST数据集训练LeNet

我将数据预处理部分进行了修改

将

    # 原READ IMAGE
    # gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)#将输入的 BGR 格式彩色图像转换为灰度图。
    # gray = 255 - cv.resize(gray, (28, 28), interpolation=cv.INTER_LINEAR)#调整图像尺寸并进行反色处理
    # X = torch.Tensor(gray.reshape(1, 28, 28).tolist())##将处理后的图像转换为 PyTorch 张量。
    # X = X.to(device)

OpenCV 默认以 BGR 顺序读取图像，COLOR_BGR2GRAY转换公式为：灰度值 = 0.299 *B + 0.587 * G + 0.114 * R

输出gray为 NumPy 数组，像素值范围为0-255（0 为黑，255 为白）。

 尺寸调整：cv.resize将灰度图缩放到 28×28 像素，使用双线性插值（INTER_LINEAR）保持平滑。

反色处理：255 - gray将像素值反转，例如：原黑色（0）→ 白色（255）

原灰色（128）→ 灰色（128）

• 原白色（255）→ 黑色（0）

 改为

    # 修改后的图片预处理
    # 对图片进行transform转换 尺寸变换为28X28 转换为Tensor数据类型 Normlize归一化处理(新加入)
    transform = torchvision.transforms.Compose([
        torchvision.transforms.Resize((28, 28)),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize([0.13], [0.30])
    ])
    # transform转换
    X = transform(img)
    X = X.to(device)

 进行图片大小变化，tensor类型转换，均值0.13，标准差0.30归一化，确保像素值的尺度统一，避免大数值像素对模型的影响强于小数值像素

将打开图片的方式cv2，

        # 以cv2打开
        # img = cv.imread("my_predict_image/" + str(i) + ".png")

改成Image.open 

        # 用Image.open打开
        img_path = f"my_predict_image/{i}.png"
        img = Image.open(img_path).convert("L")

 然后用保存的MNIST数据集训练好的LeNet模型，预测图片

 由原来的0.2提升到0.9

MINIST数据集训练AlexNet

代码

model.py

import torch
from torch import nn


# 在model.py中修改
class AlexNet(nn.Module):
    def __init__(self):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),  # 增加卷积核数量
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(128, 256, kernel_size=3, padding=1),  # 增加卷积核数量
            nn.ReLU(inplace=True),

            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),

            nn.Conv2d(256, 128, kernel_size=3, padding=1),  # 增加卷积核数量
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((3, 3))
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(128 * 3 * 3, 512),  # 增加全连接层的神经元数量
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(512, 512),
            nn.ReLU(inplace=True),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        # print("展平后尺寸:", x.shape)  # 打印展平后的尺寸，用于调试
        x = self.classifier(x)
        return x


if __name__ == '__main__':
    model = AlexNet()
    X = torch.rand(size=(1, 1, 28, 28), dtype=torch.float32)
    print("输入尺寸:", X.shape)
    for layer in model.features:
        X = layer(X)
        print(layer.__class__.__name__, '\toutput shape: \t', X.shape)
    X = model.avgpool(X)
    print("avgpool 输出尺寸:", X.shape)
    X = torch.flatten(X, 1)
    print("展平后尺寸:", X.shape)
    print("最终输出尺寸:", model(X).shape)

train.py

import torch
from torch import nn
from torchvision import datasets
from torchvision.transforms import ToTensor
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torch.optim.lr_scheduler import StepLR  # 导入学习率调度器
from model import AlexNet
import numpy as np

def confusion_matrix_scores(confusion_matrix):
    confusion_matrix = np.array(confusion_matrix)
    num_classes = confusion_matrix.shape[0]
    total_samples = np.sum(confusion_matrix)

    per_class_values = []
    for i in range(num_classes):
        tp_i = confusion_matrix[i, i]
        col_i_sum = np.sum(confusion_matrix[:, i])
        precision_i = tp_i / col_i_sum if col_i_sum != 0 else 0
        row_i_sum = np.sum(confusion_matrix[i, :])
        recall_i = tp_i / row_i_sum if row_i_sum != 0 else 0
        if precision_i + recall_i != 0:
            f1_i = 2 * precision_i * recall_i / (precision_i + recall_i)
        else:
            f1_i = 0
        per_class_values.append((precision_i, recall_i, f1_i))

    accuracy = np.trace(confusion_matrix) / total_samples
    precisions, recalls, f1s = zip(*per_class_values)
    micro_precision = np.mean(precisions)
    micro_recall = np.mean(recalls)
    micro_f1 = np.mean(f1s)

    print(f"总准确率：{100*accuracy:.2f}%", end=" ")
    print(f"总精确率：{100*micro_precision:.2f}%", end=" ")
    print(f"总召回率：{100*micro_recall:.2f}%", end=" ")
    print(f"总F1分数：{100*micro_f1:.2f}%")

# DATASET
train_data = datasets.MNIST(
    root="./data",
    train=True,
    download=True,
    transform=ToTensor()
)

test_data = datasets.MNIST(
    root="./data",
    train=False,
    download=True,
    transform=ToTensor()
)

# 参数
batch_size = 128
lr = 0.001  # 调整学习率
epochs = 20

# PREPROCESS
train_dataloader = DataLoader(dataset=train_data, batch_size=batch_size)
test_dataloader = DataLoader(dataset=test_data, batch_size=batch_size)
for X, y in train_dataloader:
    print(X.shape)
    print(y.shape)
    break

# MODEL
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AlexNet().to(device)

# TRAIN MODEL
loss_func = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(params=model.parameters(), lr=lr)
scheduler = StepLR(optimizer, step_size=10, gamma=0.1)  # 学习率调度器

writer = SummaryWriter("logs" + "/e" + str(epochs) + "_bs" + str(batch_size) + "_lr" + str(lr))

def train(dataloader, model, loss_func, optimizer, epoch):
    model.train()
    data_size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)
        y_hat = model(X)
        loss = loss_func(y_hat, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    loss, current = loss.item(), batch * len(X)
    print(f"第{epoch + 1}轮训练\tloss: {loss:>7f}", end="\t")
    writer.add_scalar("train_loss", loss, epoch + 1)

# Test model
def test(dataloader, model, loss_fn):
    confusion_matrix = [[0 for _ in range(10)] for _ in range(10)]
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            pred_classes = pred.argmax(1)
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
            for pred_class, true_class in zip(pred_classes, y):
                confusion_matrix[true_class.item()][pred_class.item()] += 1
    test_loss /= num_batches
    correct /= size
    writer.add_scalar("test_loss", test_loss, epoch + 1)
    writer.add_scalar("test_accuracy", 100 * correct, epoch + 1)
    confusion_matrix_scores(confusion_matrix)

if __name__ == "__main__":
    print("e:" + str(epochs) + "bs:" + str(batch_size) + "lr:" + str(lr))
    for epoch in range(epochs):
        train(train_dataloader, model, loss_func, optimizer, epoch)
        test(test_dataloader, model, loss_func)
        scheduler.step()  # 更新学习率

    # Save models
    torch.save(model.state_dict(), "AlexNet_" + "e" + str(epochs) + "bs" + str(batch_size) + "lr" + str(lr) + ".pth")
    print("e:" + str(epochs) + "bs:" + str(batch_size) + "lr:" + str(lr))
    print("Saved PyTorch AlexNet State")
    writer.close()

添加学习率调度器：使用 StepLR 学习率调度器，每 10 个 epoch 将学习率降低为原来的 0.1 倍。

更新学习率：在每个 epoch 结束时调用 scheduler.step() 来更新学习率。

predict.py

import torch
import cv2 as cv
import torchvision
from PIL import Image

from model import AlexNet
from matplotlib import pyplot as plt
import os
import glob

def predict(true_num ,img, model_name):
    # Loading models
    model = AlexNet()
    model.load_state_dict(torch.load("./" + model_name))

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    # 修改后的图片预处理
    # 对图片进行transform转换 尺寸变换为28X28 转换为Tensor数据类型 Normlize归一化处理(新加入)
    transform = torchvision.transforms.Compose([
        torchvision.transforms.Resize((28, 28)),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize([0.13], [0.30])
    ])
    # transform转换
    X = transform(img)
    X = X.to(device)


    with torch.no_grad():
        pred = model(X)
        # print(pred[0].argmax(0))
        # print(pred)
        return pred[0].argmax(0) == true_num

def test_all_model():
    pth_files = glob.glob(os.path.join("./", '*.pth'))
    file_names = [os.path.basename(pth_file) for pth_file in pth_files]
    print(file_names)
    for model_name in file_names:
        count = 0
        for i in range(10):
            img = cv.imread("my_predict_image/" + str(i) + ".png")
            if predict(i, img, model_name):
                count += 1
        print(model_name, "上的准确度：", count/10)

if __name__ == "__main__":
    count = 0
    for i in range(10):
        # img = cv.imread("my_predict_image/" + str(i) + ".png")
        img_path = f"my_predict_image/{i}.png"
        img = Image.open(img_path).convert("L")
        model_name="AlexNet_e20bs128lr0.001.pth"
        if predict(i, img, model_name):
            count += 1
    print(f"{model_name}, 上的准确度：{count/10}", )

运行train.py

 运行predict.py时

 在执行矩阵乘法时，输入矩阵的形状不匹配。具体来说，self.classifier 里第一个全连接层 nn.Linear(128 * 3 * 3, 512) 期望输入的特征维度是 128 * 3 * 3，但实际输入的特征维度是 128x9，这表明输入到全连接层的数据形状和模型定义的不相符。

假设原始图像预处理后是 [1, 28, 28]（channel=1，height=28，width=28 ），缺少批量维度。

unsqueeze(0)：在第 0 维增加一个维度，将张量变为 [1, 1, 28, 28]，此时 batch_size=1，符合模型输入的维度要求。

即在X = transform(img)后加 .unsqueeze(0)

准确度是 1

原AlexNet模型

1. 输入层（原图）

接收尺寸为 224×224×3（长 × 宽 × 通道数，3 代表 RGB 彩色图像）的原始图像数据，作为网络输入 。

2. 卷积层 C1

• 用 96 个 11×11 卷积核，以 步长（Stride）4 在输入图像上滑动做卷积运算，提取基础特征（如边缘、纹理等）。
• 输出特征图尺寸经计算为 55×55（公式：\((224 - 11) / 4 + 1 = 55\) ），通道数 96，即得到 55×55×96 的特征图 。

3. 池化层（Max pooling）

对 C1 输出的特征图做 最大池化（Max pooling），池化核尺寸通常为 3×3、步长 2（图中未显式标，但 AlexNet 标准流程有），进一步压缩特征、降低计算量，同时保留关键特征，输出仍为 55×55×96（不同参数设置可能有差异，需结合实际） 。

4. 卷积层 C2

• 用 256 个 5×5 卷积核，在池化后特征图上做卷积，提取更复杂特征。
• 输出特征图尺寸经计算（结合池化后输入等）变为 27×27，通道数 256，即 27×27×256 。

5. 池化层（Max pooling）

对 C2 输出做最大池化，池化核 3×3、步长 2，压缩后输出 13×13×256 特征图（尺寸计算：\((27 - 3) / 2 + 1 = 13\) ）。

6. 卷积层 C3

• 用 384 个 3×3 卷积核，在 C2 池化后的特征图上卷积，挖掘更抽象特征。
• 输出保持通道数 384、尺寸 13×13，即 13×13×384 。

7. 卷积层 C4

结构同 C3，同样用 384 个 3×3 卷积核，输出 13×13×384 特征图，继续丰富特征表达 。

8. 卷积层 C5

• 用 256 个 3×3 卷积核，输出 13×13×256 特征图 。

9. 池化层（Max pooling）

对 C5 输出做最大池化，池化核 3×3、步长 2，输出压缩为 6×6×256（或图中简化示意的 13 相关，实际以标准 AlexNet 为准，核心是降维 ），进一步提炼关键特征 。

10. 全连接层 FC6

将池化后二维特征图 展平为一维向量（长度为 6×6×256 等，依实际尺寸），通过 4096 个神经元 的全连接层，把特征映射到高维空间，学习特征间复杂关联，输出 4096 维向量 。

11. 全连接层 FC7

与 FC6 类似，也是 4096 个神经元的全连接层，对 FC6 输出进一步非线性变换、特征融合，输出 4096 维向量 。

12. 全连接层 FC8

用 1000 个神经元 的全连接层（对应 ImageNet 1000 类分类任务 ），通过 Softmax 等激活函数，输出 1000 维概率分布，代表输入图像属于各类别的概率，完成分类预测 。

该模型相较于原始AlexNet的优势；

修改后的 AlexNet 模型针对 MNIST 数据集做了专门优化，相较于原始 AlexNet 有以下显著优势：

1. 适配小尺寸输入图像

原始 AlexNet：
设计用于处理 224×224 的 RGB 图像，第一层卷积核为 11×11，步长 4，会导致 28×28 的 MNIST 图像在第一次卷积后尺寸急剧缩小至约 5×5，丢失大量细节。

修改后模型：
将第一层卷积核改为 3×3，步长 1，保留更多空间信息（28×28 输入→28×28 输出），更适合提取手写数字的精细特征（如笔画断点、连接方式）。

2. 减少参数量，提升训练效率

原始 AlexNet：
全连接层维度为 128×6×6→2048→2048→1000，参数量庞大（约 6000 万参数），对 MNIST 这种小数据集易过拟合，且训练速度慢

修改后模型：

全连接层改为 128×3×3→512→512→10，参数量减少约 90%

3. 优化特征提取策略

卷积核数量调整：
修改后的模型在中间层增加了卷积核数量（如 128→256→256），增强特征表达能力，更适合捕捉手写数字的多样性（如不同书写风格、倾斜角度）。

池化策略优化：
使用 2×2 池化核替代原始的 3×3，在小尺寸图像上保留更多空间分辨率，防止过度下采样导致的信息丢失。

4. 增加模型稳定性

引入 AdaptiveAvgPool2d：
通过AdaptiveAvgPool2d((3, 3))强制输出固定尺寸的特征图，确保全连接层输入维度稳定，避免因输入尺寸波动导致的错误。

Dropout 层增强泛化：
在全连接层添加 Dropout（默认 p=0.5），随机丢弃神经元，减少对特定特征的依赖，提升模型泛化能力。

5. 灰度图处理更高效

输入通道适配：
原始 AlexNet 要求 3 通道 RGB 输入，而 MNIST 是灰度图（1 通道）。修改后的模型将输入通道改为 1，避免不必要的通道扩展，简化计算流程。

6. 训练速度与准确率平衡

实验验证：
在 MNIST 上，修改后的模型通常能在 10-20 个 epoch 内达到 99% 以上准确率，而原始 AlexNet 易陷入过拟合。

总结：

通过缩小卷积核、调整通道数、简化全连接层，修改后的模型更符合 MNIST 数据集的特点（小尺寸、低复杂度）

MINIST数据集训练ResNet

模型

ResNet18的基本含义是，网络的基本架构是ResNet，网络的深度是18层。但是这里的网络深度指的是网络的权重层，也就是包括池化，激活，线性层。而不包括批量化归一层，池化层。

残差连接

 残差连接，这个概念的核心思想是通过添加额外的连接来解决深度神经网络训练中的梯度消失和梯度爆炸等问题，F(x)里面有权重w，梯度下降求的就是w，当对F(x)求导（对样本求导）等于0时，此时F(x)+x的导数，结果为1，此时计算式就不等于0，如果残差映射(F(x))的结果的维度与跳跃连接(x)的维度不同，那咱们是没有办法对它们两个进行相加操作的，必须对x进行升维操作，让他俩的维度相同时才能计算。

代码

model.py

import torch
from torch import nn

class Residual(nn.Module):
    def __init__(self, input_channels, num_channels,
                 use_1x1conv=False, strides=1):
        super().__init__()
        self.conv1 = nn.Conv2d(input_channels, num_channels,
                               kernel_size=3, padding=1, stride=strides)
        self.conv2 = nn.Conv2d(num_channels, num_channels,
                               kernel_size=3, padding=1)
        if use_1x1conv:
            self.conv3 = nn.Conv2d(input_channels, num_channels,
                                   kernel_size=1, stride=strides)
        else:
            self.conv3 = None
        self.bn1 = nn.BatchNorm2d(num_channels)
        self.bn2 = nn.BatchNorm2d(num_channels)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, X):
        Y = self.relu(self.bn1(self.conv1(X)))
        Y = self.bn2(self.conv2(Y))
        if self.conv3:
            X = self.conv3(X)
        Y += X
        return self.relu(Y)


def resnet_block(input_channels, num_channels, num_residuals,
                 first_block=False):
    blk = []
    for i in range(num_residuals):
        if i == 0 and not first_block:
            blk.append(Residual(input_channels, num_channels,
                                use_1x1conv=True, strides=2))
        else:
            blk.append(Residual(num_channels, num_channels))
    return blk


class ResNet18(nn.Module):
    def __init__(self):
        super(ResNet18, self).__init__()
        b1 = nn.Sequential(nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3),
                           nn.BatchNorm2d(64), nn.ReLU(),
                           nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
        b2 = nn.Sequential(*resnet_block(64, 64, 2, first_block=True))
        b3 = nn.Sequential(*resnet_block(64, 128, 2))
        b4 = nn.Sequential(*resnet_block(128, 256, 2))
        b5 = nn.Sequential(*resnet_block(256, 512, 2))
        self.net = nn.Sequential(b1, b2, b3, b4, b5,
                                 nn.AdaptiveAvgPool2d((1, 1)),
                                 nn.Flatten(), nn.Linear(512, 10))

    def forward(self, x):
        logits = self.net(x)
        return logits


if __name__ == '__main__':
    model = ResNet18()
    X = torch.rand(size=(256, 1, 28, 28), dtype=torch.float32)
    for layer in model.net:
        X = layer(X)
        print(layer.__class__.__name__, '\toutput shape: \t', X.shape)
    X = torch.rand(size=(1, 1, 28, 28), dtype=torch.float32)
    print(model(X))

• Residual 类：
  实现 “残差块”（Residual Block），是 ResNet 核心组件。通过 shortcut 连接（conv3 可选的 1×1 卷积），解决深层网络梯度消失、训练难问题，让网络能学习残差（Y += X 实现），结构包含：

  • 2 个 3×3 卷积层（conv1、conv2） + BN 层（bn1、bn2） + ReLU 激活；
  • 可选的 1×1 卷积（conv3）用于匹配维度，保证 shortcut 连接可加。

• resnet_block 函数：
  构建 “残差块组”，通过循环堆叠 Residual 块，控制每个组的通道数、残差块数量。首块（first_block）和组内首块（i == 0 且非首组）会特殊处理（如用 1×1 卷积升维、 stride=2 降采样），实现通道数翻倍、特征图缩小。

• ResNet18 类：
  定义完整的 ResNet-18 网络结构，由 5 个阶段（b1-b5）组成：

  • b1：7×7 卷积 + BN + ReLU + 3×3 最大池化，初步提取特征、缩小尺寸；
  • b2-b5：调用 resnet_block 构建残差块组，逐步提升通道数（64→128→256→512）、缩小特征图；
  • 末尾：自适应平均池化（AdaptiveAvgPool2d）将特征图缩为 1×1，展平（Flatten）后接全连接层（Linear），输出 10 维分类结果（对应 10 分类任务）。

train.py

import torch
from torch import nn
from torchvision import datasets
from torchvision.transforms import ToTensor
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from model import ResNet18
import numpy as np


def confusion_matrix_scores(confusion_matrix):
    # 将 confusion_matrix 转换为 NumPy 数组，方便进行矩阵操作
    confusion_matrix = np.array(confusion_matrix)
    # 获取类别数量
    num_classes = confusion_matrix.shape[0]
    # 计算样本总数
    total_samples = np.sum(confusion_matrix)

    # 存储每个类别的精确率，召回率和F1分数
    per_class_values = []
    for i in range(num_classes):
        # 计算类别 i 的真阳性数（TP）
        tp_i = confusion_matrix[i, i]

        # 计算类别 i 的精确率（Precision）
        col_i_sum = np.sum(confusion_matrix[:, i])  # 计算第 i 列的总和
        precision_i = tp_i / col_i_sum if col_i_sum != 0 else 0  # 添加除零检查

        # 计算类别 i 的召回率（Recall）
        row_i_sum = np.sum(confusion_matrix[i, :])  # 计算第 i 行的总和
        recall_i = tp_i / row_i_sum if row_i_sum != 0 else 0  # 添加除零检查

        # 计算类别 i 的F1分数
        if precision_i + recall_i != 0:
            f1_i = 2 * precision_i * recall_i / (precision_i + recall_i)
        else:
            f1_i = 0

        # 将计算结果存储到列表per_class_values中
        per_class_values.append((precision_i, recall_i, f1_i))
        # 打印类别 i 的各项指标
        # print(f"类别{i}的精确率：{100*precision_i:.2f}%，召回率：{100*recall_i:.2f}%，F1分数：{100*f1_i:.2f}%")

    # 计算总的准确率（Accuracy）
    accuracy = np.trace(confusion_matrix) / total_samples  # np.trace(confusion_matrix) 返回对角线元素之和，即所有TP的总和
    # 将每个类别的精确率、召回率和F1分数解包
    precisions, recalls, f1s = zip(*per_class_values)

    # 计算微平均（micro-average）的精确率、召回率和F1分数
    micro_precision = np.mean(precisions)
    micro_recall = np.mean(recalls)
    micro_f1 = np.mean(f1s)

    print(f"总准确率：{100*accuracy:.2f}%", end=" ")
    print(f"总精确率：{100*micro_precision:.2f}%", end=" ")
    print(f"总召回率：{100*micro_recall:.2f}%", end=" ")
    print(f"总F1分数：{100*micro_f1:.2f}%")


# DATASET
train_data = datasets.MNIST(
    root="D:\深度学习\数据集",
    train=True,
    download=True,
    transform=ToTensor()
)

test_data = datasets.MNIST(
    root="D:\深度学习\数据集",
    train=False,
    download=True,
    transform=ToTensor()
)

# 参数
batch_size = 128
lr = 0.0025
epochs = 20

# PREPROCESS
train_dataloader = DataLoader(dataset=train_data, batch_size=batch_size)
test_dataloader = DataLoader(dataset=test_data, batch_size=batch_size)
for X, y in train_dataloader:
    print(X.shape)  # torch.Size([256, 1, 28, 28])
    print(y.shape)  # torch.Size([256])
    break

# MODEL
device = "cuda" if torch.cuda.is_available() else "cpu"
model = ResNet18().to(device)

# TRAIN MODEL
loss_func = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(params=model.parameters(), lr=lr)

writer = SummaryWriter("logs" + "/e" + str(epochs) + "_bs" + str(batch_size) + "_lr" + str(lr))


def train(dataloader, model, loss_func, optimizer, epoch):
    model.train()
    data_size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)

        y_hat = model(X)
        loss = loss_func(y_hat, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    loss, current = loss.item(), batch * len(X)
    print(f"第{epoch + 1}轮训练\tloss: {loss:>7f}", end="\t")
    writer.add_scalar("train_loss", loss, epoch + 1)


# Test model
def test(dataloader, model, loss_fn):
    confusion_matrix = [[0 for _ in range(10)] for _ in range(10)]
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            pred_classes = pred.argmax(1)
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

            for pred_class, true_class in zip(pred_classes, y):
                confusion_matrix[true_class.item()][pred_class.item()] += 1

    test_loss /= num_batches
    correct /= size
    # print(f"Test Error: Accuracy: {(100 * correct):>0.1f}%, Average loss: {test_loss:>8f}")
    writer.add_scalar("test_loss", test_loss, epoch + 1)
    writer.add_scalar("test_accuracy", 100 * correct, epoch + 1)
    # for i in confusion_matrix:
    #     print(i)
    confusion_matrix_scores(confusion_matrix)


if __name__ == "__main__":
    print("e:" + str(epochs) + "bs:" + str(batch_size) + "lr:" + str(lr))
    for epoch in range(epochs):
        train(train_dataloader, model, loss_func, optimizer, epoch)
        test(test_dataloader, model, loss_func)

    # Save models
    torch.save(model.state_dict(), "ResNet18_" + "e" + str(epochs) + "bs" + str(batch_size) + "lr" + str(lr) + ".pth")
    print("e:" + str(epochs) + "bs:" + str(batch_size) + "lr:" + str(lr))
    print("Saved PyTorch ResNet18 State")
    writer.close()

predict.py

import torch
import cv2 as cv
import torchvision
from PIL import Image

from model import ResNet18
from matplotlib import pyplot as plt
import os
import glob


def predict(true_num, img, model_name):
    # Loading models
    model = ResNet18()
    model.load_state_dict(torch.load("./" + model_name))

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

    # 修改后的图片预处理
    # 对图片进行transform转换 尺寸变换为28X28 转换为Tensor数据类型 Normlize归一化处理(新加入)
    transform = torchvision.transforms.Compose([
        torchvision.transforms.Resize((28, 28)),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize([0.13], [0.30])
    ])
    # transform转换
    X = transform(img).unsqueeze(0)
    X = X.to(device)

    model.eval()
    with torch.no_grad():
        pred = model(X)
        return pred[0].argmax(0) == true_num


def test_all_model():
    pth_files = glob.glob(os.path.join("./", '*.pth'))
    file_names = [os.path.basename(pth_file) for pth_file in pth_files]
    print(file_names)
    for model_name in file_names:
        count = 0
        for i in range(10):
            img = cv.imread("my_predict_image/" + str(i) + ".png")
            if predict(i, img, model_name):
                count += 1
        print(model_name, "上的准确度：", count / 10)


if __name__ == "__main__":
    count = 0
    for i in range(10):
        # 用Image.open打开
        img_path = f"my_predict_image/{i}.png"
        img = Image.open(img_path).convert("L")
        model_name = "ResNet18_e20bs128lr0.0025.pth"
        if predict(i, img, model_name):
            count += 1
    print(f"{model_name}, 上的准确度：{count / 10}", )

与原ResNet18的比较：

1. 输入通道数适配 MNIST 数据集

• 原始 ResNet18 通常用于处理 3 通道的 RGB 图像（如 ImageNet 数据集），因此第一个卷积层的输入通道数为 3。
• 本代码中，ResNet18类的第一个卷积层输入通道数为 1（nn.Conv2d(1, 64, ...)），这是为了适配 MNIST 数据集的单通道灰度图像（28×28 像素）。

2. 输出类别数适配 MNIST 任务

• 原始 ResNet18 用于 1000 类的 ImageNet 图像分类，最终全连接层输出维度为 1000。
• 本代码中，由于 MNIST 是 10 类手写数字分类任务，最终全连接层输出维度为 10（nn.Linear(512, 10)）。

3. 网络块数量与残差连接细节

• 原始 ResNet18 的每个阶段包含 2 个残差块（共 4 个阶段，总计 8 个残差块，加上初始卷积和池化层，总层数为 18），本代码的resnet_block函数实现遵循这一结构（每个阶段num_residuals=2）。
• 但残差块中归一化和激活函数的顺序与原始 ResNet 一致（卷积→批归一化→ReLU），无明显差异。

4. 训练相关的适配

• 原始 ResNet18 训练时使用的图像预处理（如归一化均值和标准差）针对 ImageNet 数据集（均值[0.485, 0.456, 0.406]，标准差[0.229, 0.224, 0.225]）。
• 本代码中，predict.py的图像预处理针对 MNIST 数据集，使用单通道的归一化参数（Normalize([0.13], [0.30])），与 MNIST 的统计特性匹配。

结果：

15轮，128批次，lr0.0025

 30轮，128批次，lr0.0025