train_all.py

import os
import argparse
import collections
import random
import sys
from pathlib import Path

import numpy as np
import PIL
import torch
import torchvision
from sconf import Config
from prettytable import PrettyTable

from domainbed.datasets import get_dataset
from domainbed import hparams_registry
from domainbed.lib import misc
from domainbed.lib.writers import get_writer
from domainbed.lib.logger import Logger
from domainbed.trainer import train


def str2bool(string):
    if string == "True":
        return True
    else:
        return False
    
def main():
    parser = argparse.ArgumentParser(
        description="Domain generalization", allow_abbrev=False
    )
    parser.add_argument("name", type=str)
    parser.add_argument("configs", nargs="*")
    parser.add_argument("--data_dir", type=str, default="datadir/")
    parser.add_argument("--dataset", type=str, default="PACS")
    parser.add_argument("--algorithm", type=str, default="ERM")
    parser.add_argument(
        "--trial_seed",
        type=int,
        default=0,
        help="Trial number (used for seeding split_dataset and random_hparams).",
    )
    parser.add_argument(
        "--hparams_seed",
        type=int,
        default=0,
        help="Hparams seed number (used for seeding split_dataset and random_hparams).",
    )
    parser.add_argument("--seed", type=int, default=0, help="Seed for everything else")
    parser.add_argument(
        "--steps",
        type=int,
        default=None,
        help="Number of steps. Default is dataset-dependent.",
    )
    parser.add_argument(
        "--checkpoint_freq",
        type=int,
        default=None,
        help="Checkpoint every N steps. Default is dataset-dependent.",
    )

    parser.add_argument("--test_envs", type=int, nargs="+", default=None)
    parser.add_argument("--holdout_fraction", type=float, default=0.2)
    parser.add_argument(
        "--model_save", default=None, type=int, help="Model save start step"
    )
    parser.add_argument("--tb_freq", default=10)
    parser.add_argument("--debug", action="store_true", help="Run w/ debug mode")
    parser.add_argument(
        "--show", action="store_true", help="Show args and hparams w/o run"
    )
    parser.add_argument(
        "--evalmode",
        default="fast",
        help="[fast, all]. if fast, ignore train_in datasets in evaluation time.",
    )
    parser.add_argument(
        "--prebuild_loader", action="store_true", help="Pre-build eval loaders"
    )

    # [#] our args
    parser.add_argument("--lmd", type=float, default=0.5)
    parser.add_argument("--tmp", type=float, default=0.1)
    parser.add_argument("--clip_backbone", type=str, default="")
    parser.add_argument("--backbone", type=str, default="") #vit-base
    parser.add_argument("--prompt_style", type=int, default=2)
    parser.add_argument("--pretrained", default=None, type=str)
    parser.add_argument("--sweep", type=str, default="", help="Hparam Sweep")
    parser.add_argument("--output_dir", default=None, type=str)
    parser.add_argument("--miro_swad", default=None, type=str)
    parser.add_argument("--hpt_st", type=int, default=0)
    parser.add_argument("--cls_pth", type=str, default="")
    parser.add_argument("--check_acc", action="store_true")
    parser.add_argument("--open_set", action="store_true")
    parser.add_argument("--model_load", type=int, default=0)
    parser.add_argument("--swad_fix", action="store_false")
    parser.add_argument("--addn_proj", default="False", type=str)

    # [#] cutmix args
    parser.add_argument("--beta", type=float, default=0.0)
    parser.add_argument("--cutmix_prob", type=float, default=0.0)
    args, left_argv = parser.parse_known_args()
    args.deterministic = True

    # setup hparams
    # hparams = hparams_registry.default_hparams(args.algorithm, args.dataset)
    if args.hpt_st == 0:
        hparams = hparams_registry.default_hparams(args.algorithm, args.dataset)
    elif args.hpt_st == 1:
        print("HP Tuning strategy 1 .....")
        hparams = hparams_registry.random_hparams(
            args.algorithm,
            args.dataset,
            misc.seed_hash(args.hparams_seed, args.trial_seed),
        )
    else:
        print("HP Tuning strategy 2 .....")
        hparams = hparams_registry.random_hparams_st2(
            args.algorithm,
            args.dataset,
            misc.seed_hash(args.hparams_seed, args.trial_seed),
        )

    keys = ["config.yaml"] + args.configs
    keys = [open(key, encoding="utf8") for key in keys]
    hparams = Config(*keys, default=hparams)
    hparams.argv_update(left_argv)

    # setup debug
    if args.debug:
        args.checkpoint_freq = 5
        args.steps = 10
        args.name += "_debug"

    timestamp = misc.timestamp()
    args.unique_name = args.name
    args.unique_name = f"{timestamp}_{args.name}"

    # [#] Our args
    hparams["prompt_style"] = args.prompt_style
    hparams["pretrained"] = args.pretrained
    hparams["cutmix_prob"] = args.cutmix_prob
    hparams["addn_proj"] = str2bool(args.addn_proj)
    if args.test_envs is not None:
        args.test_envs = [[env] for env in args.test_envs]
    if args.dataset == "DomainNet":
        hparams["test_batchsize"] = 64
    # hparams["beta"] = args.beta
    hparams["lmd"] = args.lmd
    hparams["cls_pth"] = args.cls_pth
    hparams["tmp"] = args.tmp
    hparams["backbone"] = args.backbone
    if args.clip_backbone != "":
        hparams["clip_backbone"] = args.clip_backbone
    else:
        hparams["clip_backbone"] = "ViT-B/16"

    # path setup
    args.work_dir = Path(".")
    args.data_dir = Path(args.data_dir)

    if args.sweep == "":
        args.out_root = args.work_dir / Path("train_output") / args.dataset
        args.out_dir = args.out_root / args.unique_name
        args.out_dir.mkdir(exist_ok=True, parents=True)
        logger = Logger.get(args.out_dir / "log.txt")
    else:
        print("Creating sweep directory ...")
        args.out_dir = Path(args.output_dir)
        args.out_dir.mkdir(exist_ok=True, parents=True)
        logger = Logger.get(args.out_dir / "log.txt")

    # writer = get_writer(args.out_root / "runs" / args.unique_name)
    if args.debug:
        logger.setLevel("DEBUG")
    cmd = " ".join(sys.argv)
    logger.info(f"Command :: {cmd}")

    logger.nofmt("Environment:")
    logger.nofmt("\tPython: {}".format(sys.version.split(" ")[0]))
    logger.nofmt("\tPyTorch: {}".format(torch.__version__))
    logger.nofmt("\tTorchvision: {}".format(torchvision.__version__))
    logger.nofmt("\tCUDA: {}".format(torch.version.cuda))
    logger.nofmt("\tCUDNN: {}".format(torch.backends.cudnn.version()))
    logger.nofmt("\tNumPy: {}".format(np.__version__))
    logger.nofmt("\tPIL: {}".format(PIL.__version__))

    # Different to DomainBed, we support CUDA only.
    assert torch.cuda.is_available(), "CUDA is not available"

    logger.nofmt("Args:")
    for k, v in sorted(vars(args).items()):
        logger.nofmt("\t{}: {}".format(k, v))

    logger.nofmt("HParams:")
    for line in hparams.dumps().split("\n"):
        logger.nofmt("\t" + line)

    if args.show:
        exit()

    # seed
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.backends.cudnn.deterministic = args.deterministic
    torch.backends.cudnn.benchmark = not args.deterministic

    # Dummy datasets for logging information.
    # Real dataset will be re-assigned in train function.
    # test_envs only decide transforms; simply set to zero.
    dataset, _in_splits, _out_splits = get_dataset([0], args, hparams)

    # print dataset information
    logger.nofmt("Dataset:")
    logger.nofmt(
        f"\t[{args.dataset}] #envs={len(dataset)}, #classes={dataset.num_classes}"
    )
    for i, env_property in enumerate(dataset.environments):
        logger.nofmt(f"\tenv{i}: {env_property} (#{len(dataset[i])})")
    logger.nofmt("")

    n_steps = args.steps or dataset.N_STEPS
    checkpoint_freq = args.checkpoint_freq or dataset.CHECKPOINT_FREQ
    logger.info(f"n_steps = {n_steps}")
    logger.info(f"checkpoint_freq = {checkpoint_freq}")

    org_n_steps = n_steps
    n_steps = (n_steps // checkpoint_freq) * checkpoint_freq + 1
    logger.info(f"n_steps is updated to {org_n_steps} => {n_steps} for checkpointing")

    if not args.test_envs:
        args.test_envs = [[te] for te in range(len(dataset))]
    logger.info(f"Target test envs = {args.test_envs}")

    ###########################################################################
    # Run
    ###########################################################################
    all_records = []
    results = collections.defaultdict(list)
    
    if args.check_acc:
        for test_env in args.test_envs:
            dom_name = dataset.environments[int(test_env[0])]
            result = train(
                test_env,
                args=args,
                hparams=hparams,
                n_steps=n_steps,
                checkpoint_freq=checkpoint_freq,
                logger=logger,
                writer=None,
            )
            if args.cls_pth != "":
                all_records.append((dom_name, result[1]["test_in"]))
            else:
                all_records.append((dom_name, result))
        
        if args.cls_pth != "":
            columns = ["Acc. Table"] + dataset.environments
            acc_table = PrettyTable(columns)
            row = ["Avg. Cls"]
            for i in range(len(all_records)):
                name, acc_val = all_records[i]
                row += [f"{acc_val * 100:.2f}"]
            acc_table.add_row(row)
        else:
            columns = ["Domain"] + hparams["class_names"]
            acc_table = PrettyTable(columns)
            for i in range(len(all_records)):
                name, acc_dict = all_records[i]
                row = [name] + [f"{acc.item() * 100:.2f}" for acc in list(acc_dict.values())]
                acc_table.add_row(row)
        print(acc_table)
        exit()

    for test_env in args.test_envs:
        res, records = train(
            test_env,
            args=args,
            hparams=hparams,
            n_steps=n_steps,
            checkpoint_freq=checkpoint_freq,
            logger=logger,
            writer=None,
        )
        all_records.append(records)
        for k, v in res.items():
            results[k].append(v)

    # log summary table
    logger.info("=== Summary ===")
    logger.info(f"Command: {' '.join(sys.argv)}")
    logger.info("Unique name: %s" % args.unique_name)
    logger.info("Out path: %s" % args.out_dir)
    logger.info("Algorithm: %s" % args.algorithm)
    logger.info("Dataset: %s" % args.dataset)

    table = PrettyTable(["Selection"] + dataset.environments + ["Avg."])
    for key, row in results.items():
        row.append(np.mean(row))
        row = [f"{acc:.3%}" for acc in row]
        table.add_row([key] + row)
    logger.nofmt(table)


if __name__ == "__main__":
    main()