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