from absl import app, flags, logging from absl.flags import FLAGS import os import shutil import tensorflow as tf from core.yolov4 import YOLO, decode, compute_loss, decode_train from core.dataset import Dataset from core.config import cfg import numpy as np from core import utils from core.utils import freeze_all, unfreeze_all flags.DEFINE_string('model', 'yolov4', 'yolov4, yolov3') flags.DEFINE_string('weights', './scripts/yolov4.weights', 'pretrained weights') flags.DEFINE_boolean('tiny', False, 'yolo or yolo-tiny') def main(_argv): physical_devices = tf.config.experimental.list_physical_devices('GPU') if len(physical_devices) > 0: tf.config.experimental.set_memory_growth(physical_devices[0], True) trainset = Dataset(FLAGS, is_training=True) testset = Dataset(FLAGS, is_training=False) logdir = "./data/log" isfreeze = False steps_per_epoch = len(trainset) first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS global_steps = tf.Variable(1, trainable=False, dtype=tf.int64) warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch # train_steps = (first_stage_epochs + second_stage_epochs) * steps_per_period input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3]) STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS) IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH freeze_layers = utils.load_freeze_layer(FLAGS.model, FLAGS.tiny) feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny) if FLAGS.tiny: bbox_tensors = [] for i, fm in enumerate(feature_maps): if i == 0: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) else: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) else: bbox_tensors = [] for i, fm in enumerate(feature_maps): if i == 0: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) elif i == 1: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) else: bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE) bbox_tensors.append(fm) bbox_tensors.append(bbox_tensor) model = tf.keras.Model(input_layer, bbox_tensors) model.summary() if FLAGS.weights == None: print("Training from scratch") else: if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights": utils.load_weights(model, FLAGS.weights, FLAGS.model, FLAGS.tiny) else: model.load_weights(FLAGS.weights) print('Restoring weights from: %s ... ' % FLAGS.weights) optimizer = tf.keras.optimizers.Adam() if os.path.exists(logdir): shutil.rmtree(logdir) writer = tf.summary.create_file_writer(logdir) # define training step function # @tf.function def train_step(image_data, target): with tf.GradientTape() as tape: pred_result = model(image_data, training=True) giou_loss = conf_loss = prob_loss = 0 # optimizing process for i in range(len(freeze_layers)): conv, pred = pred_result[i * 2], pred_result[i * 2 + 1] loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i) giou_loss += loss_items[0] conf_loss += loss_items[1] prob_loss += loss_items[2] total_loss = giou_loss + conf_loss + prob_loss gradients = tape.gradient(total_loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) tf.print("=> STEP %4d/%4d lr: %.6f giou_loss: %4.2f conf_loss: %4.2f " "prob_loss: %4.2f total_loss: %4.2f" % (global_steps, total_steps, optimizer.lr.numpy(), giou_loss, conf_loss, prob_loss, total_loss)) # update learning rate global_steps.assign_add(1) if global_steps < warmup_steps: lr = global_steps / warmup_steps * cfg.TRAIN.LR_INIT else: lr = cfg.TRAIN.LR_END + 0.5 * (cfg.TRAIN.LR_INIT - cfg.TRAIN.LR_END) * ( (1 + tf.cos((global_steps - warmup_steps) / (total_steps - warmup_steps) * np.pi)) ) optimizer.lr.assign(lr.numpy()) # writing summary data with writer.as_default(): tf.summary.scalar("lr", optimizer.lr, step=global_steps) tf.summary.scalar("loss/total_loss", total_loss, step=global_steps) tf.summary.scalar("loss/giou_loss", giou_loss, step=global_steps) tf.summary.scalar("loss/conf_loss", conf_loss, step=global_steps) tf.summary.scalar("loss/prob_loss", prob_loss, step=global_steps) writer.flush() def test_step(image_data, target): with tf.GradientTape() as tape: pred_result = model(image_data, training=True) giou_loss = conf_loss = prob_loss = 0 # optimizing process for i in range(len(freeze_layers)): conv, pred = pred_result[i * 2], pred_result[i * 2 + 1] loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i) giou_loss += loss_items[0] conf_loss += loss_items[1] prob_loss += loss_items[2] total_loss = giou_loss + conf_loss + prob_loss tf.print("=> TEST STEP %4d giou_loss: %4.2f conf_loss: %4.2f " "prob_loss: %4.2f total_loss: %4.2f" % (global_steps, giou_loss, conf_loss, prob_loss, total_loss)) for epoch in range(first_stage_epochs + second_stage_epochs): if epoch < first_stage_epochs: if not isfreeze: isfreeze = True for name in freeze_layers: freeze = model.get_layer(name) freeze_all(freeze) elif epoch >= first_stage_epochs: if isfreeze: isfreeze = False for name in freeze_layers: freeze = model.get_layer(name) unfreeze_all(freeze) for image_data, target in trainset: train_step(image_data, target) for image_data, target in testset: test_step(image_data, target) model.save_weights("./checkpoints/yolov4") if __name__ == '__main__': try: app.run(main) except SystemExit: pass