通过

使用 TensorFlow 训练物体检测模型

  • 项目

现在我们配置了 TensorFlow，将使用 YOLO 体系结构来训练物体检测模型。 YOLO 是一种神经网络，可在单个评估中根据某张图像预测范围框和类别概率。 YOLO 模型每秒可处理超过 60 帧，这使其成为用于检测视频中的物体的出色体系结构。 可在此处找到有关 YOLO 工作原理的详细信息。

使用 YOLO

首先，下载此 YOLO 示例文件，其中包含入门帮助程序脚本。

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

  1. 使用 YOLO 预先训练的模型权重。 预先训练的模型已在一个包含 80 个日常物体类别（如公共汽车、人、三明治等）的大型数据集上进行了训练。如果想下载 ONNX 格式的预先训练的 YOLO 模型，可在此处进行下载。 然后，可以转到本教程的最后一个阶段，了解如何将该模型集成到应用中。

  2. 使用自定义数据集实现迁移学习。 迁移学习是一种方法，它使用经过训练的模型作为起点来训练可解决各种相关任务的模型。 本教程将使用包含 80 个类的预先训练的 YOLO 权重，使用 VOC 数据集 训练包含 20 个类的模型。 如果想使用自定义类创建自己的数据集，请参阅此处的说明。

  3. 从头开始训练 YOLO。 不建议使用此方法，因为这种方法很难达到一致的结果。 最初的 YOLO 论文在训练整个网络之前在 imagenet（包含数十万张照片）上训练了暗网。

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

继续第二个选项，并通过以下步骤实现迁移学习。

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

pip install -r requirements.txt

  1. 运行设置脚本以下载数据和预先训练的权重

python setup.py

  1. 转换数据集。 有关实现，请参阅 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
  1. 训练模型。 运行以下命令：
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 一起使用。