YOLOv4 on Google Colab: Train your Custom Dataset (Traffic signs) with ease | by Quang Nguyen | Towards Data Science

Colab Notebook for training YOLOv4 with custom dataset (traffic signs)
Introduction
Why YOLOv4?
Data preparation
Training with Colab
Predict with YOLOv4
Conclusion
A while ago, I wrote a tutorial on training YOLOv3 with a custom dataset (gun detection) using the free GPU provided by Google Colab. After publishing the tutorial, many people emailed me asking about problems they faced during their training process. What I noticed is that the majority of those issues come from modifying inappropriate layer architecture, running cells in incorrect directory paths or missing one/some of the required configuration steps. I also realized that the setting up process might take time and often can be quite confusing.
For those reasons, I have came up with a new approach in which most of the setting up steps are automated with a python script (Colab allows you modifying and running python script directly on its environment). With this new procedure, the only thing you need to perform locally is to prepare your custom dataset following YOLO format, while all of the remains would be executed in Colab environment.
YOLOv4 was developed by Alexey Bochkovskiy, Chien-Yao Wang, and Hong-Yuan Mark Liao. It was released in April 2020 and claimed as one of the state-of-the-art real-time object detectors at the time. According to its paper, YOLOv4 is 12% faster and 10% more accurate compare to YOLOv3. The new architecture of YOLOV4 is built with CSPDarknet54 as a backbone, which promotes the learning capability of CNN. Furthermore, the implement of universal features includes Weighted-Residual-Connections (WRC), Cross-Stage-Partial-connections (CSP), Cross mini-Batch Normalization (CmBN), Self-adversarial-training (SAT) and Mish-activation help YOLOv4 obtain a very impressive results. To get more details about YOLOv4, you can refer to the original paper here. [1]
To be able to train our custom dataset with YOLOv4, it is mandatory that our dataset follows the YOLO format. Each image from the dataset would associate with a .txt file having the same name, which contains the object classes and their coordinate following this syntax: <object-class> <x_center> <y_center> <width> <height>
There are many open-source GUI tools which can help you easily generate label file from image such as Yolo_label, OpenLabeling, Yolo_mark, BBox-Label-Tool, etc. Just simply drag and drop your mouse to create a bounding box around your objects, then the tool would generate the label file automatically.
Next, you need to create 3 files: classes.names, train.txt and test.txt. The classes.names contains the names of your objects following format:
object1_nameobject2_nameobject3_name...objectn_name
Make sure that the index of your <object-class> from label file .txt is corresponded to the index of your classes name, e.g, <object-class> of object1_name, object2_name, object3_name are 0, 1, 2 respectively.
train.txt and test.txt contains file paths of your training images and testing images. The idea is that you split your dataset into training set and test sets so that you will train your model on the training set and validate it on the test set. If the loss on your training set is high, it means your model is underfitting and you will need to train longer. If the loss is low on your training set and high on your test set, it means your model is overfitting and you will need to add more data. Depending on the number of images in your dataset that your validation set can be extracted from around 5% (small dataset) to 30% (large dataset) of the total dataset. These two files share the same format syntax as data/<image_path>.
In this tutorial, I will use the Traffic Signs Dataset in YOLO format from Kaggle. You don’t have to download it now since we will directly download it from Colab environment later. Examples of some of its images and labels can be seen as below (notice that one image can have multiple objects):
00058.txt: 3 0.5919117647058824 0.518125 0.027941176470588237 0.05125 00011.txt: 0 0.5477941176470589 0.46 0.03676470588235294 0.075 1 0.5477941176470589 0.379375 0.051470588235294115 0.09625 00334.txt: 0 0.25441176470588234 0.59375 0.033823529411764704 0.0575 0 0.5724264705882353 0.566875 0.027205882352941177 0.04625 00136.txt: 1 0.6150735294117647 0.52375 0.030147058823529412 0.045
In summary, the traffic sign dataset which we would use in this tutorial would be:
|- ts/|   |- 00000.jpg|   |- 00000.txt|   |- ...|- classes.names|- train.txt|- test.txt
For who is already familiar with Colab, you can jump right in my Colab notebook here and start playing around with it. However, make sure that you enable GPU in advance by select Runtime -> Change runtime type>Hardware accelerator>GPU, create directory YOLOv4_weight/backup in your Google Drive and mount Drive with Colab environment.
As mentioned before, I have implemented some modifications which help the training process more straightforward:
assert function: a function that checks whether you’re running the cell in the correct directory path. If you’re in the wrong directory, it will give you a hint to which directory you should be. In case you don’t know how to change the directory, please refer to this cd command in Linux terminal.
darknet_for_colab: darknet folder which was modified specifically to adapt with Colab environment (no MAKEFILE change necessary).
yolov4_config.py: Taking the advantage of the direct python editing feature on Colab, you can now define training parameters just by double click on yolov4_config.py and edit it (Figure 1). For example, I will set my classes=4 (our traffic sign dataset has 4 classes), max_batches=8000 (number of training iterations), batch=64 (number of samples in one batch), subdivisions=16 (number of mini_batches in one batch), etc. You can refer this page to find more details about the meaning of each parameter.
yolov4_setup.py (cell [6]): a python script which automatically generates YOLOv4 architecture config files (yolov4_custom_train.cfg and yolov4_custom_test.cfg) based on user-input parameters in yolov4_config.py.
The notebook below demonstrates the pipeline of YOLOv4 training procedure on Colab. I will recommend you to take a look at my Colab Notebook to see what kind of output of each cell you should expect.
# check whether GPU is provided
import os
!nvcc --version

# remove the existing folder if have
#!rm -r darknet_for_colab

# download and compile darknet_for_colab
assert os.getcwd()=='/content', 'Directory should be "/content" instead of "{}"'.format(os.getcwd())
!git clone https://github.com/quangnhat185/darknet_for_colab.git
%cd darknet_for_colab
!make
!chmod +x ./darknet

assert os.getcwd()=='/content/darknet_for_colab', 'Directory should be "/content/darknet_for_colab" instead of "{}"'.format(os.getcwd())
!wget https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.conv.137

%cd data
assert os.getcwd()=='/content/darknet_for_colab/data', 'Directory should be "/content/darknet_for_colab/data" instead of "{}"'.format(os.getcwd())

# download custom data of common traffic signs
!wget --no-check-certificate "https://onedrive.live.com/download?cid=A86CBC7F31A1C06B&resid=A86CBC7F31A1C06B%21109&authkey=ACmKGNgx4G2TT4A" -O ts.zip
!unzip ts.zip
!rm -f ts.zip
!ls
%cd ..

assert os.getcwd()=='/content/darknet_for_colab', 'Directory should be "/content/darknet_for_colab" instead of "{}"'.format(os.getcwd())
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import glob

def read_label(image_path):
  file_name = image_path.replace('.jpg', '.txt')
  with open(file_name, 'rt') as file:
    print(os.path.basename(file_name) + ': \n' + file.read())

image_path = glob.glob("data/ts/*.jpg")
fig = plt.figure(figsize=(12,8))
cols = 2
rows = 2
grid = gridspec.GridSpec(nrows=rows, ncols=cols, figure=fig)
for i in range(cols*rows):
  fig.add_subplot(grid[i])
  image=plt.imread(image_path[i])
  plt.title(os.path.basename(image_path[i]))
  plt.axis(False)
  plt.imshow(image)
  read_label(image_path[i])

plt.savefig("dataset_examples.jpg", dpi=300)

Double click on file yolov4_config.py to modify the hyperpameters directly from Colab environment
E.g: I will train my dataset with these parameters:
classes= 4, 
max_batches=8000
batch=64
subdivisions=16
width=416
height=416
...
assert os.getcwd()=='/content/darknet_for_colab', 'Directory should be "/content/darknet_for_colab" instead of "{}"'.format(os.getcwd())

# Run python script to create our customize yolov4_custom_train.cfg 
# and yolov4_custom_tes.cfg in folder /cfg
!python yolov4_setup.py

assert os.getcwd()=='/content/darknet_for_colab', 'Directory should be "/content/darknet_for_colab" instead of "{}"'.format(os.getcwd())

# create Symlinks so we can save trained weight in our Google Drive
# create folder YOLOv4_weight/back in your Drive to store trained weights
!rm /content/darknet_for_colab/backup -r
!ln -s /content/drive/'My Drive'/YOLOv4_weight/backup /content/darknet_for_colab

assert os.getcwd()=='/content/darknet_for_colab', 'Directory should be "/content/darknet_for_colab" instead of "{}"'.format(os.getcwd())

# Training with YOLOv4
!./darknet detector train data/yolov4.data cfg/yolov4_custom_train.cfg yolov4.conv.137 -dont_show -map

# If you get CUDA out of memory adjust subdivisions above!
# adjust max batches down for shorter training above

assert os.getcwd()=='/content/darknet_for_colab', 'Directory should be "/content/darknet_for_colab" instead of "{}"'.format(os.getcwd())

# Plotting training result after 2000 epochs
import matplotlib.pyplot as plt
fig = plt.figure(figsize=(10,10))
train_result = plt.imread("chart.png")
plt.axis(False)
plt.imshow(train_result)

assert os.getcwd()=='/content/darknet_for_colab', 'Directory should be "/content/darknet_for_colab" instead of "{}"'.format(os.getcwd())

# show labels of our custom traffic sign dataset
!cat data/classes.names

assert os.getcwd()=='/content/darknet_for_colab', 'Directory should be "/content/darknet_for_colab" instead of "{}"'.format(os.getcwd())

# download example images an video
!wget --no-check-certificate "https://onedrive.live.com/download?cid=A86CBC7F31A1C06B&resid=A86CBC7F31A1C06B%21120&authkey=AM5VslNNW9a8aO8" -O examples.zip
!unzip examples.zip
!rm -r examples.zip

assert os.getcwd()=='/content/darknet_for_colab', 'Directory should be "/content/darknet_for_colab" instead of "{}"'.format(os.getcwd())

# cfg/coco.data was harcode in darknet, thus we need 
# to duplicate one with the same content of data/yolov4.data
%cp data/yolov4.data cfg/coco.data

assert os.getcwd()=='/content/darknet_for_colab', 'Directory should be "/content/darknet_for_colab" instead of "{}"'.format(os.getcwd())

# coco.names is hardcoded somewhere in the detector
# the processed image would be save as file 'prediction.jpg'
img_path = "examples/test_image_4.jpg"
!./darknet detect cfg/yolov4_custom_test.cfg backup/yolov4_custom_train_last.weights {img_path} -dont-show

assert os.getcwd()=='/content/darknet_for_colab', 'Directory should be "/content/darknet_for_colab" instead of "{}"'.format(os.getcwd())

# plot predict image
import matplotlib.pyplot as plt 
fig = plt.figure(figsize=(12,12))
plt.axis(False)
processed_image = plt.imread("./predictions.jpg")
plt.imshow(processed_image)

assert os.getcwd()=='/content/darknet_for_colab', 'Directory should be "/content/darknet_for_colab" instead of "{}"'.format(os.getcwd())

# video processed
!python darknet_video.py -v examples/test_video.mp4 -c cfg/yolov4_custom_test.cfg -w backup/yolov4_custom_train_last.weights -o output.mp4

Although I did define my max_batches=8000, both accuracy and loss from training result didn’t improve much after 2000 iterations (Figure 2). The number of iterations might increase if you train with more classes or harder-to-learn dataset. Make sure that you monitor the loss and accuracy by double click on file chart.png or just simply look at the training statistic from the output cell.
After obtain the training weights, there are several ways to deploy YOLOv4 with third-party frameworks including OpenCV, Keras, Pytorch, etc. However, those are beyond the scope of this tutorial. You can predict an image or video using the same Colab working space which you’ve used for training process (Figure 3 and Figure 4). Cell 13 and 15 gives you an explicit way to predict an image and video respectively. By default, the predicted image is saved at predictions.jpg while the predicted video would be saved at output.mp4.
This tutorial has introduced a new approach which allows you training your custom dataset with YOLOv4 on Google Colab with ease. All modifications relating to neural network architecture and training parameters are automated and can be performed within Colab environment, while unit tests are integrated to debug common compiling errors.
However, I would like to point out that using an automated tool has its own price. The fact that you have no idea what is happening behind might be disturbing, especially people who want to have a better intuition about YOLOv4 architecture. In that case, you can refer to the original GitHub repository of YOLOv4 or read my article on training YOLOv3 on Colab.
Train YOLOv4 on Colab notebook
Darknet for colab repository
YOLOv4 weights for traffic sign detection (2000 iterations)
Traffic signs dataset in YOLO format
[1] Bochkovskiy, Alexey, Chien-Yao Wang, and Hong-Yuan Mark Liao. “YOLOv4: Optimal Speed and Accuracy of Object Detection.” arXiv preprint arXiv:2004.10934 (2020).