You Only Look Once v4 with TensorFlow and DALI

This example presents a sample implementation of a YOLOv4 network, based on the following paper – Alexey Bochkovskiy, Chien-Yao Wang, Hong-Yuan Mark Liao, YOLOv4: Optimal Speed and Accuracy of Object Detection.

The implementation is based on the code available in the darknet repository.

To run training from scratch on all available GPUs with DALI, run the following command:

python src/main.py train /coco/train2017 /coco/annotations/instances_train2017.json \
  -b 8 -e 6 -s 1000 -o output.h5 \
  --pipeline dali-gpu --multigpu --use_mosaic

To save checkpoints each epoch, add the following flag to the command:

--ckpt_dir /ckpt

Resume the training by providing an appropriate path to the -w flag:

-w /ckpt/epoch_3.h5

To perform evaluation every second epoch, add the following flags:

--eval_file_root /coco/val2017 \
--eval_annotations /coco/annotations/instances_val2017.json \
--eval_frequency 2 --eval_steps 500

To evaluate trained model, run the following command:

python src/main.py eval /coco/val2017 /coco/annotations/instances_val2017.json \
  -w output.h5 -b 1 -s 5000

To perform an inference and display the results on screen, run the following command:

python src/main.py infer image.png -w output.h5 -c coco-labels.txt

Requirements

• This example was tested using python 3.8 and tensorflow 2.4.1 and should work on later versions.

• The model requires COCO 2017 dataset for training.

• Besides TensorFlow, the following python packages are also required:

  matplotlib tensorflow-addons pycocotools

Usage

Training

usage: main.py train [-h] file_root annotations
  [--batch_size BATCH_SIZE] [--epochs EPOCHS] [--steps STEPS] [--output OUTPUT]
  [--start_weights START_WEIGHTS] [--log_dir LOG_DIR] [--ckpt_dir CKPT_DIR]
  [--pipeline PIPELINE] [--multigpu] [--use_mosaic] [--learning_rate LEARNING_RATE]
  [--eval_file_root EVAL_FILE_ROOT] [--eval_annotations EVAL_ANNOTATIONS]
  [--eval_steps EVAL_STEPS] [--eval_frequency EVAL_FREQUENCY]
  [--seed SEED]

Argument	Description
-h, –help	show this help message and exit
file_root	path to folder containing train coco2017 images
annotations	path to instances_train2017.json file
–batch_size BATCH_SIZE, -b BATCH_SIZE	number of images per training step, default = 8
–epochs EPOCHS, -e EPOCHS	number of training epochs, default = 5
–steps STEPS, -s STEPS	number of training steps per epoch, default = 1000
–output OUTPUT, -o OUTPUT	path to a .h5 output file for trained model, default = output.h5
–start_weights START_WEIGHTS, -w START_WEIGHTS	initial weights file in h5 or YOLO format
–log_dir LOG_DIR	path to a directory for TensorBoard logs
–ckpt_dir CKPT_DIR	path to a directory for checkpoint files
–pipeline PIPELINE	either dali_gpu, dali_cpu or numpy
–multigpu	if present, training is run using all available GPUs
–use_mosaic	if present, mosaic data augmentation is used
–learning_rate LEARNING_RATE	learning rate for training, default = 1e-3
–eval_file_root EVAL_FILE_ROOT	path to folder containing val coco2017 images
–eval_annotations EVAL_ANNOTATIONS	path to instances_val2017.json file
–eval_steps EVAL_STEPS	number of images per evaluation step, default = 5000
–eval_frequency EVAL_FREQUENCY	number of training epochs between each evaluation, default = 5
–seed SEED	seed for DALI and TensorFlow

Inference

usage: main.py infer [-h] image [--weights WEIGHTS] [--classes CLASSES]
                     [--output OUTPUT]

Argument	Description
-h, –help	show this help message and exit
image	path to an image to perform inference on
–weights WEIGHTS, -w WEIGHTS	path to a trained weights file in h5 or YOLO format
–classes CLASSES, -c CLASSES	path to a coco-labels.txt file
–output OUTPUT, -o OUTPUT	path to an output image

Evaluation

usage: main.py eval [-h] file_root annotations [--weights WEIGHTS]
                    [--batch_size BATCH_SIZE] [--steps STEPS]

Argument	Description
-h, –help	show this help message and exit
file_root	path to folder containing val coco2017 images
annotations	path to instances_val2017.json file
–weights WEIGHTS, -w WEIGHTS	path to a trained weights file in h5 or YOLO format
–batch_size BATCH_SIZE, -b BATCH_SIZE	number of images per evaluation step, default = 1
–steps STEPS, -s STEPS	number of evaluation steps, default = 1000