Training custom model about keras_cv_attention_models HOT 10 CLOSED

import tensorflow as tf
from tensorflow import keras
from keras_cv_attention_models import yolov8, hornet
from keras_cv_attention_models.coco import init_dataset, eval_func, losses
from keras_cv_attention_models.imagenet import train, init_global_strategy, init_lr_scheduler

input_shape = (320, 320, 3)
num_classes = 2  # Specify the actual num_classes
batch_size = 4
anchors_mode = "anchor_free"
anchor_pyramid_levels = [3, 5]
rescale_mode = "raw01"
lr_base = 1e-3
lr_decay_steps = 30
data_name = "coco/2017"  # Or custom one "datasets/coco_dog_cat/detections.json"

strategy = init_global_strategy(enable_float16=True)

with strategy.scope():
    # Build own model
    backbone = hornet.HorNetTiny(num_classes=0, input_shape=input_shape)
    model = yolov8.YOLOV8_S(
        backbone=backbone, input_shape=input_shape, anchors_mode='anchor_free', regression_len=4, pretrained=None, num_classes=num_classes
    )
    # Add another output
    another_output = keras.layers.GlobalAveragePooling2D(name="another")(model.inputs[0])
    model = keras.models.Model(model.inputs[0], model.outputs + [another_output])

    train_dataset = init_dataset(
        data_name=data_name, input_shape=input_shape, rescale_mode=rescale_mode, batch_size=batch_size, anchors_mode=anchors_mode, anchor_pyramid_levels=anchor_pyramid_levels,
    )[0]
    # Add another output for dataset
    train_dataset = train_dataset.map(lambda xx, yy: (xx, (yy, tf.random.uniform([batch_size, 3]))))

    lr_scheduler, lr_total_epochs = init_lr_scheduler(lr_base, lr_decay_steps)
    loss = losses.AnchorFreeLoss(input_shape, pyramid_levels=anchor_pyramid_levels)
    metrics = losses.ClassAccuracyWithBboxWrapper(loss)
    optimizer = keras.optimizers.AdamW(learning_rate=lr_base, weight_decay=0.002)

    # Add another loss and metrics
    another_loss = keras.losses.MeanSquaredError()
    model.compile(optimizer, loss=[loss, another_loss], metrics={"outputs_fp32": metrics, 'another': None})
    coco_ap_eval = eval_func.COCOEvalCallback(data_name, rescale_mode=rescale_mode, batch_size=batch_size)

    other_fit_kwargs = {}  # Set custom ones for `model.fit`
    latest_save, hist = train(
        compiled_model=model,
        epochs=lr_total_epochs,
        train_dataset=train_dataset,
        test_dataset=None,  # COCO eval using coco_ap_eval callback, set `validation_data` for `model.fit` to None
        initial_epoch=0,
        lr_scheduler=lr_scheduler,
        basic_save_name="coco_test",
        init_callbacks=[coco_ap_eval],
        logs=None,
        **other_fit_kwargs,
    )

train_dataset = init_dataset(...)[0]
another_dataset = init_dataset(...)[0]
# Zip them
tt = tf.data.Dataset.zip((train_dataset, another_dataset))

(aa, bb), (cc, dd) = tt.as_numpy_iterator().next()
print(aa.shape, bb.shape, cc.shape, dd.shape)
# (4, 320, 320, 3) (4, 320, 320, 3) (4, 100, 7) (4, 100, 7)

# Map and rearrange, here keeps only the first lables output, but repeat twice
tt = tt.map(lambda one, another: ((one[0], another[0]), (one[1], one[1])))
(images, another_images), (labels, another_labels) = tt.as_numpy_iterator().next()
print(images.shape, another_images.shape, labels.shape, another_labels.shape)
# (4, 320, 320, 3) (4, 320, 320, 3) (4, 100, 7) (4, 100, 7)

another_input = keras.layers.Input(model.input_shape[1:])
another_output = model(another_input)
model = keras.models.Model([model.inputs[0], another_input], [model.outputs[0], another_output])
...
model.compile(optimizer, loss=[loss, another_loss], metrics={model.output_names[0]: metrics, model.output_names[1]: None})
latest_save, hist = train(compiled_model=model, train_dataset=train_dataset, ...)

def on_epoch_end(self, epoch=0, logs=None):
    if not self.built:
        if self.dataset_kwargs["rescale_mode"] == "auto":
            self.dataset_kwargs["rescale_mode"] = getattr(self.model, "rescale_mode", "torch")
        self.model = models.Model(self.model.inputs[0], self.model.outputs[-1])  # <- rebuild model
        self.build(self.model.input_shape, self.model.output_shape)

from keras_cv_attention_models.coco import anchors_func
print(anchors_func.get_anchor_free_anchors(input_shape=(512, 512, 3)).shape)
# (5376, 4)