def train_step(image_data, target): with tf.GradientTape() as tape: pred_result = yolo(image_data, training=True) giou_loss = conf_loss = prob_loss = 0 # optimizing process grid = 3 if not TRAIN_YOLO_TINY else 2 for i in range(grid): conv, pred = pred_result[i * 2], pred_result[i * 2 + 1] loss_items = compute_loss(pred, conv, *target[i], i, CLASSES=TRAIN_CLASSES) 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, yolo.trainable_variables) optimizer.apply_gradients(zip(gradients, yolo.trainable_variables)) # update learning rate # about warmup: https://arxiv.org/pdf/1812.01187.pdf&usg=ALkJrhglKOPDjNt6SHGbphTHyMcT0cuMJg global_steps.assign_add(1) if global_steps < warmup_steps: # and not TRAIN_TRANSFER: lr = global_steps / warmup_steps * TRAIN_LR_INIT else: lr = TRAIN_LR_END + 0.5 * (TRAIN_LR_INIT - 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() return global_steps.numpy(), optimizer.lr.numpy(), giou_loss.numpy( ), conf_loss.numpy(), prob_loss.numpy(), total_loss.numpy()
def validate_step(image_data, target): with tf.GradientTape() as tape: pred_result = yolo(image_data, training=False) giou_loss=conf_loss=prob_loss=0 # optimizing process grid = 3 if not TRAIN_YOLO_TINY else 2 for i in range(grid): conv, pred = pred_result[i*2], pred_result[i*2+1] loss_items = compute_loss(pred, conv, *target[i], i, CLASSES=TRAIN_CLASSES) giou_loss += loss_items[0] conf_loss += loss_items[1] prob_loss += loss_items[2] total_loss = giou_loss + conf_loss + prob_loss return giou_loss.numpy(), conf_loss.numpy(), prob_loss.numpy(), total_loss.numpy()