上位机知识篇---Jetson Nano的深度学习GPU推理

Jetson Nano 是 NVIDIA 推出的低成本边缘计算开发板（类似小型电脑），主打轻量化深度学习部署，适合学生、开发者做边缘 AI 项目（如摄像头实时识别、小型机器人视觉等）。虽然它的 GPU 算力有限（128 核 Maxwell 架构，算力 0.5 TFLOPS），比不上台式机 GPU，但足够跑通 "环境配置→模型训练→部署" 的全流程。下面用通俗易懂的语言介绍每个步骤。

一、准备工作：硬件和系统

1. 必备硬件

Jetson Nano 开发板：核心设备，自带 GPU。
电源：必须用 5V 4A 直流电源（Micro USB 接口容易供电不足，建议用 DC 接口），否则训练或推理时可能突然断电。
SD 卡：至少 32GB（推荐 64GB），用来装系统。
散热风扇：Jetson Nano 运行时会发热，长时间训练需要散热，否则会降频变慢。
摄像头：USB 摄像头或 Jetson 专用摄像头（如 Raspberry Pi Camera），用于部署时实时采集图像。
显示器、键盘、鼠标：初期配置用，后期可通过 SSH 远程操作。

2. 刷写系统（JetPack）

Jetson Nano 需要专用系统 JetPack，它是 NVIDIA 为 Jetson 系列定制的系统包，自带：

操作系统（基于 Ubuntu）；
GPU 驱动、CUDA（用于 GPU 计算）、cuDNN（加速深度学习）等底层工具，不用手动装。

步骤：

从 NVIDIA 官网下载 JetPack 镜像（比如最新的 JetPack 4.6，支持大部分深度学习框架）；
用软件（如 Etcher）将镜像写入 SD 卡；
把 SD 卡插入 Jetson Nano，通电开机，按提示完成系统初始化（设用户名、密码等）。

二、环境配置：安装深度学习工具

Jetson Nano 是 ARM 架构（和普通电脑的 x86 架构不同），不能直接装普通的 PyTorch 或 TensorFlow，需要用 NVIDIA 预编译的版本。这里以最常用的 PyTorch 为例。

1. 换源（加速下载）

默认软件源在国外，下载慢，先换成国内源（比如清华源）：

bash 复制代码

# 备份原文件
sudo cp /etc/apt/sources.list /etc/apt/sources.list.bak
# 编辑源文件，替换为清华源（适合 JetPack 对应的 Ubuntu 版本，比如 18.04）
sudo nano /etc/apt/sources.list
# 替换后更新
sudo apt update && sudo apt upgrade -y

2. 安装 PyTorch

Jetson Nano 只能用 NVIDIA 提供的 ARM 版本 PyTorch，步骤：

查看 JetPack 版本（系统自带，可在 "设置→详细信息" 里看，比如 JetPack 4.6 对应 Ubuntu 18.04）；
从 NVIDIA 官网找对应版本的 PyTorch 安装包（比如 JetPack 4.6 对应 PyTorch 1.10.0），下载 .whl 文件；
安装依赖和 PyTorch：

bash 复制代码

# 安装依赖
sudo apt install -y python3-pip libopenblas-base libopenmpi-dev
# 安装 PyTorch（文件名替换为你下载的）
pip3 install torch-1.10.0-cp36-cp36m-linux_aarch64.whl

3. 安装配套库

torchvision ：PyTorch 的视觉工具库，需要和 PyTorch 版本匹配，同样用预编译包：
bash 复制代码
```
pip3 install torchvision-0.11.1-cp36-cp36m-linux_aarch64.whl
```
OpenCV ：处理图像，JetPack 通常预装了，若没有：
bash 复制代码
```
sudo apt install -y python3-opencv
```
其他工具 ：数据处理用的 numpy，可视化用的 matplotlib 等：
bash 复制代码
```
pip3 install numpy matplotlib pandas
```

4. 验证环境

打开 Python 终端，检查是否能正常调用 GPU：

python 复制代码

import torch
# 查看是否支持 CUDA（返回 True 则成功）
print(torch.cuda.is_available())  # 输出 True
# 查看 GPU 信息（显示 Jetson Nano 的 GPU）
print(torch.cuda.get_device_name(0))  # 输出 "NVIDIA Tegra X1"

三、数据准备：适合 Jetson Nano 的 "轻量化" 数据

Jetson Nano 内存（4GB）和存储有限，不能用太大的数据集（比如 ImageNet 有 1400 万张图，根本放不下）。建议：

用小数据集：比如自己拍 1000 张图片（如 "猫、狗、鸟" 三类），或用公开的小数据集（如 Kaggle 上的宠物识别小数据集）；
预处理数据：缩小图像尺寸（比如从 1024x1024 缩到 224x224），减少存储和计算量；
本地存储：把数据存在 SD 卡或外接 U 盘（别存在系统盘，避免占满空间）。

四、模型训练：选对模型，降低难度

Jetson Nano 的 GPU 算力弱，不适合训练大型模型（如 ResNet101、YOLOv5 完整版），否则可能几天都训练不完，甚至内存溢出。建议用以下策略：

1. 选轻量级模型

优先用为边缘设备设计的小模型：

分类任务：MobileNetV2、EfficientNet-Lite0/1（参数少，计算快）；
检测任务：YOLOv5s（YOLOv5 的 smallest 版本）、MobileNet-SSD。

2. 用迁移学习（核心技巧）

不要 "从头训练"，而是用预训练模型微调（站在巨人的肩膀上）：

原理：预训练模型（如 MobileNetV2）已经在大数据集上学过通用特征（边缘、纹理），我们只需用少量数据训练最后几层，适配自己的任务（如识别 "猫、狗"）；
好处：训练快（几小时就能完成），数据量少也能出效果。

3. 训练时的优化技巧

减小 batch size：一次喂给模型的图片数量（比如设为 4 或 8，普通 GPU 可能用 32），避免内存不足；
冻结大部分层：只训练最后 1-2 层（比如冻结 MobileNetV2 的前 90% 层），减少计算量；
降低输入尺寸：比如把图像从 224x224 降到 160x160（精度会降一点，但速度快很多）；
监控资源 ：用 nvidia-smi 命令查看 GPU 占用，若超过 90% 可能要调低参数。

训练代码示例（简单分类任务）

假设用 MobileNetV2 微调，识别 "猫、狗、鸟" 三类：

python 复制代码

import torch
from torchvision import models, transforms
from torch.utils.data import DataLoader, Dataset
from PIL import Image
import os

# 1. 定义数据集（自己的图片存在 "data/train" 下，每个类别一个文件夹）
class CustomDataset(Dataset):
    def __init__(self, root, transform):
        self.root = root
        self.transform = transform
        self.classes = os.listdir(root)  # 类别：cat, dog, bird
        self.images = []
        for label, cls in enumerate(self.classes):
            for img_name in os.listdir(f"{root}/{cls}"):
                self.images.append((f"{root}/{cls}/{img_name}", label))
    
    def __len__(self):
        return len(self.images)
    
    def __getitem__(self, idx):
        img_path, label = self.images[idx]
        img = Image.open(img_path).convert("RGB")
        return self.transform(img), label

# 2. 数据预处理（尺寸调小，适合 Nano）
transform = transforms.Compose([
    transforms.Resize((160, 160)),  # 输入尺寸从 224 降到 160
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 3. 加载数据
train_dataset = CustomDataset("data/train", transform)
train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)  # batch size=4

# 4. 加载预训练模型，修改输出层（3类）
model = models.mobilenet_v2(pretrained=True)
# 冻结大部分层（只训练最后一层）
for param in model.features.parameters():
    param.requires_grad = False
# 修改输出层（原来输出 1000 类，现在改 3 类）
model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, 3)
model = model.cuda()  # 放到 GPU 上

# 5. 训练
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.classifier.parameters(), lr=0.001)  # 只优化最后一层

model.train()
for epoch in range(10):  # 训练 10 轮
    total_loss = 0
    for images, labels in train_loader:
        images, labels = images.cuda(), labels.cuda()
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"第 {epoch+1} 轮，损失：{total_loss/len(train_loader):.4f}")

# 保存模型
torch.save(model.state_dict(), "nano_model.pth")

五、模型部署：让模型在 Jetson Nano 上 "实时工作"

训练好的模型需要部署到 Jetson Nano 上，比如用摄像头实时识别物体。核心是优化推理速度（让模型跑得更快）。

1. 模型优化：用 TensorRT 加速

TensorRT 是 NVIDIA 的推理加速工具，能把模型 "压缩优化"，在 Jetson 上速度可提升 2-5 倍。步骤：

把 PyTorch 模型转为 ONNX 格式（中间格式）；
用 TensorRT 把 ONNX 转为优化后的引擎（.engine 文件）。

示例代码（转换模型）：

python 复制代码

# 1. 加载训练好的模型
model = models.mobilenet_v2(pretrained=False)
model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, 3)
model.load_state_dict(torch.load("nano_model.pth"))
model.eval()
model = model.cuda()

# 2. 转为 ONNX 格式
dummy_input = torch.randn(1, 3, 160, 160).cuda()  # 输入示例
torch.onnx.export(
    model, dummy_input, "model.onnx",
    input_names=["input"], output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

# 3. 用 TensorRT 优化（需要安装 tensorrt）
import tensorrt as trt
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, TRT_LOGGER)
with open("model.onnx", "rb") as f:
    parser.parse(f.read())
# 构建优化引擎
config = builder.create_builder_config()
engine = builder.build_serialized_network(network, config)
with open("model.engine", "wb") as f:
    f.write(engine)

2. 实时推理：用摄像头测试

用优化后的模型，结合 OpenCV 读取摄像头画面，实时预测：

python 复制代码

import cv2
import torch
import numpy as np

# 加载优化后的模型（这里简化为用 PyTorch 模型，实际用 TensorRT 更快）
model = models.mobilenet_v2(pretrained=False)
model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, 3)
model.load_state_dict(torch.load("nano_model.pth"))
model.eval()
model = model.cuda()

# 预处理函数
def preprocess(img):
    img = cv2.resize(img, (160, 160))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # OpenCV 是 BGR，转为 RGB
    img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
    img = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img)
    return img.unsqueeze(0).cuda()

# 打开摄像头
cap = cv2.VideoCapture(0)  # 0 是默认摄像头
classes = ["cat", "dog", "bird"]

while True:
    ret, frame = cap.read()  # 读取一帧画面
    if not ret:
        break
    
    # 预处理并预测
    input_tensor = preprocess(frame)
    with torch.no_grad():
        output = model(input_tensor)
        pred = torch.argmax(output).item()
        label = classes[pred]
    
    # 在画面上显示结果
    cv2.putText(frame, label, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    cv2.imshow("Jetson Nano Detection", frame)
    
    # 按 q 退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

六、常见问题和注意事项

供电不足：训练或推理时突然关机，一定换 5V 4A 电源；
内存溢出：训练时报错 "out of memory"，减小 batch size 或输入尺寸；
速度慢：推理帧率低（如每秒 5 帧），用 TensorRT 优化，或换更小的模型；
散热：摸起来很烫，加装散热风扇，否则会降频（性能变低）。

总结（通俗易懂版）

用 Jetson Nano 做深度学习，就像 "用小型电动车跑长途"------ 虽然动力不如大卡车（台式机 GPU），但规划好路线（选轻量模型、用迁移学习）、做好保养（优化参数、散热），照样能跑完全程。流程就是：先装好系统和工具，用小数据和轻量模型训练，最后优化模型让它在摄像头前 "实时干活"。适合做一些有趣的小项目，比如 "智能垃圾桶（识别垃圾类型）""宠物监控（识别猫还是狗）" 等。