1. 从零到一为什么选择猫狗分类作为边缘AI部署的入门项目猫狗分类这个项目在深度学习圈子里几乎快被玩烂了。但如果你把视角从“跑通一个模型”切换到“让模型在真实硬件上跑起来”事情就完全不一样了。我选择这个组合——TensorFlow 2 训练猫狗分类模型然后部署到 Jetson Nano 上——不是因为它简单恰恰是因为它麻雀虽小五脏俱全数据管道的构建、迁移学习的策略、模型格式的转换、推理引擎的适配、边缘设备的资源约束每一个环节都能让你踩到坑而每一个坑都对应着实际工程中绕不开的问题。先说清楚这个项目到底在做什么。猫狗分类是一个典型的二分类图像任务输入是一张猫或狗的照片输出是它属于哪一类的概率。TensorFlow 2 是目前最主流的深度学习框架之一它的 Keras 高级API让模型搭建变得非常直观。而 Jetson Nano 是一块只有信用卡大小的边缘计算开发板搭载了128核 Maxwell 架构的 GPU算力大约在 472 GFLOPS 左右功耗却只有 5W 到 10W。把训练好的模型塞进这块小板子里跑推理就是所谓的“边缘部署”。那为什么非要部署到边缘设备上直接放在服务器上跑不行吗这里涉及几个很现实的考量。第一是延迟数据在本地处理不需要上传到云端再等结果返回响应速度可以做到毫秒级。第二是隐私图像数据不出本地对于很多场景来说这是硬性要求。第三是成本边缘设备一次性投入之后不需要持续支付云端算力费用。第四是可靠性网络断了也能正常工作。这些理由放在工业质检、智能安防、零售分析等场景里每一个都足够支撑你去做边缘部署。这个项目适合谁如果你已经学过一些深度学习的基础知识能用 TensorFlow 或 PyTorch 搭个简单的 CNN但对模型部署这件事还比较陌生那这个项目就是为你准备的。如果你是一个嵌入式开发者想了解怎么把 AI 模型跑在自己的硬件上这个项目同样适合。甚至如果你只是一个对边缘 AI 感兴趣的爱好者跟着走一遍也能对整条链路有个清晰的认知。我个人的经验是很多人训练模型很熟练但一到部署环节就卡住了。原因在于训练和部署之间有一道“格式鸿沟”——训练用的是浮点权重、Python 运行时、完整的框架依赖而部署需要的是精简的推理引擎、高效的算子实现、尽可能小的内存占用。跨越这道鸿沟需要你理解模型转换、推理优化、硬件适配这几个关键环节。这个项目就是帮你把这道鸿沟填平的。2. 环境搭建与数据准备把地基打牢再动工2.1 训练环境的选择与配置训练环境我建议直接用带 GPU 的机器不管是本地的台式机还是云端的实例都行。TensorFlow 2 的安装现在比以前省心多了pip 一条命令就能搞定但有几个细节需要注意。首先是 CUDA 和 cuDNN 的版本匹配问题。TensorFlow 2.x 每个版本对 CUDA 和 cuDNN 的要求都不一样装错了就是各种报错。我的做法是先确定 TensorFlow 的版本然后去查官方文档的兼容性表格再倒推安装对应版本的 CUDA 和 cuDNN。比如 TensorFlow 2.10 之前的版本在 Windows 上还支持 GPU之后的版本 Windows 原生就不支持了只能用 WSL2。这个坑我踩过折腾了一下午才发现是版本问题。# 创建虚拟环境强烈建议不要污染系统环境 python -m venv tf2_env source tf2_env/bin/activate # Linux/Mac # tf2_env\Scripts\activate # Windows # 安装 TensorFlow 2以 2.13 为例 pip install tensorflow[and-cuda]2.13.0 # 验证 GPU 是否可用 python -c import tensorflow as tf; print(tf.config.list_physical_devices(GPU))如果输出里能看到 GPU 设备说明环境没问题。如果只有 CPU那就得检查 CUDA 版本和驱动了。注意虚拟环境一定要用因为 TensorFlow 的依赖链条很长不同项目对版本的要求可能冲突。我见过太多人因为系统环境里装了好几个版本的 numpy 和 protobuf 导致各种奇怪的错误。2.2 猫狗数据集的获取与组织猫狗数据集最经典的来源是 Kaggle 上的 Dogs vs. Cats 竞赛数据包含 25000 张带标签的训练图片猫和狗各 12500 张。图片尺寸不一有横有竖分辨率也参差不齐。这个数据集的好处是量大、标注干净缺点是有些图片质量不高存在重复和噪声。下载下来之后目录结构通常是这样data/ ├── train/ │ ├── cats/ │ │ ├── cat.0.jpg │ │ ├── cat.1.jpg │ │ └── ... │ └── dogs/ │ ├── dog.0.jpg │ ├── dog.1.jpg │ └── ... └── validation/ ├── cats/ └── dogs/我习惯用tf.keras.utils.image_dataset_from_directory来加载数据它会自动根据子目录名生成标签非常方便。但有一个细节要注意默认情况下它会把图片缩放到统一尺寸如果你的原始图片长宽比差异很大直接缩放会导致变形。我的做法是先统一缩放到一个稍大的尺寸比如 256x256然后在数据增强阶段再做随机裁剪到 224x224这样既能保证输入尺寸一致又能保留一定的空间信息。import tensorflow as tf IMG_SIZE (224, 224) BATCH_SIZE 32 train_ds tf.keras.utils.image_dataset_from_directory( data/train, image_sizeIMG_SIZE, batch_sizeBATCH_SIZE, label_modebinary, shuffleTrue, seed42 ) val_ds tf.keras.utils.image_dataset_from_directory( data/validation, image_sizeIMG_SIZE, batch_sizeBATCH_SIZE, label_modebinary, shuffleFalse )2.3 数据管道的性能优化数据管道的效率直接决定了训练速度。很多人模型结构调了半天结果发现瓶颈在数据加载上。TensorFlow 的tf.dataAPI 提供了几个关键操作来优化管道cache()把数据缓存到内存或磁盘避免每个 epoch 都重新读取prefetch()让数据加载和模型计算重叠进行shuffle()打乱数据顺序注意 buffer_size 要足够大map()配合num_parallel_calls并行处理数据增强AUTOTUNE tf.data.AUTOTUNE train_ds train_ds.cache().shuffle(1000).prefetch(buffer_sizeAUTOTUNE) val_ds val_ds.cache().prefetch(buffer_sizeAUTOTUNE)这里有个经验之谈如果你的数据集不大比如几千张图cache()直接放内存就行。但如果数据集很大内存放不下就得用cache(filename)缓存到磁盘或者干脆不用 cache靠 prefetch 来掩盖 IO 延迟。我在 25000 张图的规模下测试过缓存到内存大概占用 2-3GB一般机器扛得住。3. 模型设计与训练策略迁移学习是王道3.1 为什么不用从零搭建 CNN从零搭建一个 CNN 来做猫狗分类理论上可行但实际效果往往不如迁移学习。原因很简单猫狗数据集只有 25000 张图对于训练一个深度网络来说数据量偏少容易过拟合。而迁移学习可以利用在 ImageNet 上预训练好的权重这些权重已经学到了大量的底层特征——边缘、纹理、形状——这些特征对猫狗分类同样有效。我试过从零训练一个 5 层卷积的网络验证集准确率卡在 85% 左右就上不去了。换成 MobileNetV2 做迁移学习同样的数据准确率轻松到 97% 以上。这个差距在部署到边缘设备时更明显因为边缘设备对模型大小和计算量有严格限制迁移学习让你可以用更小的模型达到更好的效果。3.2 MobileNetV2 的结构与选择理由MobileNetV2 是 Google 在 2018 年提出的轻量级网络核心创新是倒残差结构Inverted Residual和线性瓶颈Linear Bottleneck。简单来说它在传统残差块的“压缩-卷积-扩张”基础上做了反转先用 1x1 卷积升维再用 3x3 深度可分离卷积提取特征最后用 1x1 卷积降维。这样做的好处是在保持表达能力的同时大幅减少计算量。为什么选 MobileNetV2 而不是 ResNet50 或 VGG16看一组数据就明白了模型参数量计算量 (FLOPs)ImageNet Top-1VGG16138M15.5G71.3%ResNet5025.6M4.1G76.1%MobileNetV23.4M0.3G72.0%MobileNetV2 (1.4x)6.9M0.6G75.0%MobileNetV2 的参数量只有 ResNet50 的七分之一计算量只有十三分之一但精度差距并不大。在 Jetson Nano 这种算力有限的设备上这个差距直接决定了你能不能跑到实时帧率。3.3 迁移学习的两种策略与实操迁移学习有两种常见策略特征提取和微调。特征提取是把预训练模型当作固定的特征提取器只训练新加的分类头。具体做法是加载 MobileNetV2 的卷积基去掉顶部的全连接层冻结所有层然后在上面加一个全局平均池化层和一个二分类的全连接层。base_model tf.keras.applications.MobileNetV2( input_shape(224, 224, 3), include_topFalse, weightsimagenet ) base_model.trainable False # 冻结卷积基 model tf.keras.Sequential([ base_model, tf.keras.layers.GlobalAveragePooling2D(), tf.keras.layers.Dropout(0.2), tf.keras.layers.Dense(1, activationsigmoid) ])微调是在特征提取训练收敛之后解冻卷积基的顶部若干层用很小的学习率继续训练。这样可以让预训练的特征稍微适应你的特定数据分布通常能再提升一两个百分点的准确率。# 先训练分类头 model.compile( optimizertf.keras.optimizers.Adam(1e-3), lossbinary_crossentropy, metrics[accuracy] ) history model.fit(train_ds, validation_dataval_ds, epochs10) # 然后解冻顶部层做微调 base_model.trainable True for layer in base_model.layers[:-30]: layer.trainable False model.compile( optimizertf.keras.optimizers.Adam(1e-5), # 学习率要小 lossbinary_crossentropy, metrics[accuracy] ) fine_tune_history model.fit(train_ds, validation_dataval_ds, epochs10)注意微调时的学习率一定要小通常比初始训练小一到两个数量级。学习率太大会把预训练学到的权重直接破坏掉效果反而变差。我一般用 1e-5 起步如果 loss 震荡就再降。3.4 数据增强与正则化数据增强是防止过拟合的利器。对于猫狗分类常用的增强手段包括随机水平翻转、随机旋转、随机缩放、随机裁剪、亮度调整等。TensorFlow 2 提供了tf.keras.layers.RandomFlip、RandomRotation、RandomZoom等层可以直接嵌入模型里。data_augmentation tf.keras.Sequential([ tf.keras.layers.RandomFlip(horizontal), tf.keras.layers.RandomRotation(0.1), tf.keras.layers.RandomZoom(0.1), tf.keras.layers.RandomBrightness(0.1), ])这里有个细节数据增强层只在训练时生效推理时会自动跳过。所以你可以放心地把它作为模型的一部分导出时也不会影响推理结果。除了数据增强Dropout 和权重衰减也是常用的正则化手段。我在分类头里加了一个 Dropout(0.2)在微调阶段加了 L2 正则化验证集上的过拟合现象明显减轻。3.5 训练过程的监控与调参训练过程中要盯紧几个指标训练 loss、验证 loss、训练准确率、验证准确率。如果训练 loss 持续下降但验证 loss 开始上升那就是过拟合了需要加强正则化或提前停止。如果两个 loss 都不下降那可能是学习率太小或者模型容量不够。我习惯用EarlyStopping和ReduceLROnPlateau两个回调来自动化管理训练过程callbacks [ tf.keras.callbacks.EarlyStopping( monitorval_loss, patience5, restore_best_weightsTrue ), tf.keras.callbacks.ReduceLROnPlateau( monitorval_loss, factor0.5, patience3, min_lr1e-7 ), tf.keras.callbacks.ModelCheckpoint( best_model.h5, monitorval_accuracy, save_best_onlyTrue ) ]EarlyStopping的 patience 设 5 意味着验证 loss 连续 5 个 epoch 不下降就停止训练并恢复到最佳权重。ReduceLROnPlateau在验证 loss 停滞时自动降低学习率帮助模型跳出局部最优。这两个配合使用基本不用手动调学习率。4. 模型转换与优化从 Keras 到 TensorRT 的必经之路4.1 为什么不能直接拿 Keras 模型去部署Keras 模型.h5 或 SavedModel 格式包含的是完整的计算图和浮点权重依赖 TensorFlow 运行时来执行。在 Jetson Nano 上直接跑 TensorFlow 推理当然可以但性能很差。原因有几个TensorFlow 的运行时开销大很多算子在 GPU 上的实现没有针对 Maxwell 架构优化而且默认的浮点精度是 FP32对 Nano 的算力来说负担太重。所以我们需要把模型转换成更适合边缘推理的格式。在 Jetson 平台上最主流的选择是 TensorRT。TensorRT 是 NVIDIA 的推理优化引擎它会做几件事层融合把多个算子合并成一个、精度校准FP32 转 FP16 或 INT8、内核自动调优针对具体硬件选择最优的 CUDA 内核。经过 TensorRT 优化后推理速度通常能提升 2-5 倍。4.2 转换路径的选择ONNX 还是直接 TF-TRT从 TensorFlow 模型到 TensorRT有两条路可走。第一条路是 TensorFlow 模型 → ONNX → TensorRT。ONNX 是开放的模型交换格式tf2onnx 工具可以把 TensorFlow 模型转成 ONNX然后用 TensorRT 的 ONNX 解析器加载。这条路的好处是通用性强ONNX 生态成熟调试工具多。缺点是转换过程中可能遇到不支持的算子需要手动处理。第二条路是直接用 TF-TRTTensorFlow-TensorRT。TensorFlow 集成了 TensorRT 的支持可以在加载模型时自动把兼容的子图转成 TensorRT 引擎。这条路的好处是流程简单不需要中间格式。缺点是灵活性差有些自定义层可能不被支持而且 TF-TRT 的版本兼容性也是个坑。我个人的选择是走 ONNX 这条路。原因是 ONNX 的转换过程更可控出了问题容易定位。而且 ONNX 模型可以在 PC 上先用 onnxruntime 验证一遍确认输出和原始模型一致再放到 Jetson 上转 TensorRT减少在边缘设备上调试的时间。4.3 实操TensorFlow 模型转 ONNX先安装 tf2onnxpip install tf2onnx onnx onnxruntime然后转换python -m tf2onnx.convert \ --saved-model saved_model_dir \ --output model.onnx \ --opset 13这里有几个关键参数需要注意。--opset指定 ONNX 的算子集版本版本越高支持的算子越多但 TensorRT 对高版本 opset 的支持可能不完整。我一般用 opset 13兼容性比较好。--saved-model指定 SavedModel 目录如果你只有 .h5 文件需要先转成 SavedModel 格式。转换完成后用 onnxruntime 验证一下import onnxruntime as ort import numpy as np sess ort.InferenceSession(model.onnx) input_name sess.get_inputs()[0].name dummy_input np.random.randn(1, 224, 224, 3).astype(np.float32) output sess.run(None, {input_name: dummy_input}) print(output[0].shape) # 应该是 (1, 1)如果输出形状正确说明转换成功。如果报错通常是某个算子不被支持需要查 tf2onnx 的文档看有没有替代方案。4.4 在 Jetson Nano 上用 TensorRT 加速推理Jetson Nano 上已经预装了 TensorRT但版本可能比较老。先确认一下dpkg -l | grep tensorrt然后用trtexec工具把 ONNX 转成 TensorRT 引擎/usr/src/tensorrt/bin/trtexec \ --onnxmodel.onnx \ --saveEnginemodel.trt \ --fp16 \ --workspace1024 \ --verbose--fp16开启半精度推理Jetson Nano 的 GPU 对 FP16 有原生支持速度比 FP32 快不少精度损失通常在可接受范围内。--workspace指定工作空间大小单位 MBNano 的内存有限1024MB 是个比较稳妥的值。--verbose打印详细日志方便排查问题。转换成功后用 Python 加载 TensorRT 引擎做推理import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit import numpy as np TRT_LOGGER trt.Logger(trt.Logger.WARNING) def load_engine(engine_path): with open(engine_path, rb) as f, trt.Runtime(TRT_LOGGER) as runtime: return runtime.deserialize_cuda_engine(f.read()) def infer(engine, input_data): with engine.create_execution_context() as context: # 分配输入输出内存 input_shape input_data.shape output_shape (1, 1) d_input cuda.mem_alloc(input_data.nbytes) d_output cuda.mem_alloc(np.prod(output_shape) * 4) cuda.memcpy_htod(d_input, input_data) context.execute_v2([int(d_input), int(d_output)]) output np.empty(output_shape, dtypenp.float32) cuda.memcpy_dtoh(output, d_output) return output注意TensorRT 引擎是和硬件绑定的在 PC 上生成的引擎不能直接拿到 Jetson Nano 上用必须在目标设备上重新生成。这个坑我踩过白白浪费了一个下午。5. 部署实战在 Jetson Nano 上跑通完整推理流程5.1 Jetson Nano 的系统准备Jetson Nano 拿到手之后第一件事是烧录系统镜像。NVIDIA 官方提供了基于 Ubuntu 的 JetPack SDK包含了 CUDA、cuDNN、TensorRT 等所有必要的组件。烧录需要用 balenaEtcher 之类的工具把镜像写到 microSD 卡上然后插到 Nano 上启动。启动之后有几件事必须做第一换源。默认的软件源在国内访问很慢换成国内镜像源能大幅提升 apt 安装的速度。第二扩大 swap。Nano 只有 4GB 内存训练肯定不够但推理时如果模型稍大也可能爆内存。创建一个 4GB 的 swap 文件能缓解这个问题sudo fallocate -l 4G /swapfile sudo chmod 600 /swapfile sudo mkswap /swapfile sudo swapon /swapfile # 永久生效 echo /swapfile none swap sw 0 0 | sudo tee -a /etc/fstab第三设置电源模式。Nano 默认是 5W 模式GPU 频率被限制。如果散热条件允许可以切换到 10W 模式sudo nvpmodel -m 0 # 10W 模式 sudo jetson_clocks # 锁定最高频率注意10W 模式发热量明显增大一定要加散热片或者小风扇。我试过不加散热跑 10W 模式不到十分钟就过热降频了反而比 5W 模式还慢。5.2 推理脚本的编写与优化部署到 Nano 上的推理脚本需要处理几个问题图像预处理、模型推理、结果后处理。图像预处理要和训练时保持一致。训练时用的是 224x224 的输入归一化到 [0, 1] 或者 [-1, 1]取决于模型。MobileNetV2 的预处理是归一化到 [-1, 1]这个细节如果搞错了推理结果会完全不对。def preprocess_image(image_path): img tf.io.read_file(image_path) img tf.image.decode_jpeg(img, channels3) img tf.image.resize(img, (224, 224)) img tf.cast(img, tf.float32) / 127.5 - 1.0 # [-1, 1] return tf.expand_dims(img, axis0)推理部分用 TensorRT 引擎执行前面已经写过。后处理就是把 sigmoid 输出转成类别标签def postprocess(output): prob output[0][0] label dog if prob 0.5 else cat confidence prob if prob 0.5 else 1 - prob return label, confidence5.3 性能测试与瓶颈分析部署完成之后一定要做性能测试。我通常关注三个指标单帧推理延迟、吞吐量FPS、内存占用。在 Jetson Nano 10W 模式下MobileNetV2 的 FP16 TensorRT 引擎单帧推理延迟大约在 15-20ms也就是 50-60 FPS。这个速度对于大多数实时应用来说已经足够了。如果换成 FP32延迟会增加到 30-40ms差不多慢一倍。内存占用方面TensorRT 引擎加载后大约占用 200-300MB 内存加上系统和 Python 运行时总共 1GB 左右。Nano 的 4GB 内存完全够用。如果性能不达标可以从几个方向排查是不是用了 FP32 而不是 FP16是不是没有开 jetson_clocks是不是输入尺寸太大了是不是模型本身太重了比如用了 ResNet50每一个因素都可能成为瓶颈。5.4 常见部署问题与解决方案在实际部署过程中我遇到过不少问题这里整理成一张速查表问题现象可能原因解决方案推理结果全是同一类预处理不一致检查归一化参数和输入尺寸TensorRT 引擎加载失败引擎与硬件不匹配在目标设备上重新生成引擎推理速度远低于预期未开启 FP16 或未锁频加 --fp16运行 jetson_clocks内存不足报错模型太大或 swap 不够换更小的模型扩大 swapONNX 转换报错算子不支持查 tf2onnx 文档替换或自定义算子输出形状不对输入维度不匹配确认 batch size 和输入尺寸还有一个容易被忽略的问题TensorRT 引擎的序列化和反序列化。引擎文件可能比较大几十 MB加载需要时间。如果每次推理都重新加载引擎启动延迟会很高。正确的做法是在程序启动时加载一次然后复用。6. 从项目到产品还能怎么扩展6.1 多分类与目标检测的延伸猫狗二分类只是一个起点。把这个框架扩展到多分类比如识别不同品种的猫狗或者目标检测不仅识别是什么还要框出位置思路是一样的。多分类只需要把最后的 Dense 层从 1 个输出改成 N 个输出激活函数从 sigmoid 改成 softmax损失函数从 binary_crossentropy 改成 categorical_crossentropy。目标检测则需要换用 SSD、YOLO 等检测架构但部署流程——训练、转 ONNX、转 TensorRT、在 Nano 上推理——是完全一致的。6.2 视频流实时推理静态图片推理跑通之后下一步自然是接摄像头做视频流推理。Jetson Nano 支持 CSI 摄像头和 USB 摄像头用 OpenCV 读取视频帧逐帧送入 TensorRT 引擎推理再把结果画在帧上显示出来。这里的关键是处理好帧率和延迟的平衡。如果推理速度跟不上摄像头帧率可以跳帧处理或者降低输入分辨率。6.3 模型量化INT8 的诱惑与陷阱FP16 之后还能进一步量化到 INT8。INT8 的理论速度是 FP16 的两倍模型大小也减半。但 INT8 需要校准数据集来确定量化参数而且精度损失可能比较明显。我试过在猫狗分类上做 INT8 量化准确率从 97% 掉到了 94% 左右速度提升了大约 60%。如果你的应用对精度要求不那么苛刻INT8 是值得尝试的。但如果精度是硬指标那就老老实实用 FP16。6.4 实际部署中的工程化考量从 demo 到产品还有不少工程化的工作要做。比如如何做模型的版本管理和热更新如何监控推理服务的健康状态如何处理异常输入损坏的图片、非图像文件如何做批量推理提升吞吐这些问题在 demo 阶段可以忽略但在实际部署中每一个都需要认真对待。我个人的体会是边缘 AI 部署这件事训练模型只是冰山一角水面下的工程化工作才是真正决定项目成败的关键。把猫狗分类部署到 Jetson Nano 上跑通你收获的不仅仅是一个能识别猫狗的模型而是一整套可以复用到其他项目上的部署方法论。这套方法论包括如何选择适合边缘设备的模型架构、如何做模型格式转换、如何优化推理性能、如何排查部署问题。这些东西才是这个项目最大的价值所在。最后分享一个小技巧在 Jetson Nano 上调试 TensorRT 的时候善用trtexec的--dumpProfile参数它会打印每一层的耗时帮你快速定位性能瓶颈。很多时候你以为瓶颈在卷积层结果发现是某个 reshape 或者 transpose 操作拖慢了整体速度。这种细节只有实际跑过一遍才能体会到。