【边缘AI】如何构建智能门铃的人体检测模型这篇文章我们分享了模型的架构和设计，今天我们基于 YOLOv3-Tiny 的 来开发这个模型。包括模型定义、数据集加载、训练循环、评估以及量化感知训练 (QAT) 的集成。
考虑到代码的完整性和可读性，我将代码分成几个核心部分：
- 配置 (Config)
- 模型定义 (Model Definition)**
- 数据集与数据加载器 (Dataset & Dataloader)
- 工具函数 (Utility Functions)
- 训练脚本 (Training Script)
- 量化感知训练 (QAT) 集成
评估 (Evaluation)
5. 训练脚本 (Training Script)
这是核心训练逻辑。
# train.py
import torch
import torch.optim as optim
from tqdm import tqdm
from torch.utils.tensorboard import SummaryWriter
import torch.quantization
from config import Config
from model import YOLOv3Tiny
from dataset import YOLODataset, custom_collate_fn
from utils import non_max_suppression, cells_to_boxes, mean_average_precision, save_checkpoint, load_checkpoint
# TensorBoard
writer = SummaryWriter(log_dir="runs/yolov3_tiny_person_detector")
def train_fn(train_loader, model, optimizer, loss_fn, scaler, epoch):
loop = tqdm(train_loader, leave=True)
losses = []
for batch_idx, (x, y) in enumerate(loop):
x = x.to(Config.DEVICE) # x: (batch_size, 3, H, W)
y = y.to(Config.DEVICE) # y: (num_targets, 6) -> (batch_idx, class, x, y, w, h)
# 混合精度训练 (可选，但推荐用于加速)
with torch.cuda.amp.autocast():
predictions, total_loss = model(x, y) # 在模型内部计算损失
losses.append(total_loss.item())
optimizer.zero_grad()
scaler.scale(total_loss).backward()
scaler.step(optimizer)
scaler.update()
# 更新 tqdm 进度条
mean_loss = sum(losses) / len(losses)
loop.set_postfix(loss=mean_loss)
# 记录到 TensorBoard
writer.add_scalar("Training Loss", total_loss.item(), global_step=epoch * len(train_loader) + batch_idx)
def evaluate_fn(val_loader, model, iou_threshold, conf_threshold, map_iou_threshold):
model.eval()
all_pred_boxes = [] # [[train_idx, class_pred, obj_conf, x1, y1, x2, y2], ...]
all_true_boxes = [] # [[train_idx, class_pred, x1, y1, x2, y2], ...]
loop = tqdm(val_loader, leave=True, desc="Evaluating")
for batch_idx, (x, labels) in enumerate(loop):
x = x.to(Config.DEVICE)
with torch.no_grad():
predictions = model(x) # (batch_size, num_total_predictions, 5 + num_classes)
# 将模型输出转换为 NMS 和 mAP 所需的格式
batch_boxes_list = cells_to_boxes(predictions, img_size=Config.IMAGE_SIZE)
for i, boxes in enumerate(batch_boxes_list): # 遍历每个图像的预测
# 应用 NMS
nms_boxes = non_max_suppression(
boxes,
iou_threshold=iou_threshold,
conf_threshold=conf_threshold,
num_classes=Config.NUM_CLASSES,
)
for nms_box in nms_boxes:
# [x1, y1, x2, y2, obj_conf, class_conf, class_pred]
# 转换到 mAP 格式: [train_idx, class_pred, obj_conf, x1, y1, x2, y2]
all_pred_boxes.append([batch_idx + i * Config.BATCH_SIZE, # 确保 batch_idx 唯一
int(nms_box[6]), # class_pred
nms_box[4], # obj_conf
nms_box[0], nms_box[1], nms_box[2], nms_box[3]])
# 处理真实标签
for target_idx in range(labels.shape[0]):
batch_id = int(labels[target_idx, 0].item())
class_id = int(labels[target_idx, 1].item())
x_center, y_center, width, height = labels[target_idx, 2:].tolist()
# 将归一化 [0, 1] 的中心点+宽高转换为绝对像素 x1, y1, x2, y2
x1 = (x_center - width / 2) * Config.IMAGE_SIZE
y1 = (y_center - height / 2) * Config.IMAGE_SIZE
x2 = (x_center + width / 2) * Config.IMAGE_SIZE
y2 = (y_center + height / 2) * Config.IMAGE_SIZE
all_true_boxes.append([batch_id + i * Config.BATCH_SIZE, class_id, x1, y1, x2, y2])
# 计算 mAP
mAP = mean_average_precision(
all_pred_boxes,
all_true_boxes,
iou_threshold=map_iou_threshold,
num_classes=Config.NUM_CLASSES,
)
model.train() # 评估完后切换回训练模式
return mAP
def main():
model = YOLOv3Tiny(num_classes=Config.NUM_CLASSES, config_anchors=Config.ANCHORS).to(Config.DEVICE)
optimizer = optim.Adam(model.parameters(), lr=Config.LEARNING_RATE)
scaler = torch.cuda.amp.GradScaler() # 混合精度训练
# 数据加载
train_dataset = YOLODataset(
Config.TRAIN_IMG_DIR,
Config.TRAIN_LABEL_DIR,
Config.ANCHORS,
img_size=Config.IMAGE_SIZE,
num_classes=Config.NUM_CLASSES,
is_train=True,
augment_config=Config.AUGMENTATION_CONFIG,
)
train_loader = DataLoader(
dataset=train_dataset,
batch_size=Config.BATCH_SIZE,
num_workers=os.cpu_count() // 2, # 根据你的CPU核心数调整
shuffle=True,
collate_fn=custom_collate_fn,
pin_memory=True, # 提高数据传输效率
)
val_dataset = YOLODataset(
Config.VAL_IMG_DIR,
Config.VAL_LABEL_DIR,
Config.ANCHORS,
img_size=Config.IMAGE_SIZE,
num_classes=Config.NUM_CLASSES,
is_train=False, # 验证集不进行数据增强
)
val_loader = DataLoader(
dataset=val_dataset,
batch_size=Config.BATCH_SIZE,
num_workers=os.cpu_count() // 2,
shuffle=False,
collate_fn=custom_collate_fn,
pin_memory=True,
)
if Config.LOAD_MODEL:
load_checkpoint(Config.LOAD_MODEL_PATH, model, optimizer)
# --- QAT 设置 ---
if Config.QUANT_MODE:
print("--- Enabling Quantization Aware Training ---")
# 准备模型进行 QAT
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm') # 适用于 CPU 后端
# 可以考虑 'qnnpack' 用于 ARM CPU
# model.qconfig = torch.quantization.get_default_qat_qconfig('qnnpack')
# 融合模块 (Conv+BN+ReLU)，这样 QAT 可以在融合层上进行量化
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_fused = torch.quantization.prepare_qat(model_fused, inplace=True)
model = model_fused # 使用 QAT 准备好的模型进行训练
print("Model fused and prepared for QAT.")
# 校准 (可选，但推荐在 QAT 之前进行少量校准以初始化统计量)
print("Calibrating model for QAT...")
model.eval()
with torch.no_grad():
for i, (x, _) in enumerate(tqdm(train_loader, desc="Calibration")):
if i >= Config.NUM_CALIBRATION_BATCHES:
break
x = x.to(Config.DEVICE)
model(x)
print("Calibration finished.")
model.train() # 回到训练模式
# --- 训练循环 ---
best_mAP = 0.0
for epoch in range(Config.NUM_EPOCHS):
print(f"Epoch {epoch+1}/{Config.NUM_EPOCHS}")
train_fn(train_loader, model, optimizer, None, scaler, epoch) # Loss_fn 已经在模型内部处理
if Config.SAVE_MODEL:
if epoch % 5 == 0 or epoch == Config.NUM_EPOCHS - 1: # 每5个epoch保存一次，或在最后保存
save_checkpoint(model, optimizer, filename=Config.MODEL_PATH.replace(".pth.tar", f"_epoch{epoch+1}.pth.tar"))
# 评估模型
mAP = evaluate_fn(val_loader, model, Config.CONF_THRESHOLD, Config.NMS_IOU_THRESH, Config.MAP_IOU_THRESH)
print(f"Validation mAP: {mAP:.4f}")
writer.add_scalar("Validation mAP", mAP, global_step=epoch)
if mAP > best_mAP:
best_mAP = mAP
if Config.SAVE_MODEL:
save_checkpoint(model, optimizer, filename=Config.MODEL_PATH)
print(f"New best mAP: {best_mAP:.4f}, model saved to {Config.MODEL_PATH}")
writer.close()
# --- 训练完成后，如果进行了 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.jit.trace 或 torch.save(state_dict)
# 对于部署到 TFLite，通常需要先导出 ONNX 或直接使用 PyTorch Mobile/TorchScript
# 这里只是保存 PyTorch 格式的量化模型，方便验证
torch.save(quantized_model.state_dict(), "yolov3_tiny_person_quantized.pth")
print("Quantized model saved to yolov3_tiny_person_quantized.pth")
# 验证量化模型大小 (可选)
# from torch.optim.lr_scheduler import ReduceLROnPlateau
# from torch.quantization.quantize_fx import prepare_fx, convert_fx
# from torch.fx.graph_module import GraphModule
# from torch.quantization.qconfig import QConfig
# 通过导出 TorchScript 或 ONNX 来进行 TFLite 转换
# 这一步通常在训练结束后单独进行
print("To export to TFLite, you'll need to use TorchScript/ONNX export:")
print("Example: torch.jit.trace(quantized_model, dummy_input).save('model.pt')")
print("Then use TF Lite Converter with Torch-TF Bridge or ONNX-TF for conversion.")
if __name__ == "__main__":
# 在运行前，请确保在 `data/doorbell_person_dataset` 下有 `images/train`, `labels/train`, `images/val`, `labels/val` 目录
# 并放入你的图像和YOLO格式的标注文件。
main()
6. 量化感知训练 (QAT) 的进一步说明
在上面的 train.py 中，我已经集成了 PyTorch 的 QAT 流程。关键步骤：
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm'): 设置量化配置。fbgemm是针对服务器 CPU 后端的，对于 ARM CPU 来说，更推荐使用qnnpack，即model.qconfig = torch.quantization.get_default_qat_qconfig('qnnpack')。torch.quantization.fuse_modules: 将串联的Conv -> BN -> ReLU层融合为一个操作。这有助于减少量化误差并提高推理效率。你可能需要根据你的ConvBlock结构调整融合的层名。torch.quantization.prepare_qat(model_fused, inplace=True): 在模型中插入假量化 (fake quantization) 模块。在训练过程中，这些模块会模拟量化对模型权重和激活值的影响。- 校准 (Calibration)：虽然 QAT 会在训练期间学习量化参数，但在
prepare_qat之后进行少量数据校准可以帮助初始化量化统计量（如 min/max 范围），这通常能带来更好的训练稳定性。 torch.quantization.convert(model, inplace=False): 在 QAT 训练结束后，调用此函数将模型转换为真正的量化模型。它会将权重转换为 INT8，并将浮点操作替换为量化操作。
从 PyTorch 量化模型到 TFLite 部署
PyTorch 直接导出的 INT8 模型不能直接在 TFLite 上运行，你需要一个中间步骤：
-
导出 TorchScript：
-
在 QAT 训练并
convert得到quantized_model后，你可以将其导出为 TorchScript：dummy_input = torch.randn(1, 3, Config.IMAGE_SIZE, Config.IMAGE_SIZE).to(Config.DEVICE) traced_model = torch.jit.trace(quantized_model, dummy_input) traced_model.save("yolov3_tiny_quantized.pt") -
然后，你可以尝试使用 PyTorch 提供的 Torch-TensorFlow Bridge 将 TorchScript 模型转换为 TensorFlow Graph，再从 TensorFlow Graph 转换为 TFLite。这个流程目前仍在发展中，可能需要一些手动调整。
-
-
导出 ONNX：
-
将 PyTorch 模型导出为 ONNX 格式：
dummy_input = torch.randn(1, 3, Config.IMAGE_SIZE, Config.IMAGE_SIZE).to(Config.DEVICE) torch.onnx.export(quantized_model, dummy_input, "yolov3_tiny_quantized.onnx", opset_version=13, # 选择合适的 opset 版本 do_constant_folding=True, input_names=['input'], output_names=['output']) -
然后使用 ONNX-TensorFlow 工具将 ONNX 模型转换为 TensorFlow SavedModel 格式：
# 命令行安装 onnx-tf pip install onnx onnx_tf # 在 Python 中 from onnx_tf.backend import prepare import onnx onnx_model = onnx.load("yolov3_tiny_quantized.onnx") tf_rep = prepare(onnx_model) tf_rep.export_graph("yolov3_tiny_quantized_tf_savedmodel") -
最后，使用 TensorFlow Lite Converter 将 SavedModel 转换为 TFLite：
import tensorflow as tf converter = tf.lite.TFLiteConverter.from_saved_model("yolov3_tiny_quantized_tf_savedmodel") converter.optimizations = [tf.lite.Optimize.DEFAULT] # 如果模型已经是 INT8 量化，这里不需要指定 representative_dataset # 但如果 ONNX 到 TF 转换过程中丢失了量化信息，可能需要 Post-Training Integer Quantization # converter.representative_dataset = representative_data_gen # 参见之前给出的示例 # 确保输出张量类型是 INT8 (如果量化成功) converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] tflite_model = converter.convert() with open("yolov3_tiny_person_detector_quantized.tflite", "wb") as f: f.write(tflite_model)
-
运行
-
准备数据集：将你的图像和 YOLO 格式的标注文件放在
data/doorbell_person_dataset/images/train,labels/train,images/val,labels/val目录下。 -
创建文件：将上述代码分别保存为
config.py,model.py,dataset.py,utils.py,train.py。 -
运行训练：
python train.py
重要提示：
- 锚框 (Anchors)：这里提供的锚框是 YOLOv3 针对 COCO 数据集训练的默认锚框。对于你的智能门铃人体检测数据集，建议运行 K-means 聚类算法来生成最适合你数据集的锚框。这会显著提高模型的检测性能。
- 超参数调优：学习率、批大小、数据增强参数和训练轮次都需要根据你的数据集进行仔细的调优。
- 预训练权重：加载预训练的 YOLOv3-Tiny 权重（例如在 ImageNet 或 COCO 上训练的）将大大加速训练并提高性能。你可能需要找到并下载一个 PyTorch 版本的 YOLOv3-Tiny 预训练权重。
- 环境：确保你的 PyTorch 安装支持 CUDA（如果你有 GPU），否则它将默认在 CPU 上运行，训练速度会慢很多。
- 调试：在
YOLOLayer中损失计算可能比较复杂，如果遇到NaN或收敛问题，请仔细检查损失函数和数据处理部分。
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
