简介这份资源面向深度学习入门者与计算机视觉方向的初学者围绕PyTorch框架下卷积神经网络的实现展开帮助读者理解CNN的基本结构与训练流程。包内共10个文件以4个py脚本和2个pt模型文件为主另含MNIST数据集的图像与标签文件压缩包约19.83MB脚本涵盖网络定义、自动求导、线性层与容器用法等模块pt文件可直接加载已训练权重。内容以MNIST手写数字识别为实验场景完整呈现从数据预处理、模型搭建到损失函数与优化器配置、训练及测试集评估的全过程读者可据此掌握卷积层、池化层、全连接层与ReLU的组合方式并理解归一化、批加载与准确率统计等关键环节。目前已有1605人学习下载适合希望快速跑通第一个CNN项目、并在此基础上调整网络结构与训练策略的读者参考。1. 从一张手写数字图说起PyTorch 卷积神经网络到底能解决什么手头有一批 28×28 的灰度图可能是手写数字、工业质检里的划痕、或者心电图的单导联波形。你想让模型自己学会区分它们而不是靠人写 if-else 规则。这时候 PyTorch 加卷积神经网络CNN就是最常被拿出来的组合。PyTorch 负责自动求导和 GPU 加速CNN 负责从局部像素里提取边缘、纹理、形状这些层级特征。它适合谁适合已经会一点 Python、想从全连接网络往前再走一步的工程师也适合需要把模型导出成 ONNX 再部署到 C 或边缘设备的人。这一章不堆公式只把“为什么是卷积而不是全连接”这件事用工程视角说清楚后面几章再落到环境搭建、最小训练脚本、参数调优和踩坑记录。卷积的核心动作是滑动窗口加权求和。一张 28×28 的图如果直接拉平接全连接层第一层权重矩阵就是 784×512参数量接近 40 万而且每个像素的位置信息被彻底打散。卷积层用 3×3 的核在空间上滑动参数量只有 9 个加上通道数却能在整张图上共享同一组权重。这种参数共享和局部连接让 CNN 对平移有一定容忍度也让训练数据需求大幅下降。PyTorch 的nn.Conv2d把卷积、偏置、初始化都封装好了你只需要关心in_channels、out_channels、kernel_size、stride、padding这几个参数。但真正跑起来之前环境里的 CUDA 和 PyTorch 版本必须对得上否则连第一个卷积都跑不动。2. 把 PyTorch 环境搭到能跑卷积CUDA、cuDNN 与版本对齐2.1 先确认显卡驱动和 CUDA 版本再选 PyTorch 安装命令很多人装 PyTorch 的顺序是反的先pip install torch跑起来发现torch.cuda.is_available()返回 False再回头折腾驱动。正确的顺序是先用nvidia-smi看驱动支持的 CUDA 最高版本再去 PyTorch 官网找对应命令。比如驱动显示 CUDA 12.8你可以装 cu12.8 或 cu12.6 的 PyTorch但不要装 cu11.8否则大概率不匹配。# 查看显卡驱动和最高支持的 CUDA 版本 nvidia-smi # 输出示例 # CUDA Version: 12.8 # 说明驱动可以支持到 CUDA 12.8 的运行时nvidia-smi右上角的 CUDA Version 是驱动能支持的最高运行时版本不是你已经安装的 CUDA Toolkit 版本。PyTorch 自带的 CUDA 运行时是打包在 wheel 里的所以只要你装的 PyTorch 版本对应的 CUDA 不超过这个上限就能直接用。如果你在 WSL 里搭环境nvidia-smi同样能识别到显卡但需要 Windows 侧驱动足够新WSL 内的 CUDA 库由 PyTorch wheel 提供不需要单独装 CUDA Toolkit。2.2 用 conda 或 venv 隔离环境再执行安装命令我一般用 conda 建一个干净环境避免和系统里的 TensorFlow 或其他 PyTorch 版本打架。Python 版本选 3.10 或 3.11太新的 3.13 有些包还没跟上。# 创建并激活环境 conda create -n cnn_lab python3.10 -y conda activate cnn_lab # 安装 PyTorch以 CUDA 12.8 为例具体命令以官网为准 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 # 验证安装 python -c import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))--index-url指向 PyTorch 官方 wheel 源比默认 PyPI 更稳。如果你没有 GPU或者只是想先跑通 CPU 版本把cu128换成cpu即可。验证时如果torch.cuda.is_available()是 False先看驱动版本再看是不是装成了 CPU 版。torch.cuda.get_device_name(0)能打印出显卡型号说明 CUDA 通路已经打通。2.3 安装完成后用一个小卷积做冒烟测试不要等到训练脚本写完才发现环境有问题。装完立刻跑一个最小卷积确认前向和反向都能走通。import torch import torch.nn as nn # 定义一个单层卷积输入 1 通道输出 8 通道核 3x3 conv nn.Conv2d(in_channels1, out_channels8, kernel_size3, padding1).cuda() x torch.randn(4, 1, 28, 28).cuda() # batch4, 1通道, 28x28 y conv(x) print(y.shape) # 期望 torch.Size([4, 8, 28, 28]) # 反向传播冒烟 loss y.sum() loss.backward() print(conv.weight.grad.shape) # 期望 torch.Size([8, 1, 3, 3])这段代码做了三件事把卷积层搬到 GPU、构造一个随机输入、执行前向和反向。padding1让 3×3 卷积的输出空间尺寸和输入保持一致这是后面搭网络时最常用的配置。如果这一步报 CUDA out of memory说明显卡被其他进程占着或者 batch 太大把 4 改成 1 再试。如果报RuntimeError: CUDA error: no kernel image is available基本就是 PyTorch 的 CUDA 版本和显卡算力不匹配换一个 wheel 源重装。3. 用 PyTorch 写一个最小 CNN 训练脚本从 DataLoader 到反向传播3.1 数据集准备与 DataLoader 的四个关键参数以 MNIST 为例torchvision.datasets已经封装好了下载和预处理。但实际项目里你面对的是自己的图片文件夹所以这里同时给出内置数据集和自定义 Dataset 的写法。import torch from torch.utils.data import DataLoader from torchvision import datasets, transforms # 预处理转张量 归一化 transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) # MNIST 全局均值和标准差 ]) train_set datasets.MNIST(root./data, trainTrue, downloadTrue, transformtransform) train_loader DataLoader( train_set, batch_size64, # 每次送进网络的样本数 shuffleTrue, # 训练集打乱验证集不要打乱 num_workers4, # 子进程数Windows 下建议设 0 pin_memoryTrue # 锁页内存加速 CPU 到 GPU 的拷贝 )batch_size直接影响显存占用和梯度稳定性。64 或 128 是常见起点显存不够就减半。shuffleTrue只在训练集用验证集和测试集必须关掉否则评估指标会波动。num_workers在 Linux 下可以设 4 或 8Windows 下如果报错就设 0。pin_memoryTrue配合.cuda()能减少数据传输开销但如果你用 CPU 训练这个参数没有意义。3.2 定义一个三层卷积网络把通道数、池化、全连接串起来下面这个网络结构是很多教程的起点两个卷积块加一个全连接分类头。我把它拆开写方便你改通道数和层数。import torch.nn as nn import torch.nn.functional as F class SimpleCNN(nn.Module): def __init__(self, num_classes10): super().__init__() # 第一个卷积块1 - 32 通道 self.conv1 nn.Conv2d(1, 32, kernel_size3, padding1) self.bn1 nn.BatchNorm2d(32) # 第二个卷积块32 - 64 通道 self.conv2 nn.Conv2d(32, 64, kernel_size3, padding1) self.bn2 nn.BatchNorm2d(64) # 池化层2x2 最大池化 self.pool nn.MaxPool2d(2, 2) # 全连接层经过两次池化28x28 - 14x14 - 7x7 self.fc1 nn.Linear(64 * 7 * 7, 128) self.fc2 nn.Linear(128, num_classes) self.dropout nn.Dropout(0.25) def forward(self, x): x self.pool(F.relu(self.bn1(self.conv1(x)))) # 28x28 - 14x14 x self.pool(F.relu(self.bn2(self.conv2(x)))) # 14x14 - 7x7 x x.view(x.size(0), -1) # 展平保留 batch 维度 x F.relu(self.fc1(x)) x self.dropout(x) x self.fc2(x) return xpadding1保证 3×3 卷积不改变空间尺寸池化层负责降采样。BatchNorm2d放在卷积和激活之间能加快收敛但推理时要记得切换model.eval()。x.view(x.size(0), -1)里的-1让 PyTorch 自动推断展平后的维度前提是前面的空间尺寸算对。如果你改了输入尺寸或池化次数fc1的输入维度要跟着改否则会报 shape 不匹配。Dropout(0.25)在全连接层前随机丢弃四分之一神经元降低过拟合推理时自动关闭。3.3 训练循环损失函数、优化器、设备搬运和日志打印训练循环的骨架就四步前向、算损失、反向、更新。但每一步都有容易翻车的地方。import torch.optim as optim device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleCNN(num_classes10).to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3) for epoch in range(5): model.train() running_loss 0.0 for i, (images, labels) in enumerate(train_loader): images, labels images.to(device), labels.to(device) optimizer.zero_grad() # 清空上一轮梯度 outputs model(images) # 前向 loss criterion(outputs, labels) # 交叉熵损失 loss.backward() # 反向传播 optimizer.step() # 更新参数 running_loss loss.item() if i % 100 99: print(fepoch {epoch1}, step {i1}, loss {running_loss/100:.4f}) running_loss 0.0optimizer.zero_grad()必须放在反向传播之前否则梯度会累加。loss.item()把张量转成 Python 浮点数避免把计算图带进日志。model.train()和model.eval()的切换影响 BatchNorm 和 Dropout 的行为验证时一定要切到eval()。如果你用CrossEntropyLoss标签不需要做 one-hot直接传类别索引即可。学习率1e-3是 Adam 的常用起点如果 loss 震荡就降到1e-4如果下降太慢就升到3e-3试试。4. 卷积神经网络调参避坑从 loss 不降到显存爆炸4.1 现象loss 一直卡在 2.3 附近不降原因通常是数据没有归一化或者标签和输出维度对不上。MNIST 的像素值在 0 到 1 之间如果不做Normalize第一层卷积的输入分布偏移太大梯度更新方向混乱。另一个常见原因是CrossEntropyLoss的输入是 logits你却在最后一层加了Softmax导致数值不稳定。解决方法是去掉模型里的Softmax让损失函数自己处理同时检查transforms.Normalize的均值和标准差是否和数据集匹配。4.2 现象训练到一半报 CUDA out of memory原因可能是 batch_size 太大、模型参数量超过显存、或者没有用torch.no_grad()包住验证循环。验证阶段不需要计算梯度如果不加torch.no_grad()中间激活值会一直保留显存很快爆掉。解决方法是把验证代码放进with torch.no_grad():块里同时把 batch_size 减半或者用torch.cuda.empty_cache()清理缓存。如果还是不够考虑用混合精度训练torch.cuda.amp能把显存占用降下来。4.3 现象验证集准确率比训练集低很多这是过拟合的典型信号。原因可能是训练集太小、模型太复杂、或者没有加正则化。解决方法是加 Dropout、加 L2 正则化在优化器里设weight_decay1e-4、或者做数据增强。PyTorch 的torchvision.transforms里可以加随机旋转、随机裁剪、颜色抖动。如果训练集本身只有几百张图再强的正则化也救不回来这时候要考虑迁移学习用预训练模型提取特征。4.4 现象模型导出 ONNX 后推理结果和 PyTorch 不一致原因通常是导出时没有切到eval()模式或者输入尺寸不固定。BatchNorm 和 Dropout 在训练和推理时的行为不同导出前必须model.eval()。另外torch.onnx.export需要指定input_names、output_names和dynamic_axes如果输入尺寸会变要把 batch 维度或空间维度标成动态。解决方法是先用固定尺寸导出验证一致性后再改动态轴。导出后用onnxruntime跑一遍对比 PyTorch 和 ONNX 的输出差值一般控制在 1e-4 以内算正常。4.5 现象DataLoader 在 Windows 下报多进程错误原因是 Windows 的 spawn 机制和 Linux 的 fork 不同num_workers大于 0 时容易卡死或报BrokenPipeError。解决方法是在if __name__ __main__:下面包住训练代码或者直接把num_workers设成 0。如果数据加载成为瓶颈可以把数据预处理提前做好存成张量文件训练时直接读减少在线增强的开销。5. 把 CNN 推到更高精度迁移学习、混合精度与 ONNX 导出验证5.1 用预训练 ResNet 做迁移学习替换最后一层全连接当你的数据集不是 MNIST 而是几百张工业缺陷图时从零训练 CNN 几乎不可能收敛。常见做法是拿 ImageNet 上预训练的 ResNet18把最后一层换成你的类别数然后冻结前面的卷积层只训练分类头。import torchvision.models as models import torch.nn as nn model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT) # 冻结所有卷积层参数 for param in model.parameters(): param.requires_grad False # 替换最后一层新层默认 requires_gradTrue num_features model.fc.in_features model.fc nn.Linear(num_features, 5) # 假设 5 个缺陷类别 model model.to(device) # 只把 fc 层的参数传给优化器 optimizer optim.Adam(model.fc.parameters(), lr1e-3)冻结参数后反向传播不会更新卷积核显存占用和训练时间都大幅下降。ResNet18_Weights.DEFAULT会自动下载预训练权重如果网络不通可以提前把权重文件放到缓存目录。训练几个 epoch 后如果分类头已经稳定可以解冻最后几个卷积块做微调学习率要调小到1e-4或更低否则会把预训练特征破坏掉。5.2 混合精度训练用 torch.cuda.amp 省显存提速度混合精度让部分计算用 float16部分用 float32能在几乎不损失精度的情况下把显存占用降三成左右。PyTorch 从 1.6 开始内置了torch.cuda.amp用法很简单。scaler torch.cuda.amp.GradScaler() for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() with torch.cuda.amp.autocast(): outputs model(images) loss criterion(outputs, labels) scaler.scale(loss).backward() # 缩放损失防止梯度下溢 scaler.step(optimizer) # 更新参数 scaler.update() # 调整缩放因子autocast()上下文里的运算会自动选择 float16 或 float32。GradScaler负责在反向传播前放大损失避免 float16 梯度变成 0。如果你在验证阶段也用autocast()记得不要调用scaler.step()。混合精度对卷积和矩阵乘法加速明显但对小 batch 或 CPU 训练没有收益。5.3 导出 ONNX 并用 onnxruntime 做数值一致性校验训练完的模型最终要落到推理引擎里ONNX 是常见中间格式。导出时要注意输入输出的名字和动态轴。import torch.onnx model.eval() dummy_input torch.randn(1, 3, 224, 224).to(device) torch.onnx.export( model, dummy_input, resnet18_defect.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}}, opset_version17 )opset_version17支持大多数现代算子如果推理引擎版本较老可以降到 12 或 13。dynamic_axes把 batch 维度标成动态这样导出后的模型可以接受任意 batch 大小。导出后用onnxruntime跑同一组输入对比 PyTorch 和 ONNX 的输出。import onnxruntime as ort import numpy as np sess ort.InferenceSession(resnet18_defect.onnx) ort_inputs {sess.get_inputs()[0].name: dummy_input.cpu().numpy()} ort_output sess.run(None, ort_inputs)[0] with torch.no_grad(): torch_output model(dummy_input).cpu().numpy() print(np.max(np.abs(ort_output - torch_output))) # 期望小于 1e-4如果差值超过 1e-3先检查导出时是否eval()再检查是否有自定义算子不被 ONNX 支持。卷积网络里常见的AdaptiveAvgPool2d在较老 opset 里可能有问题换成固定尺寸的AvgPool2d通常能解决。导出验证通过后这个 ONNX 文件就可以交给 C 或 Python 推理服务用了。我自己的习惯是每次改完网络结构先跑一遍 ONNX 导出和一致性校验再开始长训练。这样能把“训练完才发现部署不了”的后悔药提前吃掉。希望帮到你。本文还有配套的精品资源点击获取