Skip to main content

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 YOLODataset after yolo.config.parse_cfg():
train_dataset = YOLODataset(
config=yolo.config,
dataset_list="/content/drive/My Drive/Hard_Soft/NN/coco/train2017.txt",
image_path_prefix="/content/train2017",
training=True,
)

Set yolo parameters#

Use <yolo-model>.cfg(darknet).

yolo = YOLOv4()
yolo.config.parse_names("coco.names")
yolo.config.parse_cfg("yolov4-tiny.cfg")

Load pre-trained weights#

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

Freeze some layers#

# Tested on v3 only
for i in range(29):
yolo.model.get_layer(index=i).trainable = False

Set callbacks and Compile#

Use <yolo-model>.cfg(darknet).

yolo.compile()
_callbacks = [
callbacks.TerminateOnNaN(),
callbacks.TensorBoard(
log_dir="/content/drive/MyDrive/Hard_Soft/NN/logs",
update_freq=200,
histogram_freq=1,
),
SaveWeightsCallback(
yolo=yolo,
dir_path="/content/drive/MyDrive/Hard_Soft/NN/trained",
weights_type="yolo",
step_per_save=2000,
),
]

Call fit()#

yolo.fit(
train_dataset,
callbacks=_callbacks,
validation_data=val_dataset,
verbose=3, # 3: print step info
)

Full script#

  • 2021-02-13
  • OS: Ubuntu 20.04
  • TF: v2.4.1
  • yolov4: v3.0.0
from tensorflow.keras import callbacks
from yolov4.tf import YOLOv4, YOLODataset, SaveWeightsCallback
yolo = YOLOv4()
yolo.config.parse_names("coco.names")
yolo.config.parse_cfg("yolov4-tiny.cfg")
yolo.make_model()
yolo.load_weights(
"/content/drive/MyDrive/Hard_Soft/NN/yolov4/yolov4-tiny.conv.29",
weights_type="yolo",
)
yolo.summary(summary_type="yolo")
for i in range(29):
yolo.model.get_layer(index=i).trainable = False
yolo.summary()
train_dataset = YOLODataset(
config=yolo.config,
dataset_list="/content/drive/MyDrive/Hard_Soft/NN/coco/train2017.txt",
image_path_prefix="/content/train2017",
training=True,
)
val_dataset = YOLODataset(
config=yolo.config,
dataset_list="/content/drive/MyDrive/Hard_Soft/NN/coco/val2017.txt",
image_path_prefix="/content/val2017",
training=False,
)
yolo.compile()
_callbacks = [
callbacks.TerminateOnNaN(),
callbacks.TensorBoard(
log_dir="/content/drive/MyDrive/Hard_Soft/NN/logs",
update_freq=200,
histogram_freq=1,
),
SaveWeightsCallback(
yolo=yolo,
dir_path="/content/drive/MyDrive/Hard_Soft/NN/trained",
weights_type="yolo",
step_per_save=2000,
),
]
yolo.fit(
train_dataset,
callbacks=_callbacks,
validation_data=val_dataset,
verbose=3, # 3: print step info
)

TensorBoard#

tensorboard --logdir logs
Last updated on