main.py

import warnings
import numpy as np
import torch
from absl import app, flags
from pytorch_lightning import Trainer, seed_everything
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.loggers import CSVLogger, TensorBoardLogger, wandb

from fairdata import DualLoader, FairDataModule
from fairfacedata import FairFaceDataModule  
from utkfacedata import UTKFaceDataModule, DualLoader as UTKDualLoader   
from FBADSTrainer import FairBADSModule
from utils import compute_particle_diversity

warnings.filterwarnings("ignore")

GPU = 0
FLAGS = flags.FLAGS

flags.DEFINE_string("dataset_name", "lfwa_w", "Dataset name")

flags.DEFINE_bool("meta", True, "Whether to use meta-training")
flags.DEFINE_bool("gpu", True, "Whether to use gpu")
flags.DEFINE_float("eta", 1e-4, "Learning rate for SVGD")
flags.DEFINE_float("beta", 0.005, "KL weight")
flags.DEFINE_float("meta_ratio", 0.001, "meta ratio, suggest 0.001 for celeba/fairface/utkface")
flags.DEFINE_float("bias_amount", 0.4, "bias added in the train set")
flags.DEFINE_integer("max_epochs", 10, "Number of training epochs")
flags.DEFINE_integer("batch_size", 64, "Batch size")
flags.DEFINE_integer("n_particles", 10, "Number of particles")
flags.DEFINE_integer("num_workers", 4, "num of workers in data module")

# MMD 
flags.DEFINE_string("barycenter_method", "wasserstein", "Barycenter computation method: wasserstein, mmd, multi_mmd, js, or multi_js")
flags.DEFINE_string("mmd_kernel", "gaussian", "MMD kernel type: gaussian, polynomial, or laplacian")
flags.DEFINE_float("mmd_sigma", -1.0, "MMD Gaussian kernel sigma (-1 for adaptive)")
flags.DEFINE_integer("mmd_degree", 2, "MMD polynomial kernel degree")
flags.DEFINE_float("mmd_gamma", 1.0, "MMD polynomial kernel gamma")
flags.DEFINE_float("mmd_coef0", 1.0, "MMD polynomial kernel coef0")
flags.DEFINE_float("mmd_lr", 0.01, "MMD barycenter optimization learning rate")

# JS 
flags.DEFINE_float("js_lr", 0.01, "JS divergence barycenter optimization learning rate")
flags.DEFINE_integer("js_numItermax", 100, "JS divergence maximum iterations")


def main(argv):

    seeds = [1975, 831, 422]
    all_results = []

    for s in seeds:
        seed_everything(s)
        torch.set_float32_matmul_precision("high")
        accelerator = "gpu"
        device = torch.cuda.device_count() if accelerator == "gpu" else 1
        print("I use",accelerator)
        dataset_name = FLAGS.dataset_name
  
        if dataset_name == "lfwa_w":
            data_location = "./dataset"
            dm = FairDataModule(data_location,
                                dataset_name,
                                meta_ratio = FLAGS.meta_ratio,
                                bias_amount=FLAGS.bias_amount,
                                batch_size=FLAGS.batch_size,
                                num_workers = FLAGS.num_workers)
            
        elif dataset_name == "fairface":
            data_location = "./fairface"
            dm = FairFaceDataModule(data_location,
                                    batch_size=FLAGS.batch_size,
                                    bias_amount=FLAGS.bias_amount,
                                    meta_ratio=FLAGS.meta_ratio,
                                    num_workers=FLAGS.num_workers)
            
        elif dataset_name == "utkface":
            data_location = "./utkface"  
            dm = UTKFaceDataModule(data_location,
                                batch_size=FLAGS.batch_size,
                                bias_amount=FLAGS.bias_amount,
                                meta_ratio=FLAGS.meta_ratio,
                                num_workers=FLAGS.num_workers,
                                debug=False) 
            
        else:
            raise NotImplementedError

        dm.prepare_data()
        dm.setup()

        train_loader = dm.train_dataloader()
        meta_set = dm.meta_dataset
        test_loader = dm.test_dataloader()
   
        if dataset_name == "lfwa_w":
            dual_loader = DualLoader(train_loader, meta_set)
        elif dataset_name == "fairface":
            dual_loader = DualLoader(train_loader, meta_set)
        elif dataset_name == "utkface":
            dual_loader = UTKDualLoader(train_loader, meta_set)
        else:
            raise NotImplementedError
  
        Ntr = len(dm.train_dataset)
        Nmeta = len(meta_set)
        Nts = len(dm.test_dataset) if hasattr(dm, "test_dataset") else 0
        N = Ntr+Nmeta+Nts

        # This is for MMD
        mmd_kernel_params = {}
        if FLAGS.mmd_kernel == "gaussian":
            if FLAGS.mmd_sigma > 0:
                mmd_kernel_params['sigma'] = FLAGS.mmd_sigma
        elif FLAGS.mmd_kernel == "polynomial":
            mmd_kernel_params['degree'] = FLAGS.mmd_degree
            mmd_kernel_params['gamma'] = FLAGS.mmd_gamma
            mmd_kernel_params['coef0'] = FLAGS.mmd_coef0
        elif FLAGS.mmd_kernel == "laplacian":
            if FLAGS.mmd_sigma > 0:
                mmd_kernel_params['sigma'] = FLAGS.mmd_sigma
            else:
                mmd_kernel_params['sigma'] = 1.0

        if FLAGS.barycenter_method in ["mmd", "multi_mmd"]:
            mmd_kernel_params['lr'] = FLAGS.mmd_lr

        model = FairBADSModule(
            n_particles=FLAGS.n_particles,
            eta=FLAGS.eta,
            beta=FLAGS.beta,
            n_train=N,
            n_meta=Nmeta,
            global_group_distribution=dm.global_group_distribution,
            meta=FLAGS.meta,
            dataset_name=FLAGS.dataset_name,
            barycenter_method=FLAGS.barycenter_method,
            mmd_kernel_type=FLAGS.mmd_kernel,
            mmd_kernel_params=mmd_kernel_params,
            js_lr=FLAGS.js_lr,
            js_numItermax=FLAGS.js_numItermax
        )
        logger = TensorBoardLogger(save_dir="logs", name="bads_run")
        checkpoint_callback = ModelCheckpoint(
            monitor="val/accuracy",
            save_top_k=1,
            mode="max",
            filename="best-val-acc"
        )

        trainer = Trainer(
            max_epochs=FLAGS.max_epochs,
            logger=logger,
            callbacks=[checkpoint_callback],
            log_every_n_steps=10,
            accelerator=accelerator,
            devices=device,
            num_sanity_val_steps=0 
        )

        trainer.fit(model, dual_loader, dm.val_dataloader())
        trainer.test(model, dataloaders=test_loader)
        all_results.append(model.test_result_dict)

    all_metrics = {
    k: np.array([
        r[k].cpu().item() if torch.is_tensor(r[k]) else r[k] 
        for r in all_results
    ]) 
    for k in all_results[0].keys()
}

    print("=== Test Summary (mean ± std) ===")
    for key, values in all_metrics.items():
        mean = values.mean()
        std = values.std()
        print(f"{key}: {mean:.4f} ± {std:.4f}")

    with open("output_fairbads.txt", "w") as f:
        f.write("=== Test Summary (mean ± std) ===\n")
        for key, values in all_metrics.items():
            mean = values.mean()
            std = values.std()
            f.write(f"{key}: {mean:.4f} ± {std:.4f}\n")


if __name__ == "__main__":
    app.run(main)