main_mldg.py

from pathlib import Path

from fire import Fire

from flags import Flags
from model import ModelMLDG


def train(
    batch_size: int = 128,
    data_dir: str = "data",
    debug: bool = False,
    inner_loops: int = 200000,
    log_dir: str = "logs",
    lr: float = 0.0001,
    meta_step_size: float = 0.0001,  # proposal
    meta_val_beta: float = 0.0001,  # proposal
    model_dir: str = "models",
    momentum: float = 0.9,
    num_classes: int = 10,
    state_dict: str = "",
    step_size: int = 1,
    stop_gradient: bool = False,  # proposal
    test_every: int = 50,
    unseen_index: int = 0,
    weight_decay: float = 0.00005,
):
    flags = Flags(
        batch_size=batch_size,
        data_dir=Path(data_dir),
        debug=debug,
        inner_loops=inner_loops,
        log_dir=Path(log_dir),
        lr=lr,
        meta_step_size=meta_step_size,
        meta_val_beta=meta_val_beta,
        model_dir=Path(model_dir),
        momentum=momentum,
        num_classes=num_classes,
        state_dict=state_dict,
        step_size=step_size,
        stop_gradient=stop_gradient,
        test_every=test_every,
        unseen_index=unseen_index,
        weight_decay=weight_decay,
    )

    flags.create_dirs()

    model_obj = ModelMLDG(flags=flags)
    model_obj.train()

    # after training, we should test the held out domain
    model_obj.heldout_test()


if __name__ == "__main__":
    Fire({"train": train})