YOLOv4 Training on Colab

Training Step

  • 2020-10-15
  • Google Colab
  • TF: v2.3.0
  • yolov4: v1.2.1
  • coco 2017 dataset

Installation

python3 -m pip install -U yolov4

Prepare dataset

Usually, I create an image path related file using only image file names. Then, I make full path using image_path_prefix in yolo.load_dataset().

train2017.txt
000000000009.jpg 45,0.479492,0.688771,0.955609,0.595500 45,0.736516,0.247188,0.498875,0.476417 50,0.637063,0.732938,0.494125,0.510583 45,0.339438,0.418896,0.678875,0.781500 49,0.646836,0.132552,0.118047,0.096937 49,0.773148,0.129802,0.090734,0.097229 49,0.668297,0.226906,0.131281,0.146896 49,0.642859,0.079219,0.148063,0.148062
000000000025.jpg 23,0.770336,0.489695,0.335891,0.697559 23,0.185977,0.901608,0.206297,0.129554
000000000030.jpg 58,0.519219,0.451121,0.398250,0.757290 75,0.501188,0.592138,0.260000,0.456192
...
/content
โ”œโ”€โ”€ train2017
โ”‚ โ”œโ”€โ”€ 000000000009.jpg
โ”‚ โ”œโ”€โ”€ 000000000025.jpg
โ”‚ โ”œโ”€โ”€ 000000000030.jpg
โ”‚ โ””โ”€โ”€ ...
โ””โ”€โ”€ ...
# Call load_dataset() after yolo.<parameter>: anchors, classes, input_size, xyscales, ...
train_data_set = yolo.load_dataset(
"/content/drive/My Drive/Hard_Soft/NN/coco/train2017.txt",
image_path_prefix="/content/train2017",
label_smoothing=0.05
)

Set yolo parameters

yolo = YOLOv4(tiny=True)
yolo.classes = "/content/drive/My Drive/Hard_Soft/NN/coco/coco.names"
yolo.input_size = 608 # == width, height
yolo.batch_size = 32
# Call make_model() after yolo.<parameter>: anchors, classes, input_size, xyscales, ...
yolo.make_model()

Load pre-trained weights

yolo.load_weights(
"/content/drive/My Drive/Hard_Soft/NN/yolov4/yolov4-tiny.conv.29",
weights_type="yolo"
)

Set callbacks and Compile

epochs = 400
lr = 1e-4
optimizer = optimizers.Adam(learning_rate=lr)
yolo.compile(optimizer=optimizer, loss_iou_type="ciou")
def lr_scheduler(epoch):
if epoch < int(epochs * 0.5):
return lr
if epoch < int(epochs * 0.8):
return lr * 0.5
if epoch < int(epochs * 0.9):
return lr * 0.1
return lr * 0.01
_callbacks = [
callbacks.LearningRateScheduler(lr_scheduler),
callbacks.TerminateOnNaN(),
callbacks.TensorBoard(
log_dir="/content/drive/My Drive/Hard_Soft/NN/logs",
),
SaveWeightsCallback(
yolo=yolo, dir_path="/content/drive/My Drive/Hard_Soft/NN/trained",
weights_type="yolo", epoch_per_save=10
),
]

Call fit()

yolo.fit(
train_data_set,
epochs=epochs,
callbacks=_callbacks,
validation_data=val_data_set,
validation_steps=50,
validation_freq=5,
steps_per_epoch=100,
)

Full script

from tensorflow.keras import callbacks, optimizers
from yolov4.tf import SaveWeightsCallback, YOLOv4
import time
yolo = YOLOv4(tiny=True)
yolo.classes = "/content/drive/My Drive/Hard_Soft/NN/coco/coco.names"
yolo.input_size = 608
yolo.batch_size = 32
yolo.make_model()
yolo.load_weights(
"/content/drive/My Drive/Hard_Soft/NN/yolov4/yolov4-tiny.conv.29",
weights_type="yolo"
)
train_data_set = yolo.load_dataset(
"/content/drive/My Drive/Hard_Soft/NN/coco/train2017.txt",
image_path_prefix="/content/train2017",
label_smoothing=0.05
)
val_data_set = yolo.load_dataset(
"/content/drive/My Drive/Hard_Soft/NN/coco/val2017.txt",
image_path_prefix="/content/val2017",
training=False
)
epochs = 400
lr = 1e-4
optimizer = optimizers.Adam(learning_rate=lr)
yolo.compile(optimizer=optimizer, loss_iou_type="ciou")
def lr_scheduler(epoch):
if epoch < int(epochs * 0.5):
return lr
if epoch < int(epochs * 0.8):
return lr * 0.5
if epoch < int(epochs * 0.9):
return lr * 0.1
return lr * 0.01
_callbacks = [
callbacks.LearningRateScheduler(lr_scheduler),
callbacks.TerminateOnNaN(),
callbacks.TensorBoard(
log_dir="/content/drive/My Drive/Hard_Soft/NN/logs",
),
SaveWeightsCallback(
yolo=yolo, dir_path="/content/drive/My Drive/Hard_Soft/NN/trained",
weights_type="yolo", epoch_per_save=10
),
]
yolo.fit(
train_data_set,
epochs=epochs,
callbacks=_callbacks,
validation_data=val_data_set,
validation_steps=50,
validation_freq=5,
steps_per_epoch=100,
)

TensorBoard

tensorboard --logdir logs

mAP

Last updated on