train.py

import os
import tensorflow as tf
from dataset import MusnadDataset
from model import build_musnad_detector
from loss import obj_loss_fn, box_loss_fn, cls_loss_fn

BATCH_SIZE = 32
EPOCHS = 20 
LEARNING_RATE = 1e-4

def train():
    print("Loading Optimized Multi-Output Dataset...")
    train_ds = MusnadDataset("dataset", split="train", batch_size=BATCH_SIZE).get_dataset()
    val_ds = MusnadDataset("dataset", split="val", batch_size=BATCH_SIZE).get_dataset()

    print("Building Refined Multi-Head Architecture...")
    model = build_musnad_detector()
    
    optimizer = tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE)

    # Compile with separate losses, weights and accuracy reporting
    model.compile(
        optimizer=optimizer,
        loss={
            'obj_output': obj_loss_fn,
            'box_output': box_loss_fn,
            'cls_output': cls_loss_fn
        },
        loss_weights={
            'obj_output': 1.0,
            'box_output': 5.0, 
            'cls_output': 1.0
        },
        metrics={
            'cls_output': 'accuracy' 
        }
    )

    checkpoint_path = "musnad_detector_multi_head.keras"
    callbacks = [
        tf.keras.callbacks.ModelCheckpoint(checkpoint_path, monitor='val_loss', save_best_only=True, verbose=1),
        tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=4, min_lr=1e-7, verbose=1),
        tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=12, restore_best_weights=True, verbose=1)
    ]

    print("Starting Training (Each loss is visible independently)...")
    model.fit(
        train_ds,
        validation_data=val_ds,
        epochs=EPOCHS,
        callbacks=callbacks
    )

    print(f"Success! Model saved to: {checkpoint_path}")

if __name__ == "__main__":
    train()