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()
