train_pm_vae.py

import json
import os
import pickle
import random

import jax
import jax.numpy as jnp
import optax
from absl import app, flags
from bax import Trainer
from bax.callbacks import CheckpointCallback, LearningRateLoggerCallback
from ml_collections.config_flags import config_flags

from posterior_matching.models.vae import PosteriorMatchingVAE
from posterior_matching.utils import (
    configure_environment,
    load_datasets,
    cyclical_annealing_schedule,
    make_run_dir,
    TensorBoardCallback,
)

configure_environment()

config_flags.DEFINE_config_file("config", lock_config=False)


def get_beta_schedule(config):
    if "schedule" not in config:
        return lambda x: 1.0

    if config.schedule == "monotonic":
        return optax.linear_schedule(
            config.low_value,
            config.high_value,
            config.transition_steps,
            config.transition_begin,
        )

    if config.schedule == "cyclic":
        return cyclical_annealing_schedule(
            config.low_value, config.high_value, config.period, config.delay
        )


def main(_):
    config = flags.FLAGS.config

    if "seed" not in config:
        config.seed = random.randint(0, int(2e9))

    config.lock()

    train_dataset, val_dataset = load_datasets(config.data)
    is_image_data = "image" in train_dataset.element_spec
    data_key = "image" if is_image_data else "features"

    def loss_fn(step, is_training, batch):
        model = PosteriorMatchingVAE.from_config(config.model)
        out = model(batch[data_key], batch["mask"], is_training=is_training)

        beta_schedule = get_beta_schedule(config.get("beta", {}))
        beta = beta_schedule(step)

        out["beta"] = beta

        elbo = jnp.mean(out["reconstruction_ll"] - beta * out["kl"])
        matching_loss = -jnp.mean(out["matching_ll"])

        loss = -elbo + config.get("matching_coef", 1.0) * matching_loss

        return loss, jax.tree_map(jnp.mean, out)

    schedule = optax.exponential_decay(**config.lr_schedule)
    optimizer = optax.chain(
        optax.scale_by_adam(**config.get("adam", {})),
        optax.add_decayed_weights(
            config.get("weight_decay", 0.0),
            mask=lambda p: jax.tree_map(lambda x: x.ndim != 1, p),
        ),
        optax.scale_by_schedule(schedule),
        optax.scale(-1.0),
    )

    trainer = Trainer(loss_fn, optimizer, num_devices=1, seed=config.seed)

    run_dir = make_run_dir(prefix=f"pm-vae-{config.data.dataset}")
    print("Using run directory:", run_dir)

    callbacks = [
        CheckpointCallback(os.path.join(run_dir, "train_state.pkl")),
        LearningRateLoggerCallback(schedule),
        TensorBoardCallback(os.path.join(run_dir, "tb")),
    ]

    train_state = trainer.fit(
        train_dataset,
        config.steps,
        val_dataset=val_dataset,
        validation_freq=config.validation_freq,
        callbacks=callbacks,
    )

    if config.get("save_final_state", False):
        with open(os.path.join(run_dir, "train_state.pkl"), "wb") as fp:
            pickle.dump(train_state, fp)

    with open(os.path.join(run_dir, "model_config.json"), "w") as fp:
        json.dump(config.model.to_dict(), fp)


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