Compartir a través de

Entrenamiento de un modelo de detección de objetos con TensorFlow

  • Artículo

Ahora que ha configurado TensorFlow, tendrá que usar la arquitectura YOLO para entrenar el modelo de detección de objetos. YOLO es una red neuronal que predice los cuadros delimitadores y las probabilidades de clase a partir de una imagen en una sola evaluación. Los modelos de YOLO pueden procesar más de 60 fotogramas por segundo, por lo que es una arquitectura excelente para detectar objetos en vídeos. Puede encontrar más información sobre cómo funciona YOLO aquí.

Uso de YOLO

En primer lugar, descargue este archivo de ejemplo de YOLO, que contiene scripts auxiliares para empezar a trabajar.

Al usar YOLO, tiene tres opciones:

  1. Use pesos de modelos previamente entrenados para YOLO. El modelo entrenado previamente se ha entrenado en un conjunto de datos grande con 80 clases (categorías) de objetos cotidianos como autobuses, personas, sándwiches, etc. Si quiere descargar un modelo YOLO entrenado previamente en formato ONNX, puede hacerlo aquí. A continuación, puede continuar con la fase final de este tutorial para aprender a integrar ese modelo en una aplicación.

  2. Implemente el aprendizaje de transferencia con un conjunto de datos personalizado. El aprendizaje de transferencia es un método para usar un modelo entrenado como punto de partida para entrenar un modelo que resuelve una tarea diferente pero que está relacionada. En este tutorial se usarán las ponderaciones YOLO previamente entrenadas con 80 clases, para entrenar un modelo con 20 clases con el conjunto de datos VOC. Si quiere crear su propio conjunto de datos con clases personalizadas, consulte las instrucciones aquí.

  3. Entrenamiento de YOLO desde cero. Esta técnica no se recomienda, ya que es muy difícil de converger. El documento de YOLO original entrenó Darknet en ImageNet (que contiene cientos de miles de fotos) antes de entrenar también toda la red.

Implemente el aprendizaje de transferencia en pesos de YOLO previamente entrenados en el conjunto de datos de VOC:

A continuación, seleccione la segunda opción e implemente el aprendizaje de transferencia con los pasos siguientes.

  1. En una ventana de miniconda, vaya al directorio de ejemplo de YOLO y ejecute el siguiente comando para instalar todos los paquetes PIP necesarios para YOLO.

pip install -r requirements.txt

  1. Ejecute el script de instalación para descargar los datos y las ponderaciones previamente entrenadas.

python setup.py

  1. Transforme el conjunto de datos. Consulte tools/voc2012.py para la implementación: este formato se basa en la API de detección de objetos de tensorflow. Muchos campos no son necesarios, pero aquí se han rellenado por cuestiones de compatibilidad con la API oficial.
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. Entrene el modelo. Ejecute los comandos siguientes:
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

Ahora ya tiene un modelo entrenado de nuevo con 20 clases, listo para su uso.

Pasos siguientes

Como ya ha creado un modelo de TensorFlow, es necesario convertirlo al formato ONNX para poder usarlo con las API de Windows Machine Learning.