简介面向深度学习初学者与PyTorch使用者这份源码资源覆盖常见模型实现与训练脚本可辅助理解框架核心操作并快速上手实战。压缩包共140个文件以45个ipynb交互式笔记和37个py脚本为主体同时包含14张png示意图、9篇md说明文档以及数据表、配置等辅助材料整体约29.33MB便于按需查阅与本地运行。多数代码可直接跑通遇到环境或版本问题时可参考配套说明或搜索解决方案。包内还保留了Git仓库结构方便追踪修改与版本对比。目前已有1943人学习适合想通过源码研读提升PyTorch编程能力的读者。1. 深度学习PyTorch源码把训练管线拆成零件拿来就能跑有次我接到一个图像分类的活数据现成、模型结构也早就定好了结果光是写数据加载、写训练循环、写评估逻辑、调 checkpoint就花掉一整天最后还因为忘记切model.eval()导致推理结果惨不忍睹。这份深度学习 PyTorch 源码解决的就是这件事它把从数据到权重这条链路上的每个环节——数据管道、模型定义、训练循环、评估、保存加载——都做成了可以独立替换的模块拿到手改改配置就能训练出模型。适合刚接触 PyTorch 的开发者照着跑通全流程也适合技术老手拿它当基线框架快速验证自己的想法。整个项目读下来你会发现真正的重点不在某个模型写得多漂亮而在工程化程度够不够用。2. 架构设计目录结构、配置层与工具链2.1 目录结构模块分离与依赖方向源码用四个一级目录把职责隔开data负责数据读取和预处理models负责网络定义training负责训练和评估utils放配置、日志、检查点工具。入口main.py和推理入口inference.py在根目录配置文件统一放在configs下。. ├── data/ │ ├── dataset.py # 自定义 Dataset扫描目录并读取图片 │ └── transform.py # 训练/推理两套预处理逻辑 ├── models/ │ ├── cnn.py # 基础卷积网络 │ └── resnet.py # 经典残差网络实现 ├── training/ │ ├── trainer.py # 训练循环封装 │ └── evaluate.py # 验证集评估 ├── utils/ │ ├── config.py # yaml 配置加载与类型转换 │ ├── logger.py # 训练日志 │ └── checkpoint.py # 模型保存与恢复 ├── main.py # 训练入口 ├── inference.py # 推理入口 └── configs/ └── default.yaml # 默认训练配置这个结构最值得借鉴的是依赖方向data不依赖modelsmodels不依赖trainingtraining只依赖模型接口而不关心具体网络。我在维护这类项目时最怕的就是所有代码堆在一个文件里改数据要翻训练循环调一个学习率要滚完整个文件。这个分层方式保证了换数据集、换模型、换训练策略这三件事互不干扰。实际开发中你会明显感受到这个设计的好处验证一个新网络结构时只在models/下新增文件训练代码一行不动而调整数据增强策略时只看data/transform.py这一个文件不用去猜测哪里还藏着另一套预处理。对于要拿这份源码做二次开发的人来说理解这个依赖方向比理解任何单个函数都重要。2.2 配置层超参外置与实验复现configs/default.yaml集中管理训练相关的所有超参数。使用 yaml 而不是 argparse 或硬编码最大的优势在于每组实验都可以对应一份配置快照复现结果时直接用同一份配置不用回忆那次跑的时候学习率调成了多少。配置结构上分为data、model、training、inference四个区段data: data_root: ./data/raw batch_size: 64 num_workers: 4 image_size: 224 model: name: cnn num_classes: 10 pretrained: false training: epochs: 50 learning_rate: 0.001 optimizer: adam momentum: 0.9 scheduler: step inference: checkpoint_path: ./checkpoints/best.pth save_dir: ./outputsutils/config.py读取 yaml 并返回嵌套字典main.py通过一个函数把配置逐层取出传给对应模块。这里有一个容易被忽略的细节yaml 读取的字典键值默认是字符串部分解析器对超长数字会解析成字符串类型所以我在config.py里写了一个显式类型转换把所有数值字段int()处理一遍避免训练时在 shape 拼接处突然报类型错误。配置外置最大的现实收益是调参过程变得可追溯跑完一组实验把配置文件和权重一起存档以后找回实验结果只需要看配置。2.3 日志与检查点训练中断的后悔药utils/logger.py负责记录训练指标使用标准的 Python logging 模块同时把输出同步到控制台和文件。训练 50 个 epoch 时不可能一直盯着终端日志文件里完整保留了每个 epoch 的 loss 和 acc事后分析曲线、定位翻车点全靠它。日志的另一层价值是与检查点配合如果某次训练中途断了你能从日志里快速判断断在哪里、当时的指标是多少再决定恢复训练还是重新跑。utils/checkpoint.py是这份源码里我建议你仔细读的一个文件它实现了完整的模型保存和恢复逻辑# utils/checkpoint.py import torch def save_checkpoint(model, optimizer, epoch, loss, path): torch.save({ model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), epoch: epoch, best_loss: loss }, path) print(f[checkpoint] saved at epoch {epoch}) def load_checkpoint(model, path, optimizerNone): checkpoint torch.load(path) model.load_state_dict(checkpoint[model_state_dict]) if optimizer is not None: optimizer.load_state_dict(checkpoint[optimizer_state_dict]) return checkpoint[epoch], checkpoint[best_loss]这段代码有两个关键设计。第一保存时把 model 和 optimizer 的 state_dict 一起存进去而不是只存模型权重。只存权重的话恢复训练后优化器的动量信息和学习率调度状态全部丢失训练后半段会不稳定。第二用 dict 结构封装后续要加学习率调度器的状态直接往 dict 里塞一个键就行。加载时返回epoch和loss方便恢复训练时打印正确的当前进度。很多踩坑现场的根源都在这里这个文件就是用来兜底的。3. 核心模块拆解数据加载、模型构建与训练循环3.1 数据加载Dataset 与 DataLoader 的封装data/dataset.py把数据集按目录结构组织每个子目录名就是类别编号目录下放的图片文件属于该类别。Dataset 类在初始化时扫描所有图片路径在__getitem__中按需读取图片这种设计避免了把所有图片一次性加载进内存。# data/dataset.py import os from PIL import Image from torch.utils.data import Dataset class ImageClassificationDataset(Dataset): def __init__(self, data_root, transformNone): self.samples [] self.transform transform for label in os.listdir(data_root): class_dir os.path.join(data_root, label) if not os.path.isdir(class_dir): continue for fname in os.listdir(class_dir): self.samples.append((os.path.join(class_dir, fname), int(label))) def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, label self.samples[idx] image Image.open(img_path).convert(RGB) if self.transform: image self.transform(image) return image, label这段代码有三点值得注意。第一transform由外部传入Dataset 自己不做数据增强决策这使得同一个 Dataset 可以复用在训练和推理两种场景训练时传带随机增强的 transform验证和推理时传确定性 transform。第二int(label)的写法要求目录名是数字如果你的类别是 a、b、c 这类字符串需要在这里维护一个类别到索引的映射字典。第三__getitem__每次调用都会打开文件所以磁盘 IO 速度会直接影响训练吞吐量这也是num_workers配置尤其重要的原因。data/transform.py中定义了两个函数train_transform和val_transform。训练时使用RandomResizedCrop、RandomHorizontalFlip、ColorJitter做数据增强验证和推理时使用Resize CenterCrop。数据增强是一种正则化手段它能增加样本多样性降低模型过拟合但推理时不能用训练那套随机变换否则结果会和训练指标对不上。main.py中创建 DataLoader 的代码给出了一个合理配置基线from torch.utils.data import DataLoader train_loader DataLoader( dataset, batch_sizecfg[data][batch_size], shuffleTrue, num_workerscfg[data][num_workers], pin_memoryTrue if device cuda else False, )shuffle在训练时打开、验证时关闭这是标准配置pin_memory在 CUDA 环境下开启能减少数据从 CPU 到 GPU 的拷贝时间。DataLoader 的batch_size如果超过显存容量优先调小而不是换 GPU尤其在小数据集上做调试batch size 小一点不影响验证正确性。3.2 模型定义CNN 与残差网络的构建models/cnn.py是一个三层卷积网络专门为小尺寸彩色图像分类设计。输入是 3 通道图像经过三轮卷积和池化后通道数变成 128空间分辨率从 224 降到 28最后通过全连接层输出num_classes个类别的 logits。# models/cnn.py import torch.nn as nn class CNNClassifier(nn.Module): def __init__(self, num_classes10): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2) ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(128 * 28 * 28, 256), nn.ReLU(inplaceTrue), nn.Linear(256, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x这里最需要理解的是128 * 28 * 28这个数字。输入 224 分辨率的图每次MaxPool2d把宽高减半三次池化后是 28通道数从 3 变成 32、64、128所以全连接层的输入维度是 128 通道乘以 28 乘以 28。如果你换了一个输入尺寸这个数字就要跟着重算。我重构时会把这一层的nn.Linear换成AdaptiveAvgPool2d加 flatten 再进全连接这样输入尺寸变了也不用改代码实际效果往往还更好。models/resnet.py实现了带残差连接的经典结构。每个残差块包含两条路径主路径是两个卷积层加 ReLU旁路是一个跳跃连接直接把输入加到输出上。跳跃连接的意义是让梯度可以绕过卷积层直接回传缓解深层网络的梯度消失。这个文件和cnn.py保持一样的对外接口都接收num_classes构造参数这使得训练模块能统一调用而不用区分具体网络类型。3.3 训练循环一个 epoch 的完整节奏training/trainer.py把训练循环封装成类train_one_epoch处理单个 epoch。它在设计上把模型和优化器作为构造参数传入所以同一个 Trainer 既可以训练 CNN也可以训练残差网络。# training/trainer.py def train_one_epoch(self, train_loader, epoch): self.model.train() total_loss, correct, total 0.0, 0, 0 for batch_idx, (images, labels) in enumerate(train_loader): images images.to(self.device) labels labels.to(self.device) self.optimizer.zero_grad() outputs self.model(images) loss self.criterion(outputs, labels) loss.backward() self.optimizer.step() total_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) correct (predicted labels).sum().item() total labels.size(0) if batch_idx % 10 0: print(f[train] epoch {epoch} batch {batch_idx} loss {loss.item():.4f}) avg_loss total_loss / total acc correct / total return avg_loss, acc这十几行代码是整个项目的核心。前向传播产生outputscriterion计算 lossbackward反向传播计算梯度optimizer.step更新权重。这里有一个非常容易被忽略的顺序问题optimizer.zero_grad()必须先于loss.backward()因为 PyTorch 的梯度是累积的不置零的话上一个 batch 的梯度会累加到当前 batch 上导致更新方向混乱。loss.item()用于从张量中取出标量并断开计算图连接这样累加平均损失时不会保留所有中间计算图避免内存爆炸。torch.max(outputs, 1)取 logits 最大的索引作为预测类别。返回值是样本加权的平均损失和准确率调用方把返回值记录到日志里就可以画出 loss 和 acc 曲线。3.4 评估函数验证集上的指标计算training/evaluate.py实现验证评估。它和训练循环结构类似但有两处关键差异model.eval()和torch.no_grad()。# training/evaluate.py def evaluate(model, valid_loader, criterion, devicecpu): model.eval() total_loss, correct, total 0.0, 0, 0 with torch.no_grad(): for images, labels in valid_loader: images images.to(device) labels labels.to(device) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) correct (predicted labels).sum().item() total labels.size(0) return total_loss / total, correct / totalevaluate不需要optimizer.step所以用with torch.no_grad()把整个循环包起来PyTorch 会跳过计算图的构建推理速度更快。model.eval()会切换 BatchNorm 和 Dropout 的行为BatchNorm 在训练模式下使用当前 batch 的均值方差评估模式下使用累积的全局统计量忘记这一行会导致验证准确率明显低于预期。这份源码在每个 epoch 结束后调用evaluate然后根据准确率决定是否保存新的最佳权重。3.5 优化器与学习率调度优化器的选择和学习率策略往往决定训练的收敛质量。源码中通过main.py动态实例化优化器训练配置里写的是字符串adam或sgd# main.py 片段 optimizer_cls getattr(torch.optim, cfg[training][optimizer].title()) optimizer optimizer_cls( model.parameters(), lrcfg[training][learning_rate], momentumcfg[training].get(momentum, 0.9) ) scheduler torch.optim.lr_scheduler.StepLR( optimizer, step_size20, gamma0.1 )getattr动态获取优化器类配置里写adam就对应 Adam写sgd对应 SGD。选择优化器时Adam 收敛快、对学习率不敏感适合刚开始跑通流程SGD 配合 momentum 收敛稳定、泛化更好适合最终实验。StepLR每 20 个 epoch 把学习率乘以gamma也就是降到原来的 0.1 倍这种阶梯式下降在图像分类任务里很常见。学习率调度是训练效果调整的玄学之一StepLR节奏直观CosineAnnealingLR适合大模型或长训练周期。每个 epoch 结束后都要调用scheduler.step()它内部维护 epoch 计数器调用顺序不能跳。4. 环境搭建与运行从零跑通训练和推理4.1 环境配置虚拟环境与依赖安装建议用 conda 创建独立虚拟环境避免不同项目的依赖互相打架。Python 版本选 3.8 或以上然后安装 PyTorch 系列包conda create -n dl_train python3.8 conda activate dl_train pip install torch torchvision numpy pillow pyyaml tqdm matplotlibCPU 环境可以直接用 pip 默认安装跑小数据集完全够用。有 NVIDIA 显卡的话去 PyTorch 官网选择对应 CUDA 版本的安装命令。安装完成后验证环境python -c import torch; print(torch.__version__)如果导入报错最常见的原因是 CUDA 版本和驱动不匹配错误信息里通常会出现 libc10.so 或 libcudart.so 相关的字眼。解决方式不是反复重装而是先确认机器上驱动支持的 CUDA 版本再选对应的 PyTorch 版本。CPU 和 GPU 版本的 PyTorch 在代码层面完全一致要不要 GPU 只是安装命令的区别。环境问题大多是版本错配引起的所以在新建虚拟环境后先把 torch 单独装好、确认能导入再装其他依赖能帮你快速定位是哪一层出了问题。4.2 数据准备与配置文件修改第一次跑通流程建议使用公开的小型图像分类数据集比如十个类别的自然图片。把数据按类别放成下面的目录结构Dataset 就能直接读取data/raw/ ├── 0/ │ ├── img_001.jpg │ └── img_002.jpg ├── 1/ │ └── ... └── 9/然后修改configs/default.yaml中的两个字段data_root指向数据目录model的num_classes改成实际类别数。把training下的epochs暂时改成 5batch_size可以保持 64数据量小的话建议调成 16 或 32避免 batch 内样本分布不均匀导致训练震荡。4.3 启动训练与参数说明一切就绪后运行命令python main.py --config configs/default.yamlmain.py的执行顺序是加载配置、构建数据集和 DataLoader、实例化模型、初始化优化器和调度器、创建 Trainer、进入训练循环。每个 epoch 结束后打印验证集准确率训练中每 10 个 batch 输出一次实时 loss。下表是这份源码里最常调整的几个参数参数默认值调整建议learning_rate0.001收敛太慢可以微调到 0.01loss 震荡就降低到 0.0001batch_size64显存不足时降到 16 或 32num_workers4出现数据加载瓶颈时增加到 8epochs50初次验证链路建议 5image_size224显存紧张可降到 128需同步调整模型的 Linear 层训练完会在checkpoints/下生成best.pth和last.pth两个文件。best.pth是验证集表现最好的权重last.pth是最后一个 epoch 的权重。做实验对比时优先使用best.pth。日志文件在logs/目录下可以找到每个 epoch 的完整记录包括 train loss、val loss、train acc、val acc。4.4 推理、模型导出与可视化inference.py是独立的推理入口接收一张图片和权重文件路径输出预测类别和置信度python inference.py --image path/to/image.jpg --checkpoint checkpoints/best.pth推理核心代码有三个关键细节model.eval()切换批归一化到固定统计量torch.no_grad()关闭计算图构建softmax把 logits 归一化成概率分布。如果推理效果与训练指标差太远按第五节第 4 条的方法检查数据增强的一致性。可视化方面源码提供了一个简单工具读取 logs 目录下的记录绘制 loss 曲线# utils/plot_loss.py import matplotlib.pyplot as plt import json with open(logs/train_log.json) as f: logs json.load(f) plt.plot(logs[train_loss], labeltrain_loss) plt.plot(logs[val_loss], labelval_loss) plt.legend() plt.savefig(logs/loss_curve.png)如果想用 TensorBoardPyTorch 自带的SummaryWriter可以直接接入只需要在训练循环里按 epoch 写入 scalar 值浏览器里就能看到曲线。可视化不是必需环节但它能帮你快速发现过拟合、欠拟合和 loss 不降的问题比盯着终端的数字直观得多。5. 避坑记录训练过程中的常见问题5.1 训练到一半 loss 变成 NaN现象前几个 epoch loss 正常下降第 8 到第 10 个 epoch 时 loss 突然变成 NaN后续所有输出都是 NaN训练等于报废。原因最常见的是梯度爆炸SGD 配合较大的学习率时尤其容易发生。另外输入数据中如果包含损坏的图片读出来的张量可能带了 inf 或 NaN经过几轮迭代后就会污染整个权重。解决先把学习率调低一个数量级比如从 0.001 降到 0.0001重新训练观察是否缓解。然后检查数据集中是否有打不开的图片文件逐个扫描文件完整性并把异常文件剔除。我从这个踩坑里得到的教训是训练脚本开头要加一个数据校验步骤统计每张图片能否成功解码而不是等到训练中途被 NaN 打断才发现。5.2 GPU 利用率低训练速度上不去现象nvidia-smi显示显存占用正常但 GPU 利用率很低一个 epoch 跑得比 CPU 还慢训练瓶颈明显不在计算上。原因数据加载发生在 CPU 上如果 Dataset 的__getitem__需要解码大尺寸图片CPU 处理速度跟不上 GPU 的计算速度GPU 只能空转等待数据。解决把num_workers从 0 调到 4 或 8让多个子进程并行做预处理同时开启pin_memory。如果图片本身分辨率很高在 Dataset 初始化时先统一缩放到较小尺寸减少每次读盘的 IO 量和解码时间。CPU 数据增强也是耗时操作可以换成 GPU 上的增强方案但那是优化到一定程度才需要的。5.3 加载 checkpoint 报 missing key现象训练完保存权重后用推理脚本加载报错提示 missing key 或 unexpected key模型和权重对不上。原因最常见的是训练用了 DataParallel模型权重 key 都带上了module.前缀加载到普通模型上自然对不上。另一种情况是保存时用了整个模型对象而不是state_dict加载环境不同导致结构不匹配。解决统一封装方式训练和推理要么都用 DataParallel要么都不用。保存时只存model.state_dict()。如果已经存了带module.前缀的权重写一段代码去掉前缀再加载from collections import OrderedDict new_state_dict OrderedDict() for k, v in checkpoint_items.items(): name k.replace(module., ) if k.startswith(module.) else k new_state_dict[name] v model.load_state_dict(new_state_dict)5.4 推理结果和训练准确率差距大现象训练集准确率超过百分之九十五但推理单张图片经常出错而且出错的图片看起来并不复杂。原因训练时使用了随机裁剪、随机翻转等数据增强推理时简单Resize后直接输入模型看到的图像分布不一致。另一个常见原因是训练和推理用了不同的归一化参数导致输入分布偏移。解决把训练和推理的 transform 集中到data/transform.py中统一管理验证和推理统一使用中心裁剪加相同的归一化参数。在项目里跑任何新实验前先确认推理脚本的 transform 与训练脚本一致这是我在这个项目上吃过亏后养成的检查习惯。5.5 显存 OOM但调小 batch 仍然报错现象训练时报CUDA out of memory把batch_size从 64 调到 2 依然崩一度以为显存坏了。原因输入图片分辨率过高是常见原因4K 图片直接喂进模型即使 batch 为 2 显存也扛不住。另外梯度累积逻辑如果写错比如没有正确清零梯度计算图会跨 batch 累积显存占用持续上升。解决把image_size降到 128 或 64 做测试显存占用应该明显下降。如果训练策略确实需要大 batch 而显存不够采用梯度累积每 n 个 batch 做一次optimizer.step()注意此时zero_grad()只在这个累积周期开始时调用一次。6. 二次开发改模型、调参、扩展的实操路径6.1 换掉整个骨干网络这份源码的models模块和训练逻辑完全解耦你想把 CNN 换成 Transformer 结构只要在models/下新增一个文件导出和CNNClassifier一样签名的类然后在配置里改model的name字段即可。如果你拿它当 baseline 用推荐先在models/resnet.py基础上改残差块的数量和通道数这是验证成本最低的扩展路径。6.2 把训练逻辑改成验证集早停源码现有逻辑是固定 epoch 数训练。我一般会改trainer.py在evaluate返回的验证准确率超过历史最优时保存 checkpoint连续多轮不提升就终止训练能省下不少算力。改法不复杂evaluate函数返回一个 dicttrainer 内部维护best_acc每轮结束做比较。6.3 从单卡到多卡扩展数据量上来了想用多卡训练在main.py里把模型包一层DataParallel就是最小改动方案if torch.cuda.device_count() 1: model torch.nn.DataParallel(model)注意这行代码要放在创建优化器之前并且检查点保存时按第五节第 3 条的方式统一 key 名。更工程化的方案是DistributedDataParallel但那是另一个量级的改动入门阶段不建议一步到位。我在维护这个源码的过程中最大的习惯养成是每次改模型或改数据增强第一件事就是跑一个最小配置——3 个 epoch、小 batch——确认整条链路没坏再做完整实验。从那以后我再也没有在训练到一半时发现数据读取逻辑出错而白白浪费算力。希望这份源码解析能帮你顺利跑起来也少踩几个我踩过的坑。本文还有配套的精品资源点击获取