clean up model

[hhk7734]

class Decode(Model):
    def __init__(self, anchors_ratio, num_classes: int, xyscale):
        super(Decode, self).__init__()
        self.anchors_ratio = anchors_ratio
        self.cell_ratio = None
        self.num_classes = num_classes
        self.xyscale = xyscale
        self.xy_grid = None

        self.reshape0 = layers.Reshape((-1,))
        self.concatenate = layers.Concatenate(axis=-1)

    def build(self, input_shape):
        self.cell_ratio = 1 / input_shape[1]

        self.reshape0.target_shape = (
            input_shape[1],
            input_shape[1],
            3,
            5 + self.num_classes,
        )

        # grid(1, i, j, 3, 2) => grid top left coordinates
        # [
        #     [ [[0, 0]], [[1, 0]], [[2, 0]], ...],
        #     [ [[0, 1]], [[1, 1]], [[2, 1]], ...],
        # ]
        self.xy_grid = tf.stack(
            tf.meshgrid(tf.range(input_shape[1]), tf.range(input_shape[1])),
            axis=-1,
        )  # size i, j, 2
        self.xy_grid = tf.reshape(
            self.xy_grid, (1, input_shape[1], input_shape[1], 1, 2)
        )
        self.xy_grid = tf.tile(self.xy_grid, [1, 1, 1, 3, 1])
        self.xy_grid = tf.cast(self.xy_grid, tf.float32)

    def call(self, x, training: bool = False):
        """
        @param x: Dim(batch, grid_y, grid_x,
                        anchors * (tx, ty, tw, th, score, classes))

        @return
            Dim(batch, grid_y, gird_x, anchors, (x, y, w, h, score, classes ))
        """
        x = self.reshape0(x)
        txty, twth, score, classes = tf.split(
            x, (2, 2, 1, self.num_classes), axis=-1
        )

        # x = (f(tx) + left_x) * strides / input_size
        # y = (f(ty) + top_y) * strides / input_size
        txty = (activations.sigmoid(txty) - 0.5) * self.xyscale + 0.5
        xy = (txty + self.xy_grid) * self.cell_ratio

        # w = (anchor_w * exp(tw)) / input_size
        # h = (anchor_h * exp(th)) / input_size
        wh = self.anchors_ratio * backend.exp(twth)

        if not training:
            score = activations.sigmoid(score)
            classes = activations.sigmoid(classes)

        x = self.concatenate([xy, wh, score, classes])
        return x

[hhk7734]

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

[hhk7734]

close b477be3e