Closed hhk7734 closed 4 years ago
YOLOv4
def compile(self, iou_type: str = "giou", learning_rate: float = 1e-5):
self.optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate,)
self._compiled_loss = train.make_compiled_loss(self, iou_type=iou_type)
def train_step(self, data):
"""
@param data: (x, y) => (input, ground_truth)
ground_truth == (s_truth, m_truth, l_truth)
"""
images, ground_truth = data
num_bboxes = [0, 0, 0]
for i in range(3):
num_bboxes[i] = tf.reduce_sum(ground_truth[i][..., 4])
with tf.GradientTape() as tape:
# Forward
y_pred = self(images, training=True)
xiou_loss, classes_loss = self._compiled_loss(ground_truth, y_pred)
total_loss = xiou_loss + classes_loss
loss = tf.reduce_sum(total_loss)
# Compute gradients, Backward
trainable_vars = self.trainable_variables
gradients = tape.gradient(loss, trainable_vars)
# Update weights
self.optimizer.apply_gradients(zip(gradients, trainable_vars))
return num_bboxes, classes_loss, xiou_loss, total_loss, loss