train.py

import os
import random
import toml
from sys import argv
from types import SimpleNamespace

import accelerate
from tqdm import tqdm
import wandb

import numpy as np
import torch
from torch.optim.lr_scheduler import LambdaLR

from translators.Discriminator import Discriminator

# from eval import eval_model
from utils.collate import MultiencoderTokenizedDataset, TokenizedCollator
from utils.eval_utils import EarlyStopper, eval_loop_
from utils.gan import LeastSquaresGAN, RelativisticGAN, VanillaGAN
from utils.model_utils import get_sentence_embedding_dimension, load_encoder
from utils.muon import MuonW
from utils.utils import *
from utils.streaming_utils import load_streaming_embeddings, process_batch
from utils.train_utils import rec_loss_fn, trans_loss_fn, vsp_loss_fn, get_grad_norm
from utils.wandb_logger import Logger

from datasets import load_from_disk

from fmri.data import FmriSplitConfig, load_fmri_train_and_eval_datasets


def _split_muon_params(model: torch.nn.Module):
    matrix_params = []
    scalar_vector_params = []
    for _, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if p.ndim >= 2:
            matrix_params.append(p)
        else:
            scalar_vector_params.append(p)
    return matrix_params, scalar_vector_params


def build_generator_optimizers(cfg, translator):
    opt_name = str(getattr(cfg, "optimizer", "adam")).lower()
    if opt_name == "adam":
        return [torch.optim.Adam(translator.parameters(), lr=cfg.lr, fused=False, betas=(0.5, 0.999))]
    if opt_name == "adamw":
        return [torch.optim.AdamW(
            translator.parameters(),
            lr=cfg.lr,
            betas=tuple(getattr(cfg, "adamw_betas", (0.9, 0.999))),
            eps=float(getattr(cfg, "adamw_eps", 1e-8)),
            weight_decay=float(getattr(cfg, "weight_decay", 0.01)),
        )]
    if opt_name in ("muon", "muonw"):
        matrix_params, scalar_vector_params = _split_muon_params(translator)
        optimizers = []
        if len(matrix_params) > 0:
            if hasattr(torch.optim, "Muon"):
                optimizers.append(
                    torch.optim.Muon(
                        matrix_params,
                        lr=cfg.lr,
                        momentum=float(getattr(cfg, "muon_beta", 0.95)),
                        weight_decay=float(getattr(cfg, "weight_decay", 0.01)),
                        ns_steps=int(getattr(cfg, "muon_ns_steps", 5)),
                    )
                )
            else:
                optimizers.append(
                    MuonW(
                        matrix_params,
                        lr=cfg.lr,
                        weight_decay=float(getattr(cfg, "weight_decay", 0.01)),
                        muon_beta=float(getattr(cfg, "muon_beta", 0.95)),
                        adamw_betas=tuple(getattr(cfg, "adamw_betas", (0.9, 0.95))),
                        adamw_eps=float(getattr(cfg, "adamw_eps", 1e-8)),
                        ns_steps=int(getattr(cfg, "muon_ns_steps", 5)),
                    )
                )
        if len(scalar_vector_params) > 0:
            # Best practice with Muon: keep biases/norm/scalars on AdamW (usually no WD).
            optimizers.append(
                torch.optim.AdamW(
                    scalar_vector_params,
                    lr=float(getattr(cfg, "muon_adamw_lr", cfg.lr)),
                    betas=tuple(getattr(cfg, "adamw_betas", (0.9, 0.95))),
                    eps=float(getattr(cfg, "adamw_eps", 1e-8)),
                    weight_decay=float(getattr(cfg, "muon_adamw_weight_decay", 0.0)),
                )
            )
        return optimizers
    raise ValueError(f"Unknown optimizer: {opt_name}. Expected one of: adam, adamw, muonw")


def training_loop_(
    save_dir, accelerator, gan, sup_gan, latent_gan, similarity_gan, translator, sup_dataloader, sup_iter, unsup_dataloader, sup_encs, unsup_enc, cfg, gen_optimizers, gen_schedulers, logger=None, max_num_batches=None
):
    device = accelerator.device
    if logger is None:
        logger = Logger(dummy=True)

    # wandb.watch(translator, log='all')

    if sup_iter is not None:
        dataloader_pbar = unsup_dataloader
    else:
        dataloader_pbar = zip(sup_dataloader, unsup_dataloader)


    dataloader_pbar = tqdm(dataloader_pbar, total=len(unsup_dataloader), desc="Training")

    model_save_dir = os.path.join(save_dir, 'model.pt')

    translator.train()
    for i, batches in enumerate(dataloader_pbar):
        if sup_iter is not None:
            try:
                sup_batch = next(sup_iter)
            except StopIteration:
                print('Restarting sup_dataloader...')
                sup_iter = iter(sup_dataloader)
                sup_batch = next(sup_iter)
            unsup_batch = batches
        else:
            sup_batch, unsup_batch = batches

        if max_num_batches is not None and i >= max_num_batches:
            print(f"Early stopping at {i} batches")
            break
        with accelerator.accumulate(translator), accelerator.autocast():
            assert len(set(sup_batch.keys()).intersection(unsup_batch.keys())) == 0
            if str(cfg.dataset).startswith("fmri"):
                ins = {}
                for k, v in {**sup_batch, **unsup_batch}.items():
                    if not torch.is_tensor(v):
                        raise ValueError(f"Expected tensor batch value for key={k}")
                    vv = v.to(device=device, dtype=torch.float32)
                    if cfg.normalize_embeddings:
                        vv = vv / (vv.norm(p=2, dim=1, keepdim=True) + 1e-12)
                    ins[k] = vv
            else:
                ins = {
                    **process_batch(sup_batch, sup_encs, cfg.normalize_embeddings, device),
                    **process_batch(unsup_batch, unsup_enc, cfg.normalize_embeddings, device),
                }

            recons, translations, reps = translator(
                ins, noise_level=cfg.noise_level, include_reps=True
            )

            # discriminator
            disc_r1_penalty, disc_loss, gen_loss, disc_acc_real, disc_acc_fake, gen_acc = gan.step(
                real_data=ins[cfg.unsup_emb] + torch.randn_like(ins[cfg.unsup_emb], device=ins[cfg.unsup_emb].device) * cfg.noise_level,
                fake_data=translations[cfg.unsup_emb][cfg.sup_emb] + torch.randn_like(translations[cfg.unsup_emb][cfg.sup_emb], device=translations[cfg.unsup_emb][cfg.sup_emb].device) * cfg.noise_level
            )

            sup_disc_r1_penalty, sup_disc_loss, sup_gen_loss, sup_disc_acc_real, sup_disc_acc_fake, sup_gen_acc = sup_gan.step(
                real_data=ins[cfg.sup_emb] + torch.randn_like(ins[cfg.sup_emb], device=ins[cfg.sup_emb].device) * cfg.noise_level,
                fake_data=translations[cfg.sup_emb][cfg.unsup_emb] + torch.randn_like(translations[cfg.sup_emb][cfg.unsup_emb], device=translations[cfg.sup_emb][cfg.unsup_emb].device) * cfg.noise_level,
            )

            # latent discriminator
            latent_disc_r1_penalty, latent_disc_loss, latent_gen_loss, latent_disc_acc_real, latent_disc_acc_fake, latent_gen_acc = latent_gan.step(
                real_data=reps[cfg.sup_emb],
                fake_data=reps[cfg.unsup_emb]
            )

            # similarity discriminator
            if cfg.loss_coefficient_similarity_gen > 0:
                real_sims_A = ins[cfg.sup_emb] @ ins[cfg.sup_emb].T
                fake_sims_A = (
                    translations[cfg.sup_emb][cfg.unsup_emb] @ translations[cfg.sup_emb][cfg.unsup_emb].T
                )
                real_sims_B = ins[cfg.unsup_emb] @ ins[cfg.unsup_emb].T
                fake_sims_B = (
                    translations[cfg.unsup_emb][cfg.sup_emb] @ translations[cfg.unsup_emb][cfg.sup_emb].T
                )
                similarity_r1_penalty, similarity_disc_loss, similarity_gen_loss, similarity_disc_acc_real, similarity_disc_acc_fake, similarity_gen_acc = similarity_gan.step(
                    real_data=torch.cat([real_sims_A, real_sims_B], dim=1),
                    fake_data=torch.cat([fake_sims_A, fake_sims_B], dim=1)
                )
            else:
                similarity_r1_penalty = torch.tensor(0.0)
                similarity_disc_loss = torch.tensor(0.0)
                similarity_gen_loss = torch.tensor(0.0)
                similarity_disc_acc_real = 0.0
                similarity_disc_acc_fake = 0.0
                similarity_gen_acc = 0.0
            
            rec_loss = rec_loss_fn(ins, recons, logger)
            ins_reversed = {
                cfg.sup_emb: ins[cfg.unsup_emb],
                cfg.unsup_emb: ins[cfg.sup_emb],
            }
            translations_as_recons = {
                cfg.sup_emb: translations[cfg.unsup_emb][cfg.sup_emb],
                cfg.unsup_emb: translations[cfg.sup_emb][cfg.unsup_emb],
            }
            reverse_rec_loss = rec_loss_fn(ins_reversed, translations_as_recons, logger, prefix="reverse_")

            recons_as_translations = { 
                in_name: { in_name: val } for in_name, val in recons.items() 
            }
            vsp_loss = vsp_loss_fn(ins, recons_as_translations, logger)
            if (cfg.loss_coefficient_cc_rec > 0) or (cfg.loss_coefficient_cc_trans > 0):
                cc_ins = {}
                for out_flag in translations.keys():
                    in_flag = random.choice(list(translations[out_flag].keys()))
                    cc_ins[out_flag] = translations[out_flag][in_flag].detach()
                cc_recons, cc_translations = translator(cc_ins)
                cc_rec_loss = rec_loss_fn(ins, cc_recons, logger, prefix="cc_")
                cc_trans_loss = trans_loss_fn(ins, cc_translations, logger, prefix="cc_")
                cc_vsp_loss = vsp_loss_fn(ins, cc_translations, logger)
            else:
                cc_rec_loss = torch.tensor(0.0)
                cc_trans_loss = torch.tensor(0.0)
                cc_vsp_loss = torch.tensor(0.0)

            loss = (
                + (rec_loss * cfg.loss_coefficient_rec)
                + (reverse_rec_loss * cfg.loss_coefficient_reverse_rec)
                + (vsp_loss * cfg.loss_coefficient_vsp)
                + (cc_vsp_loss * cfg.loss_coefficient_cc_vsp)
                + (cc_rec_loss * cfg.loss_coefficient_cc_rec)
                + (cc_trans_loss * cfg.loss_coefficient_cc_trans)
                + (gen_loss * cfg.loss_coefficient_gen)
                + (sup_gen_loss * cfg.loss_coefficient_gen)
                + (latent_gen_loss * cfg.loss_coefficient_latent_gen)
                + (similarity_gen_loss * cfg.loss_coefficient_similarity_gen)
            )
            exit_on_nan(loss)
            for optimizer in gen_optimizers:
                optimizer.zero_grad()
            accelerator.backward(loss)
            accelerator.clip_grad_norm_(translator.parameters(), cfg.max_grad_norm)
            grad_norm_generator = get_grad_norm(translator)
            grad_norm_discriminator = get_grad_norm(gan.discriminator)
            grad_norm_sup_discriminator = get_grad_norm(sup_gan.discriminator)
            grad_norm_latent_discriminator = get_grad_norm(latent_gan.discriminator)
            grad_norm_similarity_discriminator = get_grad_norm(similarity_gan.discriminator)

            for optimizer in gen_optimizers:
                optimizer.step()
            for scheduler in gen_schedulers:
                scheduler.step()

            metrics = {
                "disc_loss": disc_loss.item(),
                "disc_r1_penalty": disc_r1_penalty.item(),
                "sup_disc_loss": sup_disc_loss.item(),
                "sup_disc_r1_penalty": sup_disc_r1_penalty.item(),
                "latent_disc_loss": latent_disc_loss.item(),
                "latent_disc_r1_penalty": latent_disc_r1_penalty.item(),
                "similarity_disc_loss": similarity_disc_loss.item(),
                "similarity_r1_penalty": similarity_r1_penalty.item(),
                "rec_loss": rec_loss.item(),
                "reverse_rec_loss": reverse_rec_loss.item(),
                "vsp_loss": vsp_loss.item(),
                "cc_vsp_loss": cc_vsp_loss.item(),
                "cc_rec_loss": cc_rec_loss.item(),
                "cc_trans_loss": cc_trans_loss.item(),
                "gen_loss": gen_loss.item(),
                "sup_gen_loss": sup_gen_loss.item(),
                "latent_gen_loss": latent_gen_loss.item(),
                "similarity_gen_loss": similarity_gen_loss.item(),
                "loss": loss.item(),
                "grad_norm_generator": grad_norm_generator,
                "grad_norm_discriminator": grad_norm_discriminator,
                "grad_norm_sup_discriminator": grad_norm_sup_discriminator,
                "grad_norm_latent_discriminator": grad_norm_latent_discriminator,
                "grad_norm_similarity_discriminator": grad_norm_similarity_discriminator,
                "learning_rate": gen_optimizers[0].param_groups[0]["lr"],
                "disc_acc_real": disc_acc_real,
                "disc_acc_fake": disc_acc_fake,
                "latent_disc_acc_real": latent_disc_acc_real,
                "latent_disc_acc_fake": latent_disc_acc_fake,
                "gen_acc": gen_acc,
                "sup_disc_acc_real": sup_disc_acc_real,
                "sup_disc_acc_fake": sup_disc_acc_fake,
                "sup_gen_acc": sup_gen_acc,
                "similarity_disc_acc_real": similarity_disc_acc_real,
                "similarity_disc_acc_fake": similarity_disc_acc_fake,
                "similarity_gen_acc": similarity_gen_acc,
            }

            for metric, value in metrics.items():
                logger.logkv(metric, value)
            logger.dumpkvs(force=(hasattr(cfg, 'force_dump') and cfg.force_dump))
            dataloader_pbar.set_postfix(metrics)

    with open(save_dir + 'config.toml', 'w') as f:
        toml.dump(cfg.__dict__, f)
    torch.save(accelerator.unwrap_model(translator).state_dict(), model_save_dir)
    return sup_iter


def main():
    os.environ["TOKENIZERS_PARALLELISM"] = "0"
    cfg = toml.load(f'configs/{argv[1]}.toml')
    unknown_cfg = read_args(argv)
    cfg = SimpleNamespace(**{**{k: v for d in cfg.values() for k, v in d.items()}, **unknown_cfg})

    if hasattr(cfg, 'mixed_precision') and cfg.mixed_precision != 'no' and cfg.mixed_precision == 'bf16' and not torch.cuda.is_bf16_supported():
        cfg.mixed_precision = 'fp16'
        cfg.gradient_accumulation_steps = 1
        print("Note: bf16 is not available on this hardware! Reverting to fp16 and setting accumulation steps to 1.")

    # set seeds
    random.seed(cfg.seed)
    torch.manual_seed(cfg.seed)
    np.random.seed(cfg.seed)
    torch.cuda.manual_seed(cfg.seed)

    fmri_mode = str(cfg.dataset).startswith("fmri")
    use_val_set = fmri_mode or hasattr(cfg, 'val_size')

    accelerator = accelerate.Accelerator(
        mixed_precision=cfg.mixed_precision if hasattr(cfg, 'mixed_precision') and cfg.mixed_precision != 'no' else None,
        gradient_accumulation_steps=cfg.gradient_accumulation_steps
    )
    # https://github.com/huggingface/transformers/issues/26548
    accelerator.dataloader_config.dispatch_batches = False

    if hasattr(cfg, 'force_wandb_name') and cfg.force_wandb_name:
        save_dir = cfg.save_dir.format(cfg.wandb_name)
    else:
        cfg.wandb_name = ','.join([f"{k[0]}:{v}" for k, v in unknown_cfg.items()]) if unknown_cfg else cfg.wandb_name
        save_dir = cfg.save_dir.format(cfg.latent_dims if hasattr(cfg, 'latent_dims') else cfg.wandb_name)

    logger = Logger(
        project=cfg.wandb_project,
        name=cfg.wandb_name,
        dummy=(cfg.wandb_project is None) or not (cfg.use_wandb),
        config=cfg,
    )

    print("Running Experiment:", cfg.wandb_name)


    if fmri_mode:
        fmt = getattr(cfg, "fmri_format", "avg" if str(cfg.dataset).endswith("avg") else "trial")
        split_cfg = FmriSplitConfig(
            data_dir=getattr(cfg, "fmri_data_dir", "fmri"),
            fmt=fmt,
            seed=int(getattr(cfg, "fmri_seed", 42)),
            train_on_overlap=bool(getattr(cfg, "fmri_train_on_overlap", False)),
            avg_train_size=int(getattr(cfg, "fmri_overlap_train_size", getattr(cfg, "fmri_avg_train_size", 500))),
            avg_val_size=int(getattr(cfg, "fmri_overlap_eval_size", getattr(cfg, "fmri_avg_val_size", 500))),
            avg_eval_size=int(getattr(cfg, "fmri_avg_eval_size", 1000)),
            trial_train_images=int(
                getattr(cfg, "fmri_overlap_train_images", getattr(cfg, "fmri_trial_train_images", 500))
            ),
            trial_val_images=int(
                getattr(cfg, "fmri_overlap_eval_images", getattr(cfg, "fmri_trial_val_images", 500))
            ),
            trial_eval_images=int(getattr(cfg, "fmri_trial_eval_images", 1000)),
        )
        supset, unsupset, val_paired, embed_dim = load_fmri_train_and_eval_datasets(
            cfg=split_cfg,
            sup_name=cfg.sup_emb,
            unsup_name=cfg.unsup_emb,
            sup_subject=int(getattr(cfg, "fmri_sup_subject", 1)),
            unsup_subject=int(getattr(cfg, "fmri_unsup_subject", 2)),
        )
        valset_paired = val_paired

        sup_encs = {}
        unsup_enc = {}
        encoder_dims = {cfg.sup_emb: int(embed_dim), cfg.unsup_emb: int(embed_dim)}
        translator = load_n_translator(cfg, encoder_dims)
    else:
        sup_encs = {
            cfg.sup_emb: load_encoder(cfg.sup_emb, mixed_precision=cfg.mixed_precision if hasattr(cfg, 'mixed_precision') else None)
        }
        encoder_dims = {
            cfg.sup_emb: get_sentence_embedding_dimension(sup_encs[cfg.sup_emb])
        }
        translator = load_n_translator(cfg, encoder_dims)

    model_save_dir = os.path.join(save_dir, 'model.pt')
    disc_save_dir = os.path.join(save_dir, 'disc.pt')

    os.makedirs(save_dir, exist_ok=True)

    assert hasattr(cfg, 'unsup_emb')
    assert cfg.sup_emb != cfg.unsup_emb

    if not fmri_mode:
        unsup_enc = {
            cfg.unsup_emb: load_encoder(cfg.unsup_emb, mixed_precision=cfg.mixed_precision if hasattr(cfg, 'mixed_precision') else None)
        }
        unsup_dim = {
            cfg.unsup_emb: get_sentence_embedding_dimension(unsup_enc[cfg.unsup_emb])
        }
        translator.add_encoders(unsup_dim, overwrite_embs=[cfg.unsup_emb])

        assert cfg.unsup_emb not in sup_encs
        assert cfg.unsup_emb in translator.in_adapters
        assert cfg.unsup_emb in translator.out_adapters

    cfg.num_params = sum(x.numel() for x in translator.parameters())
    print("Number of parameters:", cfg.num_params)
    print("Number of *trainable* parameters:", sum(p.numel() for p in translator.parameters() if p.requires_grad))
    print(translator)

    logger = Logger(
        project=cfg.wandb_project,
        name=cfg.wandb_name,
        dummy=(cfg.wandb_project is None) or not (cfg.use_wandb),
        config=cfg,
    )

    num_workers = min(get_num_proc(), 8)
    if fmri_mode:
        # Precomputed fMRI datasets were constructed above.
        pass
    elif cfg.dataset != 'mimic':
        dset = load_streaming_embeddings(cfg.dataset)
        print(f"Using {num_workers} workers and {len(dset)} datapoints")

        dset_dict = dset.train_test_split(test_size=cfg.val_size, seed=cfg.val_dataset_seed)
        dset = dset_dict["train"]
        valset = dset_dict["test"]

        assert hasattr(cfg, 'num_points') or hasattr(cfg, 'unsup_points')
        dset = dset.shuffle(seed=cfg.train_dataset_seed)
        if hasattr(cfg, 'num_points'):
            assert cfg.num_points > 0 and cfg.num_points <= len(dset) // 2
            supset = dset.select(range(cfg.num_points))
            unsupset = dset.select(range(cfg.num_points, cfg.num_points * 2))
        elif hasattr(cfg, 'unsup_points'):
            unsupset = dset.select(range(min(cfg.unsup_points, len(dset))))
            supset = dset.select(range(min(cfg.unsup_points, len(dset)), len(dset) - len(unsupset)))
    else:
        supset = load_from_disk('data/mimic')['supervised'].shuffle(cfg.train_dataset_seed).select(range(cfg.num_points))
        unsupset = load_from_disk('data/mimic')['unsupervised'].shuffle(cfg.train_dataset_seed).select(range(cfg.num_points))
        valset = load_from_disk('data/mimic')['evaluation'].shuffle(cfg.val_dataset_seed).select(range(cfg.val_size))

        # for each, drop all columns but 'text' using remove_columns
        supset = supset.remove_columns([col for col in supset.column_names if col != 'text'])
        unsupset = unsupset.remove_columns([col for col in unsupset.column_names if col != 'text'])
        valset = valset.remove_columns([col for col in valset.column_names if col != 'text'])
        

    if not fmri_mode:
        supset = MultiencoderTokenizedDataset(
            dataset=supset,
            encoders=sup_encs,
            n_embs_per_batch=cfg.n_embs_per_batch,
            batch_size=cfg.bs,
            max_length=cfg.max_seq_length,
            seed=cfg.sampling_seed,
        )
        unsupset = MultiencoderTokenizedDataset(
            dataset=unsupset,
            encoders=unsup_enc,
            n_embs_per_batch=1,
            batch_size=cfg.bs,
            max_length=cfg.max_seq_length,
            seed=cfg.sampling_seed,
        )

    sup_dataloader = DataLoader(
        supset,
        batch_size=cfg.bs,
        num_workers=num_workers // 2,
        shuffle=True,
        pin_memory=True,
        prefetch_factor=None,
        collate_fn=TokenizedCollator(),
        drop_last=cfg.train_drop_last if hasattr(cfg, 'train_drop_last') else not fmri_mode,
    )
    unsup_dataloader = DataLoader(
        unsupset,
        batch_size=cfg.bs,
        num_workers=num_workers // 2,
        shuffle=True,
        pin_memory=True,
        prefetch_factor=None,
        collate_fn=TokenizedCollator(),
        drop_last=cfg.train_drop_last if hasattr(cfg, 'train_drop_last') else not fmri_mode,
    )

    if use_val_set:
        if fmri_mode:
            valset = valset_paired
        else:
            valset = MultiencoderTokenizedDataset(
                dataset=valset,
                encoders={ **unsup_enc, **sup_encs },
                n_embs_per_batch=2,
                batch_size=cfg.val_bs,
                max_length=cfg.max_seq_length,
                seed=cfg.sampling_seed,
            )
        val_batch_size = min(cfg.val_bs if hasattr(cfg, 'val_bs') else cfg.bs, len(valset))
        valloader = DataLoader(
            valset,
            batch_size=val_batch_size,
            num_workers=num_workers,
            shuffle=False,
            pin_memory=True,
            prefetch_factor=(8 if num_workers > 0 else None),
            collate_fn=TokenizedCollator(),
            drop_last=False,
        )
        valloader = accelerator.prepare(valloader)

    gen_optimizers = build_generator_optimizers(cfg, translator)
    
    ######################################################################################
    disc = Discriminator(
        latent_dim=translator.in_adapters[cfg.unsup_emb].in_dim,
        discriminator_dim=cfg.disc_dim,
        depth=cfg.disc_depth,
        weight_init=cfg.weight_init
    )
    disc_opt = torch.optim.Adam(disc.parameters(), lr=cfg.disc_lr, eps=cfg.eps, betas=(0.5, 0.999))

    cfg.num_disc_params = sum(x.numel() for x in disc.parameters())
    print(f"Number of discriminator parameters:", cfg.num_disc_params)
    ######################################################################################
    sup_disc = Discriminator(
        latent_dim=translator.in_adapters[cfg.sup_emb].in_dim,
        discriminator_dim=cfg.disc_dim, 
        depth=cfg.disc_depth, 
    )
    sup_disc_opt = torch.optim.Adam(sup_disc.parameters(), lr=cfg.disc_lr, eps=cfg.eps, betas=(0.5, 0.999))

    cfg.num_sup_disc_params = sum(x.numel() for x in sup_disc.parameters())
    print(f"Number of supervised discriminator parameters:", cfg.num_sup_disc_params)
    print(sup_disc)
    ######################################################################################
    latent_disc = Discriminator(
        latent_dim=cfg.d_adapter,
        discriminator_dim=cfg.disc_dim,
        depth=cfg.disc_depth,
        weight_init=cfg.weight_init
    )
    latent_disc_opt = torch.optim.RMSprop(latent_disc.parameters(), lr=cfg.disc_lr, eps=cfg.eps)
    cfg.num_latent_disc_params = sum(x.numel() for x in latent_disc.parameters())
    print(f"Number of latent discriminator parameters:", cfg.num_latent_disc_params)
    print(latent_disc)
    latent_disc_opt = torch.optim.Adam(latent_disc.parameters(), lr=cfg.disc_lr, eps=cfg.eps, betas=(0.5, 0.999))
    ######################################################################################
    similarity_disc = Discriminator(
        latent_dim=cfg.bs,
        discriminator_dim=cfg.disc_dim,
        depth=cfg.disc_depth,
        weight_init=cfg.weight_init
    )
    similarity_disc_opt = torch.optim.RMSprop(similarity_disc.parameters(), lr=cfg.disc_lr, eps=cfg.eps)
    cfg.num_similarity_disc_params = sum(x.numel() for x in similarity_disc.parameters())
    print(f"Number of similarity discriminator parameters:", cfg.num_similarity_disc_params)
    print(similarity_disc)
    similarity_disc_opt = torch.optim.Adam(similarity_disc.parameters(), lr=cfg.disc_lr, eps=cfg.eps, betas=(0.5, 0.999))
    ######################################################################################

    max_num_epochs = int(np.ceil(cfg.epochs))
    steps_per_epoch = (len(supset) + cfg.bs - 1) // cfg.bs if fmri_mode else len(supset) // cfg.bs
    total_steps = steps_per_epoch * cfg.epochs / cfg.gradient_accumulation_steps
    warmup_length = (cfg.warmup_length if hasattr(cfg, 'warmup_length') else 100)

    def lr_lambda(step):
        if step < warmup_length:
            return min(1, step / warmup_length)
        else:
            if hasattr(cfg, 'no_scheduler') and cfg.no_scheduler:
                return 1
            return 1 - (step - warmup_length) / max(1, total_steps - warmup_length)

    gen_schedulers = [LambdaLR(opt_, lr_lambda=lr_lambda) for opt_ in gen_optimizers]
    disc_scheduler = LambdaLR(disc_opt, lr_lambda=lr_lambda)
    sup_disc_scheduler = LambdaLR(sup_disc_opt, lr_lambda=lr_lambda)
    latent_disc_scheduler = LambdaLR(latent_disc_opt, lr_lambda=lr_lambda)
    similarity_disc_scheduler = LambdaLR(similarity_disc_opt, lr_lambda=lr_lambda)

    if cfg.finetune_mode:
        assert hasattr(cfg, 'load_dir')
        print(f"Loading models from {cfg.load_dir}...")
        translator.load_state_dict(torch.load(cfg.load_dir + 'model.pt', map_location='cpu'), strict=False)
        disc.load_state_dict(torch.load(cfg.load_dir + 'disc.pt', map_location='cpu'))

    prepared = accelerator.prepare(translator, *gen_optimizers, *gen_schedulers)
    translator = prepared[0]
    n_gen_opt = len(gen_optimizers)
    gen_optimizers = list(prepared[1 : 1 + n_gen_opt])
    gen_schedulers = list(prepared[1 + n_gen_opt : 1 + n_gen_opt + len(gen_schedulers)])
    disc, disc_opt, disc_scheduler = accelerator.prepare(disc, disc_opt, disc_scheduler)
    sup_disc, sup_disc_opt, sup_disc_scheduler = accelerator.prepare(sup_disc, sup_disc_opt, sup_disc_scheduler)
    latent_disc, latent_disc_opt, latent_disc_scheduler = accelerator.prepare(latent_disc, latent_disc_opt, latent_disc_scheduler)
    similarity_disc, similarity_disc_opt, similarity_disc_scheduler = accelerator.prepare(
        similarity_disc, similarity_disc_opt, similarity_disc_scheduler
    )
    sup_dataloader, unsup_dataloader = accelerator.prepare(sup_dataloader, unsup_dataloader)


    if cfg.gan_style == "vanilla":
        gan_cls = VanillaGAN
    elif cfg.gan_style == "least_squares":
        gan_cls = LeastSquaresGAN
    elif cfg.gan_style == "relativistic":
        gan_cls = RelativisticGAN
    else:
        raise ValueError(f"Unknown GAN style: {cfg.gan_style}")
    latent_gan = gan_cls(
        cfg=cfg,
        generator=translator,
        discriminator=latent_disc,
        discriminator_opt=latent_disc_opt,
        discriminator_scheduler=latent_disc_scheduler,
        accelerator=accelerator,
    )
    similarity_gan = gan_cls(
        cfg=cfg,
        generator=translator,
        discriminator=similarity_disc,
        discriminator_opt=similarity_disc_opt,
        discriminator_scheduler=similarity_disc_scheduler,
        accelerator=accelerator,
    )
    gan = gan_cls(
        cfg=cfg,
        generator=translator,
        discriminator=disc,
        discriminator_opt=disc_opt,
        discriminator_scheduler=disc_scheduler,
        accelerator=accelerator,
    )
    sup_gan = gan_cls(
        cfg=cfg,
        generator=translator,
        discriminator=sup_disc,
        discriminator_opt=sup_disc_opt,
        discriminator_scheduler=sup_disc_scheduler,
        accelerator=accelerator
    )

    sup_iter = None
    if hasattr(cfg, 'unsup_points'):
        sup_iter = iter(sup_dataloader)

    if hasattr(cfg, 'val_size') and hasattr(cfg, 'patience') and hasattr(cfg, 'min_delta'):
        early_stopper = EarlyStopper(patience=cfg.patience, min_delta=cfg.min_delta, increase=False)
        early_stopping = True
    else:
        early_stopping = False
    eval_every = max(1, int(getattr(cfg, "eval_every", 1)))

    for epoch in range(max_num_epochs):
        should_eval = use_val_set and (epoch % eval_every == 0)
        if should_eval:
            with torch.no_grad(), accelerator.autocast():
                translator.eval()
                val_res = {}
                encs = None if fmri_mode else {**sup_encs, **unsup_enc}
                recons, trans, heatmap_dict, _, _, _ = eval_loop_(cfg, translator, encs, valloader, device=accelerator.device)
                for flag, res in recons.items():
                    for k, v in res.items():
                        if k == 'cos':
                            val_res[f"val/rec_{flag}_{k}"] = v
                for target_flag, d in trans.items():
                    for flag, res in d.items():
                        for k, v in res.items():
                            if flag == cfg.unsup_emb and target_flag == cfg.unsup_emb:
                                continue
                            val_res[f"val/{flag}_{target_flag}_{k}"] = v

                if len(heatmap_dict) > 0:
                    for k,v in heatmap_dict.items():
                        if "heatmap" in k and 'top' not in k:
                            v = wandb.Image(v)
                            val_res[f"val/{k}"] = v
                        else:
                            val_res[f"val/{k} (avg. {cfg.top_k_batches} batches)"] = v
                wandb.log(val_res)
                translator.train()

            if epoch >= cfg.min_epochs and early_stopping:
                score = np.mean([v for k, v in val_res.items() if 'top_rank' in k])

                if early_stopper.early_stop(score):
                    print("Early stopping...")
                    break
                if early_stopper.counter == 0 and score < early_stopper.opt_val:
                    print(f"Saving model (counter = {early_stopper.counter})... {score} < {early_stopper.opt_val} is the best score so far...")
                    save_everything(cfg, translator, gen_optimizers[0], [gan, sup_gan, latent_gan, similarity_gan], save_dir)

        max_num_batches = None
        print(f"Epoch", epoch, "max_num_batches", max_num_batches, "max_num_epochs", max_num_epochs)
        if epoch + 1 >= max_num_epochs:
            max_num_batches = max(1, (cfg.epochs - epoch) * len(supset) // cfg.bs)
            print(f"Setting max_num_batches to {max_num_batches}")

        sup_iter = training_loop_(
            save_dir=save_dir,
            accelerator=accelerator,
            translator=translator,
            gan=gan,
            sup_gan=sup_gan,
            latent_gan=latent_gan,
            similarity_gan=similarity_gan,
            sup_dataloader=sup_dataloader,
            sup_iter=sup_iter,
            unsup_dataloader=unsup_dataloader,
            sup_encs=sup_encs,
            unsup_enc=unsup_enc,
            cfg=cfg,
            gen_optimizers=gen_optimizers,
            gen_schedulers=gen_schedulers,
            logger=logger,
            max_num_batches=max_num_batches
        )

    with open(save_dir + 'config.toml', 'w') as f:
        toml.dump(cfg.__dict__, f)


if __name__ == "__main__":
    main()