train_common.py

import glob
import json
import os
import sys
from pathlib import Path

import utils.utils as utils
from dataset import create_dataset_no_norm, create_loader, create_sampler
from scheduler import create_scheduler_each_step
from utils.utils_attack import get_attacker
from utils.utils_optimizer import get_optimizer, get_trainable_params


def t2bool(t):
    if t.lower() == "true":
        return True
    if t.lower() == "false":
        return False
    raise ValueError("Invalid value")


def ensure_dir(path):
    os.makedirs(path, exist_ok=True)


def maybe_set_output_dir_from_eval(args):
    if not getattr(args, "evaluate", False):
        return args.output_dir
    if args.eval_ckpt_path is None:
        return args.output_dir
    if os.path.isfile(args.eval_ckpt_path):
        output_dir = os.path.dirname(args.eval_ckpt_path)
        print("Output directory:", output_dir)
        args.output_dir = output_dir
    return args.output_dir


def save_args_and_configs(args, config, train_config=None):
    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    with open(os.path.join(args.output_dir, "args.json"), "w") as f:
        json.dump(vars(args), f, indent=4)
    with open(os.path.join(args.output_dir, "config.json"), "w") as f:
        json.dump(config, f, indent=4)
    if train_config is not None:
        with open(os.path.join(args.output_dir, "train_config.json"), "w") as f:
            json.dump(train_config, f, indent=4)


def setup_output_dir_basic(args, utils):
    if getattr(args, "evaluate", False):
        print("Evaluate mode")
        maybe_set_output_dir_from_eval(args)

    ensure_dir(args.output_dir)
    sys.stdout = utils.Tee(
        sys.stdout, open(os.path.join(args.output_dir, "out.txt"), "w")
    )


def setup_output_dir_with_overwrite(
    args,
    utils,
    allow_overwrite=False,
    eval_overwrite_attacks=None,
    load_prev=False,
):
    if getattr(args, "evaluate", False):
        print("Evaluate mode")
        maybe_set_output_dir_from_eval(args)

    log_stats_prev = {}
    eval_overwrite_attacks = eval_overwrite_attacks or []

    if os.path.exists(args.output_dir):
        results_paths = glob.glob(os.path.join(args.output_dir, "final_results.json"))
        if not allow_overwrite and len(eval_overwrite_attacks) == 0:
            if len(results_paths) > 0:
                print("Already exists:", args.output_dir)
                print("Skip\n\n")
                exit()
        if len(results_paths) > 0 and load_prev:
            print("Load previous results:", args.output_dir)
            with open(results_paths[0], "r") as f:
                log_stats_prev = json.load(f)
            print(log_stats_prev)

    ensure_dir(args.output_dir)
    sys.stdout = utils.Tee(
        sys.stdout, open(os.path.join(args.output_dir, "out.txt"), "w")
    )
    return log_stats_prev


def build_optimizer_and_scheduler(
    args, model, train_config, train_loader, vision_keywords=("visual", "vision")
):
    parameters = []
    for name, param in model.named_parameters():
        if getattr(args, "train_only_vision_encoder", False):
            if not any(k in name for k in vision_keywords):
                print("--------> no update: ", name)
                continue
        if not param.requires_grad:
            continue
        lr = train_config["optimizer"]["lr"]
        wd = train_config["optimizer"]["weight_decay"]
        parameters.append({"params": param, "lr": lr, "weight_decay": wd})
        print("----> update {}: lr={}, wd={}".format(name, lr, wd))
    opt_config = utils.AttrDict(train_config["optimizer"])
    optimizer = get_optimizer(parameters, opt_config)
    arg_sche = utils.AttrDict(train_config["schedular"])
    lr_scheduler, _ = create_scheduler_each_step(arg_sche, optimizer, train_loader)
    return optimizer, lr_scheduler


def build_optimizer_and_scheduler_clip(args, model, train_config, train_loader):
    if args.train_only_vision_encoder:
        print("Train only vision encoder")
        parameters = get_trainable_params(model.visual, train_config)
        for name, param in model.named_parameters():
            if "visual" not in name:
                param.requires_grad_(False)
    else:
        parameters = get_trainable_params(model, train_config)

    opt_config = utils.AttrDict(train_config["optimizer"])
    optimizer = get_optimizer(parameters, opt_config)

    arg_sche = utils.AttrDict(train_config["schedular"])
    lr_scheduler, _ = create_scheduler_each_step(arg_sche, optimizer, train_loader)
    return optimizer, lr_scheduler


def build_attacker(
    args, train_config, attack_config, model, ref_model, tokenizer, **kwargs
):
    if args.attack is None:
        return None
    return get_attacker(
        args,
        train_config,
        args.attack,
        model,
        ref_model,
        tokenizer,
        attack_config=attack_config,
        eps=args.epsilon,
        steps=args.num_iters,
        step_size=args.step_size,
        **kwargs,
    )


def build_eval_attackers(
    args, train_config, attack_list, model, ref_model, tokenizer, **kwargs
):
    return {
        att: get_attacker(
            args, train_config, att, model, ref_model, tokenizer, **kwargs
        )
        for att in attack_list
    }


def build_retrieval_dataloaders(
    args, config, batch_size_train, batch_size_test, eval_with_info=False
):
    print("Creating dataset")
    train_dataset, val_dataset, test_dataset, train_dataset_for_eval = (
        create_dataset_no_norm(
            "re",
            config,
            get_train_eval=True,
            aug_n=args.aug_n,
            aug_m=args.aug_m,
            aug_scale=(args.aug_scale, 1.0),
            n_holes=args.n_holes,
            length_ratio=args.length_ratio,
            degrees=args.degrees,
            translate=args.translate,
            scale=args.scale,
            eval_with_info=eval_with_info,
        )
    )
    datasets = [train_dataset, val_dataset, test_dataset, train_dataset_for_eval]
    if args.distributed:
        num_tasks = utils.get_world_size()
        global_rank = utils.get_rank()
        samplers = create_sampler(datasets, [True], num_tasks, global_rank)
    else:
        samplers = [None]
    samplers += [None] * 3

    train_loader, val_loader, test_loader, train_subset_loader = create_loader(
        datasets,
        samplers,
        batch_size=[batch_size_train] + [batch_size_test] * 3,
        num_workers=[4] * 4,
        is_trains=[True, False, False, False],
        collate_fns=[None] * 4,
    )
    return train_loader, val_loader, test_loader, train_subset_loader


def write_epoch_logs(output_dir, log_stats, eval_results_dict, epoch):
    log_stats.update({"output_dir": output_dir})
    with open(os.path.join(output_dir, "log.json"), "a") as f:
        json.dump(log_stats, f, indent=4)
        f.write("\n")
    print(eval_results_dict)

    with open(os.path.join(output_dir, f"eval_results_ep{epoch}.json"), "w") as f:
        json.dump(eval_results_dict, f, indent=4)


def write_final_results(output_dir, log_stats):
    with open(os.path.join(output_dir, "log.txt"), "w") as f:
        f.write(str(log_stats))
    with open(os.path.join(output_dir, "final_results.json"), "w") as f:
        json.dump(log_stats, f, indent=4)