使用 TensorFlow 训练对象检测模型

2025-06-10

配置 TensorFlow 后，我们将使用 YOLO 体系结构来训练对象检测模型。 YOLO 是一个神经网络，用于在一次计算中预测图像中的边界框和类概率。 YOLO 模型每秒可以处理超过 60 帧，使其成为检测视频中对象的出色体系结构。可以在此处找到有关 YOLO 工作原理的详细信息。

使用 YOLO

首先，下载此 YOLO 示例文件，其中包含要开始使用的帮助程序脚本。

使用 YOLO 时，我们有 3 个选项：

使用 YOLO 的预训练模型权重。 预先训练的模型已在大型数据集上训练，其中包含 80 个类（类别），用于日常对象，如公共汽车、人员、三明治等。如果要以 ONNX 格式下载预先训练的 YOLO 模型，可在此处执行此作。然后，可以继续学习本教程的最后阶段，了解如何将该模型集成到应用中。
使用自定义数据集实现转移学习。 转移学习是一种将训练的模型用作训练模型解决不同但相关的任务的起点的方法。本教程将使用预训练的，具有 80 个类别的 YOLO 权重，结合 VOC 数据集来训练一个具有 20 个类别的模型。如果要使用自定义类创建自己的数据集，请参阅此处的说明。
从头开始训练 YOLO。 不建议使用此方法，因为很难聚合。最初的YOLO论文在ImageNet上训练了darknet（包含数十万张照片），然后再训练整个网络。

在预先训练的 YOLO 权重上实现对 VOC 数据集的转移学习：

让我们继续执行第二个选项，并通过以下步骤实现转移学习。

在 miniconda 窗口中，导航到 yolo 示例目录，并运行以下命令以安装 YOLO 所需的所有 pip 包。

pip install -r requirements.txt

运行安装脚本以下载数据和预先训练的权重

python setup.py

转换数据集。请参阅 tools/voc2012.py 实现 - 此格式基于 tensorflow 对象检测 API。许多字段不是必需的，但已在此处填写，以便与官方 API 兼容。

python tools/voc2012.py \
  --data_dir './data/voc2012_raw/VOCdevkit/VOC2012' \
  --split train \
  --output_file ./data/voc2012_train.tfrecord

python tools/voc2012.py \
  --data_dir './data/voc2012_raw/VOCdevkit/VOC2012' \
  --split val \
  --output_file ./data/voc2012_val.tfrecord

训练模型。运行以下命令：

python convert.py
python detect.py --image ./data/meme.jpg # Sanity check

python train.py \
	--dataset ./data/voc2012_train.tfrecord \
	--val_dataset ./data/voc2012_val.tfrecord \
	--classes ./data/voc2012.names \
	--num_classes 20 \
	--mode fit --transfer darknet \
	--batch_size 16 \
	--epochs 10 \
	--weights ./checkpoints/yolov3.tf \
	--weights_num_classes 80

现在，你已有一个具有 20 个类的重新训练模型，可供使用。

后续步骤

创建 TensorFlow 模型后，需要将其转换为 ONNX 格式，以便与 Windows 机器学习 API 一起使用。

反馈

此页面是否有帮助？