当前位置：首页 > news >正文

基于YOLOv3-Tiny 的智能门铃的人体检测模型的实现（中）

news 2025/7/14 7:57:33

【边缘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.quantizationfrom 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)# 记录到 TensorBoardwriter.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): # 遍历每个图像的预测# 应用 NMSnms_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_prednms_box[4],      # obj_confnms_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, y2x1 = (x_center - width / 2) * Config.IMAGE_SIZEy1 = (y_center - height / 2) * Config.IMAGE_SIZEx2 = (x_center + width / 2) * Config.IMAGE_SIZEy2 = (y_center + height / 2) * Config.IMAGE_SIZEall_true_boxes.append([batch_id + i * Config.BATCH_SIZE, class_id, x1, y1, x2, y2])# 计算 mAPmAP = mean_average_precision(all_pred_boxes,all_true_boxes,iou_threshold=map_iou_threshold,num_classes=Config.NUM_CLASSES,)model.train() # 评估完后切换回训练模式return mAPdef 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 ---")# 准备模型进行 QATmodel.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:breakx = x.to(Config.DEVICE)model(x)print("Calibration finished.")model.train() # 回到训练模式# --- 训练循环 ---best_mAP = 0.0for 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 = mAPif 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 onnxonnx_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 tfconverter = 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 或收敛问题，请仔细检查损失函数和数据处理部分。

查看全文

http://www.dtcms.com/a/277556.html