自研轻量级神经网络NanoNet：普通显卡也能高效训练

1. 项目起源：为什么我会自研一个“小”神经网络

先说结论：我最近开源了一个自研的轻量级图像分类神经网络，项目名叫NanoNet（这个名字比较朴素，但很直白）。它最大的特点是——用普通消费级显卡就能完成全流程训练，不需要 A100、不需要 80G 显存，一张 RTX 3060 或者 4060 就能跑得很舒服。

先介绍下背景。我平时的工作方向是 CV 相关，经常要验证一些算法想法。但我不是那种能随时调用公司集群的幸运儿，更多时候面对的是自己家里的台式机，显卡也就是游戏卡级别的。这些年深度学习框架越来越成熟，很多人习惯了import torchvision.models.resnet50然后直接 fine-tune，模型仓库里的权重一加载就完事。但问题是，一旦你想改网络结构、想换一个轻量级 backbone、想在边缘设备上部署，这种“拿来主义”就很容易碰壁——你会发现很多开源模型的设计前提是“你有足够的显存和算力”。我甚至见过一个项目，光是 backbone 就占掉 16G 显存，我打开一看，里面还套了三层可变形卷积。

受够了这种反复折腾之后，我决定自己从零写一个神经网络项目，目标非常明确：

• 结构完全自研，从卷积层、激活函数到前向传播逻辑，都自己定义，方便理解、方便魔改；
• 训练显存友好，用普通显卡能训练，不依赖大厂算力；
• 代码开源且干净，不搞复杂的工程封装，让刚上手的人也能读懂每一行。

这里先说明一下，我这个项目不是说比 ResNet、EfficientNet 更强，而是在“个人维护、单卡可训、结构可解释”这三个维度上做到了平衡。如果你是想快速给老板跑个 baseline，那直接torchvision可能更省事；但如果你是学生、独立开发者、或者对神经网络原理有执念的人，这个项目应该能给你不少启发。

2. 核心技术拆解：NanoNet 的网络设计与显存优化方案

2.1 网络结构设计：从零构建前馈与卷积混合体

很多人在设计网络时有个误区：一上来就堆模块。我早期的版本也犯过这个错误，把注意力机制、多尺度特征融合、通道注意力全塞进去，结果在 6G 显存的卡上直接 OOM。后来我做了大量减法设计，最终的网络结构非常克制，整个网络分为三个部分：特征提取层、混合池化层、分类头。

特征提取层是基于卷积神经网络的经典范式：卷积 + 批归一化 + ReLU + 最大池化。但我在最后两个 stage 引入了类似前馈神经网络的全连接旁路结构——具体做法是把特征图拉平后经过一个隐藏层，再与卷积分支的输出做逐元素相加。这样做的好处是：既保留了卷积在局部特征上的优势，又让网络具备了一定的全局建模能力，而不是像纯 CNN 那样只能依赖感受野的堆叠。

你可能想问，为什么不直接用 Transformer 那套自注意力？因为对于“普通显卡可训练”这个前提，自注意力的计算量和显存占用对单卡用户太不友好了，尤其是序列长度一大，显存呈平方级增长。我试过在一个 224x224 输入的任务上，仅仅在最后一层加一个自注意力模块，显存就多了 2G 多，而且训练速度明显变慢。所以在 NanoNet 中，我用了“CNN 为主、全连接旁路为辅”的混合设计，既让网络有足够的表达力，又把显存占用控制在一个很理想的范围内。

2.2 模型参数量的精细化控制

模型参数量是决定训练可行性的核心因素之一。如果直接上一版带可变形卷积的 ResNet50，光这一项就有 2500 万左右的参数量，在 8G 显存的显卡上想要跑一个合理的 batch size，那是想都别想。NanoNet 最终把参数量控制在了180 万左右，是 ResNet50 的十几分之一。

具体怎么做到的呢？我用了两个技巧：

• 深度可分离卷积：相比普通卷积，它把空间卷积和通道卷积解耦，参数量和计算量都能大幅降下来。这在 MobileNet 里已经是成熟方案，我只不过是用它替代了网络中那些“不太关键”的普通卷积层。
• 通道数渐进式缩减：不像 ResNet 那样在初期就使用 64 或 128 个通道，NanoNet 从 24 通道起步，逐 stage 增加到 48、96、192。这样在输入分辨率不高的情况下，前期的计算开销就非常小了。

2.3 显存优化的关键策略

前面提到了显存友好是项目的核心卖点。我在这块踩过很多坑，总结出来的经验主要集中在四个方面：

第一，输入分辨率不求高。很多项目的默认输入是 224x224 甚至 320x320，这在分类任务上确实是主流标准，但对于普通显卡训练来说，分辨率的提升意味着特征图尺寸变大、显存占用几乎是指数级上升。我测试过将输入从 224 降到 160，准确率只损失不到 1 个百分点，但显存占用直接少了 40%。NanoNet 默认输入分辨率是 160x160，如果你显存非常紧张，还可以进一步调到 128。

第二，batch size 用小一点，配合梯度累积。很多教程会告诉你 batch size 不能太小，否则 BN 层的统计量不稳定、模型收敛差。这在理论上是正确的，但对于单卡玩家来说，batch size 往往是被显存硬性限制的。我用梯度累积（gradient accumulation）来解决——每轮只前向传播 8 张图，但梯度累积 4 个 step 再更新一次参数，效果约等于 batch size 32。这一招极大地缓解了显存压力，而且实现起来只需要十几行代码。

第三，混合精度训练不是可选项，是必备项。PyTorch 自带的torch.cuda.amp可以把大部分计算和存储切成 FP16，显存占用几乎减半。我测试过，开启 AMP 之后，原本 6G 显存才能跑的训练，4G 就能跑完，而且 FP16 带来的精度损失在分类任务上几乎可以忽略不计。

第四，避免不必要的中间变量。有些人在写自定义网络时，会把每一层的输出都保存到一个 list 里，方便可视化。这在调试阶段没问题，但在正式训练时，这些中间变量全部会留在计算图里，直到反向传播结束才会释放。我最初的版本不小心在 forward 里多保留了三层 feature map，显存直接多占了 1.5G。后来我改用“只保留需要用于后续计算的变量”，并把不需要反向的中间层包上torch.no_grad()，显存占用就降回了正常水平。

2.4 为什么选择“自研”路线而非现成模型魔改

写到这里，可能有人会问：既然 MobileNet、ShuffleNet 这些轻量级网络已经很成熟了，为什么不直接用来做分类，非要自研一个？

这个问题我也想过很多遍。我的真实想法是：在工程落地的时候，直接用现成的轻量网络确实更省心，但如果你是个人的开源项目，想要保证“任何一行代码你都能解释清楚”，那自研就是必经之路。MobileNet 里的深度可分离卷积虽然效果好，但它有着非常具体的结构细节和初始化策略，遇到问题你很难一句话说清来龙去脉。

而自研网络的好处在于可控性。我可以自由地评估一个模块对显存和准确率的影响，知道为什么加了这个层显存会涨、为什么去掉那个连接准确率会掉。长期来看，你的技术积累不会仅限于“调参”，而是能深入到“设计神经网络”的层面。这也是我把项目取名为 NanoNet 的原因——它小，但五脏俱全，你可以非常轻松地读懂结构并把它迁移到自己的任务上。

3. 完整实操指南：如何复现这个“普通显卡可训练”的项目

3.1 环境准备与依赖安装

我尽量让项目的依赖降到最简。核心环境是：

• Python 3.9 或以上版本
• PyTorch 2.x（我推荐不低于 2.0，因为 AMP 相关的 API 在 2.0 之后更稳定）
• torchvision（用于数据集加载与图像预处理）
• CUDA 11.8 或更新版本，NVIDIA 驱动 520 以上
• einops（可选，主要是让维度变换代码更可读）

安装命令就是一个常规 pip：

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install einops tensorboard

这里想提醒大家一个经常踩的坑：显卡驱动版本和 CUDA 版本的匹配问题。很多人拿到新显卡，安装了最新的驱动，然后在 PyTorch 官网复制了 cu121 的安装命令，结果 import torch 的时候报CUDA driver version is insufficient。我建议先运行nvidia-smi查看驱动支持的最高 CUDA 版本，再选择匹配的 PyTorch 预编译包。这个细节虽然基础，但我几乎每到一次新环境都会有人踩中。

3.2 数据集准备与预处理细节

数据集方面，我默认使用 CIFAR-10 作为验证性实验集，因为它的体积小、类别清晰、不需要额外下载大规模数据。但 NanoNet 本身是一个标准的图像分类网络，换成自己的数据集也只是改一下数据加载器和分类头的问题。这里给出一个通用的数据集组织方式：

data/ ├── train/ │ ├── cat/ │ ├── dog/ │ └── ... └── val/ ├── cat/ ├── dog/ └── ...

把原始图片按类别放到文件夹里，然后使用torchvision.datasets.ImageFolder直接加载即可。我强烈建议在预处理阶段做两件事：

1. 计算均值/方差并进行标准化，而不是使用 ImageNet 的固定均值方差。因为 NanoNet 接受的输入分布可能和 ImageNet 有明显差异，用固定的均值和方差可能导致稳定的性能下降。这里的代码也只需要两行：

   mean, std = [0.4914, 0.4822, 0.4465], [0.2470, 0.2435, 0.2616]

   如果你用的是自己的数据，最好跑一下torchvision.transforms.ToTensor()后统计整个训练集的均值和方差。
2. 数据增强不宜过重。很多人在小规模数据集上疯狂使用 RandomResizedCrop、ColorJitter、RandAugment，结果模型反而更难收敛。对于 NanoNet 来说，我建议使用RandomCrop + Flip + 轻度 ColorJitter即可。

3.3 训练脚本核心逻辑：AMP + 梯度累积

训练脚本是整个项目的关键部分之一，我直接贴出核心代码。这份代码经过我多轮精简，可以当作一个干净利落的单卡训练模板。

import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import datasets, transforms from torch.cuda.amp import GradScaler, autocast def train_one_epoch(model, loader, optimizer, criterion, scaler, accumulation_steps=4): model.train() total_loss, total_correct, total_num = 0, 0, 0 optimizer.zero_grad() for i, (images, labels) in enumerate(loader): images, labels = images.cuda(), labels.cuda() with autocast(): outputs = model(images) loss = criterion(outputs, labels) # 注意这里要做归一化 loss = loss / accumulation_steps scaler.scale(loss).backward() if (i + 1) % accumulation_steps == 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad() total_loss += loss.item() * accumulation_steps total_correct += (outputs.argmax(1) == labels).sum().item() total_num += images.size(0) return total_loss / len(loader), total_correct / total_num

说几个重点：

• GPU 设备：我在脚本里直接用了cuda()，如果是 CPU 环境可以改成images.cpu()，但训练速度会非常慢。建议无条件用 GPU。
• 梯度累积：代码里除以accumulation_steps是为了让梯度累积后的等效 loss 量级和单步一样，避免梯度数值突然变大。这一步如果忘了，最容易导致训练前期爆炸。
• AMP：autocast配合GradScaler，注意梯度累积时 scaler 的 step 只在累积完成后调用一次，否则梯度更新次数会多出几倍，模型可能不收敛。

3.4 模型训练参数与验证结果

我的基准实验配置如下：

• 输入分辨率：160x160
• 初始学习率：1e-3，使用余弦退火调度
• 优化器：AdamW，权重衰减 5e-4
• batch size：32（单批 8，梯度累积 4 步）
• Epochs：120
• GPU：RTX 3060 12G

在这种配置下，NanoNet 在 CIFAR-10 上达到了94.2%的准确率。作为对比，ResNet-18 在不做太多优化时大约是 95.1%，MobileNetV2 大约是 94.5%。可以看到，NanoNet 的性能确实比主流网络低一点，但它的参数量、推理耗电和显存占用都低得多。如果你觉得 94.2% 不够用，可以把输入分辨率提高到 192，并且把最后一个全连接旁路隐藏层的宽度从 512 提到 1024，实测准确率能到 95.0% 左右，显存增加也不超过 1.5G。

3.5 推理部署注意事项

项目自带了一个infer.py示例，把训练好的权重加载进来，对单张图片做预测。这一块没什么特别的，唯一想提醒的是：在推理时同样要做预处理一致性。你训练时用了RandomCrop + Flip，但推理时必须用中心裁剪，否则预测结果会不稳定。我见过很多人训练好好的，推理时准确率突然掉了十几个点，最后发现是预处理不一致。

代码大约长这样：

def preprocess(image_path): image = Image.open(image_path).convert('RGB') transform = transforms.Compose([ transforms.Resize((160, 160)), transforms.ToTensor(), transforms.Normalize(mean, std) ]) return transform(image).unsqueeze(0)

这里不要加 RandomCrop，也不要用 RandomFlip，就用 Resize 和 ToTensor。至于是否要 CenterCrop，看你的训练时是否使用了 RandomResizedCrop 之类的尺度扰动，如果没有就完全可以省略。

4. 常见问题与踩坑实录：单卡训练避坑指南

4.1 OOM（显存不足）排查思路

如果有人问我在这个项目上最常遇到的问题是什么，我一定会说：CUDA out of memory。不管你的卡是 4G、6G、8G 还是 12G，总有办法把它跑满。这里给出一套我的排查顺序：

• 看报错发生在哪一步：如果是在数据加载阶段，一般是 DataLoader 的 worker 数量开太多，或者图像没有 resize 就丢进模型了。如果是在 forward 阶段，说明网络本身太大，优先检查是不是输入分辨率设得太高。
• 逐个组件排查：用torch.cuda.max_memory_allocated()看峰值显存出现在哪个模块前后，就能判断是模型层数太深还是某个中间变量太大。我开发 NanoNet 的时候就是靠这个函数逐一优化掉多余的存储。
• 关掉梯度计算：如果你在训练过程中想跑一次验证，务必用with torch.no_grad()包裹推理过程，否则验证集上的 forward 也会产生计算图，显存直接翻倍有多。

4.2 训练 Loss 不下降或者震荡

这是一个特别常见的问题。我自己在早期训练 NanoNet 时，试过一段时间的 Loss 挂在 1.2 上下纹丝不动。排查后的原因很简单：学习率太大。很多人一上来就按照 ImageNet 预训练模型微调的经验设学习率 1e-3 甚至 1e-2，但你的网络是从零训练，参数随机初始化，学习率过高会直接导致梯度爆炸或者模型震荡在某个局部区域出不来。

建议：

• 先从 1e-3 起，如果 Loss 连续 2 个 epoch 不下降就降到 3e-4；
• 如果你用了 AdamW 并且加了权重衰减，权重衰减不宜超过 5e-4，否则小模型很容易欠拟合；
• 如果你使用了梯度累积，检查累积步数是否正确，以及是否忘记optimizer.zero_grad()。

另外一个被很多人忽略的因素是BN 层的 momentum 默认值。在 batch size 非常小（比如 8 或 4）的时候，BN 层的滑动平均会很不稳定。如果 loss 像过山车一样抖动，我建议把 BN 的momentum从默认的 0.1 调到 0.05 左右，效果很显著。

4.3 显卡能识别但训练速度极慢

这种情况大概率是没有真正调用 GPU 计算。判断方法很简单：

import torch print(torch.cuda.is_available()) # True print(torch.cuda.get_device_name(0)) # 显示你的显卡型号

如果这里显示正常，但训练时 GPU 利用率很低，多半是 DataLoader 的瓶颈。可以调大num_workers到 4 或 8，并且开启pin_memory=True。这个小改动在很多机器上可以把训练速度提升百分五十以上。

还有一种隐蔽的坑：显卡被别的进程占用了显存。你可以用nvidia-smi查看当前显卡的显存占用，如果有其他程序占着大半显存，那你只能等它释放，或者用CUDA_VISIBLE_DEVICES=0设置当前任务可见的 GPU。

4.4 模型保存与加载要注意的坑

最后提一个关于模型保存的常见问题。很多人的保存方式是：

torch.save(model.state_dict(), 'model.pth')

这样做没有错，但如果你在后面想继续训练（断点续训），你还需要保存优化器状态和 epoch 数。我的建议是保存一个字典：

torch.save({ 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'epoch': epoch, 'best_acc': best_acc, }, 'checkpoint.pth')

加载时再依次load_state_dict。这个看似简单的习惯，能帮你省下很多从头再训的时间。

另外，如果你的模型包含自定义类（比如 NanoNet 网络本身），加载权重时一定要保证类定义存在。否则你会遇到一个很经典的报错：AttributeError: 'NanoNet' object has no attribute 'conv1'，实际原因通常是网络结构定义和保存时的不一致。这里建议把网络结构定义封装成一个build_model()函数，保证从 checkpoint 加载时重建的模型结构完全一致。

4.5 快速问题排查速查表

5. 开源项目经验分享：组织好你的代码与文档，别人才能用得起来

5.1 开源仓库的目录结构建议

既然项目强调了“开源”，我就多说一点开源仓库的组织经验。很多人写开源项目只丢一个训练脚本和一个模型定义文件，然后 README 写三行就完事了。这样的项目就算再优秀，也很难辐射到更多用户。

我的 NanoNet 仓库目录结构如下：

nanonet/ ├── README.md ├── LICENSE ├── configs/ │ └── default.yaml ├── models/ │ ├── nanonet.py │ └── layers.py ├── data/ │ └── loaders.py ├── train.py ├── infer.py ├── utils/ │ ├── metrics.py │ └── checkpoints.py └── scripts/ └── train.sh

其中configs/default.yaml用 yaml 管理超参数，比如学习率、batch size、梯度累积步数。这样做的好处是你做实验时不用改 Python 代码，只改配置文件，也方便别人 fork 之后调参。

5.2 README 撰写的三个核心要点

第一，首屏就要说明项目解决的问题。标题“个人开源自研神经网络！普通显卡可训练”能抓眼球，但到了 README 里还是要用一段话说清楚：该网络参数量多少、在什么数据集上达到什么准确率、需要多少显存。这样别人在 30 秒内就能判断你的项目是否满足他的需求。

第二，给出最简复现路径。一段pip install -r requirements.txt+ 一行python train.py --config configs/default.yaml就够了。尽可能不要要求用户先做一堆繁琐的环境配置，门槛越低，使用者越多。

第三，写清楚测试过的环境。比如“本项目在 RTX 3060 12G / CUDA 11.8 / PyTorch 2.0 上验证通过”，但不要写“只支持这个环境”，否则会吓跑一批用 4060 或者 A 系列卡的人。

5.3 开源许可证与协作经验

开源一个项目很简单，但维护起来需要很多心思。我的建议是一开始就选择一个宽松的许可证，比如 MIT，这样别人可以放心用你的代码，也更容易积累社区 star 和反馈。

如果有人提 issue，尽量先要求对方提供完整的环境信息和报错日志。大多数时候不是代码问题，而是环境配置不一致。这并不说别人不会用，而是深度学习领域的环境差异实在太大了，一个 CUDA 小版本就会导致各种奇怪问题。

另外，如果有人想参与贡献，不要急着把大功能拆给他。可以先让他从修文档、补充注释做起。等项目稳定了，再开放模型结构改进类的任务。这样不仅对贡献者友好，也能保证主干代码的质量。

6. 实战落地：从 CIFAR-10 到自己的分类任务需要注意什么

6.1 迁移到自定义数据集的三步法

CIFAR-10 的 100 个测试只是热身，真正有意义的场景是把 NanoNet 迁移到你自己的数据上。我总结了一个三步迁移法：

第一步：改数据加载器。把你的图片按类别放进 train/val 文件夹，写在data/loaders.py里。

第二步：改分类头。NanoNet 最后一个分类头是nn.Linear(192, 10)。如果换成 20 类分类，就改成nn.Linear(192, 20)。这里需要注意的只是分类头的输入维度要和主干输出维度一致。

第三步：调整超参数。自定义数据集一般比 CIFAR-10 更难学，学习率要适当降低到 5e-4 或者 3e-4，训练 epoch 数延长到 150 甚至 200；如果你的数据集很小（比如几千张），建议用上预训练权重做迁移学习——虽然 NanoNet 没有 ImageNet 预训练版本，但你可以先在 CIFAR-10 上训练 100 个 epoch 得到一个特征提取器，然后再 micro-tuning 到你的目标数据集上。

6.2 边缘设备部署可行性

因为 NanoNet 的参数量只有 180 万，它的推理成本非常低。我用 ONNX 导出后，在 CPU 上跑 160x160 输入只需要约 15ms，在树莓派 4B 上大约是 80ms 到 100ms。如果你有 NPU 或 GPU 边缘设备，完全可以跑到实时级别。

ONNX 导出的代码很简单：

import torch.onnx dummy_input = torch.randn(1, 3, 160, 160) torch.onnx.export(model, dummy_input, "nanonet.onnx", input_names=["input"], output_names=["output"], dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}})

这里动态轴（dynamic_axes）建议加上，这样推理时可以用任意 batch size，灵活性更高。

7. 后续规划与个人体会

NanoNet 目前已经做到了“普通显卡可训练”和“结构自研”这两个目标，但距离一个真正的生产级开源项目还有很多路要走。我个人下一步想做的事情有几个：

第一，增加更多 backbone 变体，比如一个极度轻量版（参数量小于 100 万）和一个高精度版（加入 SE 注意力模块），方便不同场景选型。

第二，补充针对检测和分割任务的迁移示例，把 NanoNet 作为骨干网络接到简单分割网络上，验证在嵌入式场景的表现。

第三，完善训练可视化面板，让使用者能直观看到每一层的特征图、梯度的分布，进一步降低理解和调试门槛。

最后再分享一个我的真实感受：很多人以为“自己写神经网络”是一件门槛很高的事情，但当你真的从卷积开始一行一行实现时，你对深度学习的理解会有一个质变。以前调 ResNet 时，你会把某个层删掉试试效果，现在你会先想清楚这个层的设计目的、它和前后层的关系、以及它能带来多大计算开销。这种思考方式上的转变，其实比模型精度提升更有价值。

如果你手上有一张普通的游戏显卡，也想体验一下“自研网络 + 单卡训练”的完整流程，NanoNet 算是一个不错的起点。项目代码已经在开源平台放出来了，直接用git clone拉下来，改一改数据集路径，就能开始你的第一轮训练。别怕 OOM，也别怕 loss 不降，这些坑我都替你踩过了，照着上面的经验来，你的显卡能力远比你以为的更强。