到这里已经有了一个 PyTorch YOLOv3-Tiny 模型，并且通过 QAT 进行了训练和优化，并考虑了锚框聚类。现在，是时候将其部署到目标嵌入式设备上，通常这将涉及将其转换为 TensorFlow Lite (TFLite) 格式。
下面是详细的 TFLite 部署步骤，涵盖了从 PyTorch 到 TFLite 的完整流程：
TFLite 部署
- 准备 PyTorch 模型进行导出: 确保模型已完成 QAT 并转换为量化模式。
- PyTorch 模型导出为 ONNX: 将 PyTorch 模型转换为 ONNX 格式。
- ONNX 模型转换为 TensorFlow SavedModel: 使用 ONNX-TensorFlow 工具将 ONNX 模型转换为 TensorFlow 的 SavedModel 格式。
- TensorFlow SavedModel 转换为 TFLite: 使用 TensorFlow Lite Converter 将 SavedModel 转换为最终的
.tflite模型。 - 在嵌入式设备上运行 TFLite 模型: 使用 TFLite 运行时库进行推理。
详细步骤
步骤 1: 准备 PyTorch 模型进行导出
在你的 train.py 脚本中，当 Config.QUANT_MODE 为 True 且训练结束后，你会得到一个量化后的 PyTorch 模型。
# train.py (摘录，训练循环结束之后)
# --- 训练完成后，如果进行了 QAT，需要转换为量化模型 ---
if Config.QUANT_MODE:
print("--- Converting QAT model to quantized model ---")
model.eval()
# 将 QAT 模型转换为实际的量化模型 (INT8)
# 注意：这里的转换是针对 CPU 后端的
quantized_model = torch.quantization.convert(model, inplace=False)
# 保存量化模型状态字典，方便后续加载和导出
torch.save(quantized_model.state_dict(), "yolov3_tiny_person_quantized_state_dict.pth")
print("Quantized model state dict saved to yolov3_tiny_person_quantized_state_dict.pth")
# 这里需要创建一个新的 PyTorch 模型实例，并加载量化后的状态字典
# 这是为了确保模型结构是量化兼容的，并且权重是量化后的
# 重新初始化模型，并加载量化后的 state_dict
# 你的 YOLOv3Tiny 模型在定义时已经考虑了量化兼容性（例如 ConvBlock 中的 bias=not use_bn）
# 但是，torch.quantization.convert 会将量化操作插入到模型中，
# 所以最直接的方式是保存 convert 之后的整个模型 (如果可以序列化) 或其 state_dict
# 更好的是直接使用转换后的 quantized_model 变量进行下一步的 ONNX 导出
# 确保这个 `quantized_model` 是 `torch.quantization.convert` 的输出
print("Proceeding to ONNX export with the converted quantized model...")
# 后续的 ONNX 导出将使用这个 `quantized_model` 对象
else:
print("Quantization mode is OFF. Exporting unquantized model.")
quantized_model = model # 如果没有QAT，则使用原始模型
关键点: 确保你导出的 quantized_model 确实是经过 torch.quantization.convert() 处理后的模型实例，因为这个实例包含了量化操作和校准后的统计信息。
步骤 2: PyTorch 模型导出为 ONNX
这是一个通用的模型交换格式，TFLite Converter 可以通过 ONNX 间接处理 PyTorch 模型。
安装 ONNX 相关的库:
pip install onnx onnxruntime
Python 代码 (例如，可以单独写一个 export_to_onnx.py 脚本):
# export_to_onnx.py
import torch
from config import Config
from model import YOLOv3Tiny # 导入你的模型定义
def export_model_to_onnx(model_path, onnx_output_path, img_size, quant_mode=False):
"""
加载 PyTorch 模型并导出为 ONNX 格式。
model_path: 保存的 PyTorch 模型 state_dict 路径 (.pth)
onnx_output_path: 导出的 ONNX 文件路径 (.onnx)
img_size: 模型输入图像尺寸 (例如 416)
quant_mode: 是否是量化模型
"""
# 1. 实例化模型
model = YOLOv3Tiny(num_classes=Config.NUM_CLASSES, config_anchors=Config.ANCHORS)
model.to(Config.DEVICE) # 确保模型在正确的设备上
# 2. 加载模型权重
if quant_mode:
# 如果是量化模型，需要先准备QAT，然后加载 state_dict
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm') # 或 'qnnpack'
# 融合模块 (必须与训练时一致)
model_fused = torch.quantization.fuse_modules(model, [['features.0.conv', 'features.0.bn', 'features.0.activation'],
['features.2.conv', 'features.2.bn', 'features.2.activation'],
['features.4.conv', 'features.4.bn', 'features.4.activation'],
['features.6.conv', 'features.6.bn', 'features.6.activation'],
['features.8.conv', 'features.8.bn', 'features.8.activation'],
['features.10.conv', 'features.10.bn', 'features.10.activation'],
['features.12.conv', 'features.12.bn', 'features.12.activation'],
['features.13.conv', 'features.13.bn', 'features.13.activation'],
['features.14.conv', 'features.14.bn', 'features.14.activation'],
['head2_conv1.conv', 'head2_conv1.bn', 'head2_conv1.activation'],
['head2_conv2.conv', 'head2_conv2.bn', 'head2_conv2.activation'],
])
# 准备QAT
model_prepared = torch.quantization.prepare_qat(model_fused, inplace=True)
# 加载量化训练后的 state_dict
model_prepared.load_state_dict(torch.load(model_path, map_location='cpu')) # QAT后的state_dict可能需要CPU加载
# 转换为量化模式
model = torch.quantization.convert(model_prepared, inplace=True)
print("Loaded and converted quantized PyTorch model.")
else:
# 非量化模型直接加载 state_dict
model.load_state_dict(torch.load(model_path, map_location='cpu'))
print("Loaded unquantized PyTorch model.")
model.eval() # 切换到评估模式
# 3. 创建一个虚拟输入张量
# 批次大小通常为 1 进行部署
dummy_input = torch.randn(1, 3, img_size, img_size).to(Config.DEVICE)
# 4. 导出模型到 ONNX
print(f"Exporting model to ONNX: {onnx_output_path}")
torch.onnx.export(model,
dummy_input,
onnx_output_path,
verbose=False,
opset_version=13, # 推荐使用 11 或更高版本，确保兼容性
do_constant_folding=True,
input_names=['input'],
output_names=['output'],
dynamic_axes={'input': {0: 'batch_size'}, # 如果希望支持动态批次大小
'output': {0: 'batch_size'}})
print("ONNX model export complete!")
if __name__ == "__main__":
# 使用你训练好的模型路径，如果进行了QAT，就是QAT保存的 state_dict
pytorch_model_path = "yolov3_tiny_person_quantized_state_dict.pth"
onnx_output_file = "yolov3_tiny_person_quantized.onnx"
export_model_to_onnx(pytorch_model_path, onnx_output_file, Config.IMAGE_SIZE, quant_mode=Config.QUANT_MODE)
运行这个脚本: python export_to_onnx.py
步骤 3: ONNX 模型转换为 TensorFlow SavedModel
TensorFlow Lite Converter 可以直接从 ONNX 导入模型（需要 onnx-tf 库）。这会将其转换为 TensorFlow SavedModel 格式，这是 TFLite Converter 的首选输入格式。
安装 ONNX-TensorFlow:
pip install onnx-tf tensorflow==2.x # 确保安装兼容的TensorFlow版本
Python 代码 (例如，可以单独写一个 convert_onnx_to_tf.py 脚本):
# convert_onnx_to_tf.py
import onnx
from onnx_tf.backend import prepare
import tensorflow as tf
import os
def convert_onnx_to_tf_savedmodel(onnx_path, saved_model_path):
"""
将 ONNX 模型转换为 TensorFlow SavedModel 格式。
onnx_path: 输入 ONNX 文件路径 (.onnx)
saved_model_path: 输出 SavedModel 目录路径
"""
print(f"Loading ONNX model from: {onnx_path}")
onnx_model = onnx.load(onnx_path)
print("Converting ONNX model to TensorFlow SavedModel...")
# prepare 函数会创建一个 TensorFlow Backend 对象
tf_rep = prepare(onnx_model)
# 导出为 SavedModel
tf_rep.export_graph(saved_model_path)
print(f"TensorFlow SavedModel exported to: {saved_model_path}")
if __name__ == "__main__":
onnx_input_file = "yolov3_tiny_person_quantized.onnx"
tf_saved_model_dir = "yolov3_tiny_person_tf_savedmodel"
# 创建输出目录
os.makedirs(tf_saved_model_dir, exist_ok=True)
convert_onnx_to_tf_savedmodel(onnx_input_file, tf_saved_model_dir)
运行这个脚本: python convert_onnx_to_tf.py
步骤 4: TensorFlow SavedModel 转换为 TFLite
现在，我们有了 TensorFlow SavedModel，可以将其转换为最终的 .tflite 格式。
Python 代码 (例如，可以单独写一个 convert_tf_to_tflite.py 脚本):
# convert_tf_to_tflite.py
import tensorflow as tf
from config import Config
import numpy as np # 用于代表性数据集
def convert_tf_to_tflite(saved_model_path, tflite_output_path, quant_mode=False):
"""
将 TensorFlow SavedModel 转换为 TFLite 格式。
saved_model_path: 输入 SavedModel 目录路径
tflite_output_path: 输出 TFLite 文件路径 (.tflite)
quant_mode: 是否进行整数后量化 (如果模型已经是 QAT，则为 True)
"""
print(f"Loading TensorFlow SavedModel from: {saved_model_path}")
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_path)
if quant_mode:
print("Enabling default optimizations for quantization (INT8)...")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# 确保 TFLite Runtime 支持 INT8 操作
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
# 如果有自定义操作或 ONNX 转换引入了非标准 TF op，可能需要添加 SELECT_TF_OPS
# converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8, tf.lite.OpsSet.SELECT_TF_OPS]
# 如果你在 ONNX -> TF 转换过程中失去了 QAT 的量化信息
# 并且模型最终没有被正确识别为量化模型，你可能需要进行后训练整数量化 (Post-Training Integer Quantization)
# 这需要提供一个代表性数据集来校准量化范围
# def representative_data_gen():
# # 从你的训练数据集中加载少量样本（例如 100-500 张）
# # 确保数据预处理与训练时一致
# # 假设你的数据集在 Config.TRAIN_IMG_DIR
# from dataset import YOLODataset # 导入数据集
# import os
#
# # 创建一个临时的、只用于校准的数据加载器
# calibration_dataset = YOLODataset(
# Config.TRAIN_IMG_DIR,
# Config.TRAIN_LABEL_DIR,
# Config.ANCHORS,
# img_size=Config.IMAGE_SIZE,
# num_classes=Config.NUM_CLASSES,
# is_train=False, # 不进行数据增强
# )
# calibration_loader = torch.utils.data.DataLoader(
# dataset=calibration_dataset,
# batch_size=1, # 必须是1
# shuffle=False,
# num_workers=0,
# collate_fn=custom_collate_fn,
# )
#
# num_calibration_steps = Config.NUM_CALIBRATION_BATCHES # 和 QAT 校准批次数量保持一致
# print(f"Generating representative dataset for TFLite quantization ({num_calibration_steps} batches)...")
# for i, (image, _) in enumerate(calibration_loader):
# if i >= num_calibration_steps:
# break
# # 将 PyTorch tensor 转换为 NumPy 数组，并确保数据类型和范围正确
# # TFLite 通常期望 float32 输入，尽管它是量化模型
# yield [image.cpu().numpy()]
# converter.representative_dataset = representative_data_gen
# 如果是 QAT 模型，通常不需要上述的 `representative_dataset` 再次校准
# 因为 QAT 已经在训练时完成了这个过程。
# ONNX -> TF SavedModel 转换后，如果量化信息能保留下来，则 TFLite Converter 会识别。
# 如果没有保留，它会尝试进行后训练量化，这时 `representative_dataset` 就很重要了。
# 最佳实践是：QAT -> ONNX -> TF SavedModel -> TFLite，并验证 TFLite 模型的量化状态。
else:
print("Converting to TFLite without explicit quantization (float32).")
# 如果没有进行 QAT，这里会生成一个 float32 的 TFLite 模型。
# 可以在此进行 Post-Training Dynamic Range Quantization 或 Full Integer Quantization (PTIQ)
# converter.optimizations = [tf.lite.Optimize.DEFAULT] # 默认包含动态范围量化
# converter.target_spec.supported_types = [tf.float16] # 也可以转为FP16
print("Converting to TFLite...")
tflite_model = converter.convert()
with open(tflite_output_path, "wb") as f:
f.write(tflite_model)
print(f"TFLite model saved to: {tflite_output_path}")
if __name__ == "__main__":
tf_saved_model_input_dir = "yolov3_tiny_person_tf_savedmodel"
tflite_output_file = "yolov3_tiny_person_detector_quantized.tflite"
convert_tf_to_tflite(tf_saved_model_input_dir, tflite_output_file, quant_mode=Config.QUANT_MODE)
# 验证 TFLite 模型大小
tflite_model_size_mb = os.path.getsize(tflite_output_file) / (1024 * 1024)
print(f"TFLite model size: {tflite_model_size_mb:.2f} MB")
运行这个脚本: python convert_tf_to_tflite.py
步骤 5: 在嵌入式设备上运行 TFLite 模型
一旦你有了 .tflite 文件，就可以将其部署到你的 ARM Cortex-A 处理器上，并使用 TFLite 运行时库进行推理。
C++ 部署示例 (概念性代码):
// inference_on_device.cpp (伪代码)
#include <iostream>
#include <vector>
#include <string>
#include <memory> // For std::unique_ptr
// 假设已经包含了 TensorFlow Lite C++ API 的头文件
#include "tensorflow/lite/interpreter.h"
#include "tensorflow/lite/kernels/register.h"
#include "tensorflow/lite/model.h"
#include "tensorflow/lite/tools/command_line_helpers.h" // 可能有用
// 图像处理库 (例如 OpenCV)
// #include <opencv2/opencv.hpp>
// 你的模型输入尺寸
const int INPUT_IMG_SIZE = 416;
const int INPUT_CHANNELS = 3;
// 加载 TFLite 模型
std::unique_ptr<tflite::FlatBufferModel> model =
tflite::FlatBufferModel::BuildFromFile("yolov3_tiny_person_detector_quantized.tflite");
if (!model) {
std::cerr << "Failed to load model" << std::endl;
return 1;
}
// 构建解释器 (Interpreter)
tflite::ops::builtin::BuiltinOpResolver resolver;
std::unique_ptr<tflite::Interpreter> interpreter;
tflite::InterpreterBuilder(*model, resolver)(&interpreter);
if (!interpreter) {
std::cerr << "Failed to construct interpreter" << std::endl;
return 1;
}
// 分配张量 (Allocate Tensors)
// 这会为模型中的所有张量分配内存
if (interpreter->AllocateTensors() != kTfLiteOk) {
std::cerr << "Failed to allocate tensors" << std::endl;
return 1;
}
// 确保输入张量和模型期望的类型和形状匹配
// 对于量化模型，输入张量可能期望 int8 或 uint8 类型
// 在推理前需要对输入图像进行预处理，使其符合模型输入要求
TfLiteTensor* input_tensor = interpreter->input_tensor(0);
// 获取输入张量的类型和维度
// 例如，如果 input_tensor->type == kTfLiteUInt8, 则输入图像需要转换为 uint8 格式
// input_tensor->dims[0] 是 batch_size, dims[1] 是 height, dims[2] 是 width, dims[3] 是 channels
// ---------- 推理循环 (智能门铃的核心逻辑) ----------
while (true) {
// 1. 捕获图像 (例如从摄像头)
// cv::Mat frame = camera.read(); // 假设有摄像头接口
// if (frame.empty()) continue;
// 2. 图像预处理
// 调整大小: cv::resize(frame, resized_frame, cv::Size(INPUT_IMG_SIZE, INPUT_IMG_SIZE));
// 颜色空间转换: cv::cvtColor(resized_frame, resized_frame, cv::COLOR_BGR2RGB);
// 对于量化模型：
// 如果模型输入是 uint8，你需要将像素值缩放到 [0, 255] 范围
// 如果模型输入是 float32，则需要进行归一化 (除以 255.0) 和 ImageNet 标准化
// (例如：(pixel_value / 255.0 - mean) / std)
// 这里以 uint8 输入为例 (常见于边缘量化模型)
// 伪代码: 将图像数据复制到 input_tensor
// for (int i = 0; i < INPUT_IMG_SIZE * INPUT_IMG_SIZE * INPUT_CHANNELS; ++i) {
// input_tensor->data.uint8[i] = preprocessed_image_data[i];
// }
// 模拟输入数据 (实际应用中替换为摄像头数据)
std::vector<uint8_t> dummy_input_data(INPUT_IMG_SIZE * INPUT_IMG_SIZE * INPUT_CHANNELS);
// 填充一些模拟数据
// ...
memcpy(input_tensor->data.uint8, dummy_input_data.data(), dummy_input_data.size());
// 3. 运行推理
// 可以测量推理时间
// auto start_time = std::chrono::high_resolution_clock::now();
if (interpreter->Invoke() != kTfLiteOk) {
std::cerr << "Failed to invoke interpreter" << std::endl;
break;
}
// auto end_time = std::chrono::high_resolution_clock::now();
// auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count();
// std::cout << "Inference latency: " << duration << " ms" << std::endl;
// 4. 获取输出
TfLiteTensor* output_tensor = interpreter->output_tensor(0);
// 输出张量的形状和类型取决于你的模型 (YOLOv3-Tiny 的输出通常是 (1, total_boxes, 5+num_classes))
// 对于量化模型，输出张量可能是 int8 或 float32。
// 如果是 int8，你需要进行反量化 (dequantization) 才能得到实际的浮点结果：
// float real_value = (output_tensor->data.int8[i] - output_tensor->params.zero_point) * output_tensor->params.scale;
// 5. 后处理 (NMS, 阈值过滤)
// 遍历 output_tensor 中的预测框
// 根据置信度阈值和 NMS 过滤结果
// 识别出“人”的边界框
// 伪代码: 解析输出并执行NMS
// auto detections = parse_tflite_output(output_tensor, original_w, original_h);
// auto nms_results = apply_nms(detections, Config.NMS_IOU_THRESH, Config.CONF_THRESHOLD);
// 6. 决策和触发警报
// if (person_detected_in_nms_results) {
// std::cout << "Person detected!" << std::endl;
// // 触发门铃警报，发送通知等
// }
// 7. 功耗优化 (跳帧检测等)
// std::this_thread::sleep_for(std::chrono::milliseconds(FRAME_INTERVAL_MS)); // 控制处理帧率
// 或者根据运动检测器触发
}
在嵌入式设备上编译和运行：
- 获取 TensorFlow Lite C++ 库: 你需要在你的嵌入式 Linux 系统上交叉编译 TensorFlow Lite 运行时库。通常，你可以从 TensorFlow 的 GitHub 仓库下载源码，然后按照其
tensorflow/lite/tools/build_lib.sh脚本进行交叉编译。- 交叉编译: 这意味着在你的开发主机（例如 x86 Ubuntu）上为目标设备（例如 ARM 处理器）编译库。
- 取决于你的设备: 不同的 ARM 处理器可能有不同的编译标志（例如针对 NEON 指令集）。
- 集成: 将编译好的 TFLite 库以及你的推理代码集成到你的嵌入式项目（例如使用 CMake 或 Makefile）中。
- 图像输入/输出: 确保你的摄像头驱动和图像处理管道能够高效地将图像数据送入 TFLite 模型，并处理模型输出。
注意事项和优化
- 测试与验证: 在真实设备上对转换后的 TFLite 模型进行全面的测试。验证其精度是否符合预期（与 PyTorch 模型的精度损失是否在可接受范围内），推理延迟和功耗是否满足要求。
- 输入预处理: TFLite 模型的输入预处理非常关键。确保图像缩放、归一化（如果是非量化模型）或量化（如果是 INT8 模型）的方式与训练时完全一致。
- 输出后处理: 同样，模型的输出需要进行适当的后处理，包括将边界框坐标从模型输出的相对值转换为图像的绝对像素值，以及执行 NMS。
- 硬件加速 (NPU/DSP): 如果你的嵌入式芯片有 NPU (Neural Processing Unit) 或 DSP (Digital Signal Processor)，请确保你的 TFLite 运行时库已经编译并配置为利用这些硬件加速器。这通常需要在 TFLite
InterpreterBuilder中添加相应的委托 (Delegate)，例如TfLiteGpuDelegateV2(用于 GPU) 或TfLiteXNNPackDelegate(用于 XNNPACK，一个高性能 CPU 运行时)。- 对于 ARM 处理器，通常默认的 TFLite CPU 运行时就包含针对 ARM NEON 指令集的优化。
- 对于特定的 NPU，你可能需要使用芯片供应商提供的 TFLite Delegate。
- 内存优化: TFLite 提供了
SetNumThreads()来控制 CPU 推理线程数，以及SetCancellable等功能。合理配置这些参数以平衡性能和内存使用。 - 功耗管理:
- 跳帧检测: 在推理循环中实现跳帧检测或运动触发机制，是降低整体功耗的最有效方法。
- 低功耗模式: 在没有检测到人时，让门铃进入低功耗模式（例如关闭摄像头或降低帧率）。
- 模型更新: 考虑远程 OTA (Over-The-Air) 更新机制，以便未来能够更新模型或修复 bug。
你对这些部署步骤有什么疑问，或者想深入了解其中某个特定部分吗？
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
