autodl-projects/exps/LFNA/basic-maml.py

#####################################################
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.04 #
#####################################################
# python exps/LFNA/basic-maml.py --env_version v1 --inner_step 5
# python exps/LFNA/basic-maml.py --env_version v2
#####################################################
import sys, time, copy, torch, random, argparse
from tqdm import tqdm
from copy import deepcopy
from pathlib import Path

lib_dir = (Path(__file__).parent / ".." / ".." / "lib").resolve()
if str(lib_dir) not in sys.path:
    sys.path.insert(0, str(lib_dir))
from procedures import prepare_seed, prepare_logger, save_checkpoint, copy_checkpoint
from log_utils import time_string
from log_utils import AverageMeter, convert_secs2time

from utils import split_str2indexes

from procedures.advanced_main import basic_train_fn, basic_eval_fn
from procedures.metric_utils import SaveMetric, MSEMetric, ComposeMetric
from datasets.synthetic_core import get_synthetic_env, EnvSampler
from models.xcore import get_model
from xlayers import super_core

from lfna_utils import lfna_setup, TimeData


class MAML:
    """A LFNA meta-model that uses the MLP as delta-net."""

    def __init__(
        self, network, criterion, epochs, meta_lr, inner_lr=0.01, inner_step=1
    ):
        self.criterion = criterion
        # self.container = container
        self.network = network
        self.meta_optimizer = torch.optim.Adam(
            self.network.parameters(), lr=meta_lr, amsgrad=True
        )
        self.meta_lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
            self.meta_optimizer,
            milestones=[
                int(epochs * 0.8),
                int(epochs * 0.9),
            ],
            gamma=0.1,
        )
        self.inner_lr = inner_lr
        self.inner_step = inner_step
        self._best_info = dict(state_dict=None, iepoch=None, score=None)
        print("There are {:} weights.".format(self.network.get_w_container().numel()))

    def adapt(self, dataset):
        # create a container for the future timestamp
        container = self.network.get_w_container()

        for k in range(0, self.inner_step):
            y_hat = self.network.forward_with_container(dataset.x, container)
            loss = self.criterion(y_hat, dataset.y)
            grads = torch.autograd.grad(loss, container.parameters())
            container = container.additive([-self.inner_lr * grad for grad in grads])
        return container

    def predict(self, x, container=None):
        if container is not None:
            y_hat = self.network.forward_with_container(x, container)
        else:
            y_hat = self.network(x)
        return y_hat

    def step(self):
        torch.nn.utils.clip_grad_norm_(self.network.parameters(), 1.0)
        self.meta_optimizer.step()
        self.meta_lr_scheduler.step()

    def zero_grad(self):
        self.meta_optimizer.zero_grad()

    def load_state_dict(self, state_dict):
        self.criterion.load_state_dict(state_dict["criterion"])
        self.network.load_state_dict(state_dict["network"])
        self.meta_optimizer.load_state_dict(state_dict["meta_optimizer"])
        self.meta_lr_scheduler.load_state_dict(state_dict["meta_lr_scheduler"])

    def state_dict(self):
        state_dict = dict()
        state_dict["criterion"] = self.criterion.state_dict()
        state_dict["network"] = self.network.state_dict()
        state_dict["meta_optimizer"] = self.meta_optimizer.state_dict()
        state_dict["meta_lr_scheduler"] = self.meta_lr_scheduler.state_dict()
        return state_dict

    def save_best(self, score):
        success, best_score = self.network.save_best(score)
        return success, best_score

    def load_best(self):
        self.network.load_best()


def main(args):
    logger, env_info, model_kwargs = lfna_setup(args)
    model = get_model(**model_kwargs)

    dynamic_env = get_synthetic_env(mode="train", version=args.env_version)

    criterion = torch.nn.MSELoss()

    maml = MAML(
        model, criterion, args.epochs, args.meta_lr, args.inner_lr, args.inner_step
    )

    # meta-training
    last_success_epoch = 0
    per_epoch_time, start_time = AverageMeter(), time.time()
    for iepoch in range(args.epochs):
        need_time = "Time Left: {:}".format(
            convert_secs2time(per_epoch_time.avg * (args.epochs - iepoch), True)
        )
        head_str = (
            "[{:}] [{:04d}/{:04d}] ".format(time_string(), iepoch, args.epochs)
            + need_time
        )

        maml.zero_grad()
        meta_losses = []
        for ibatch in range(args.meta_batch):
            future_timestamp = dynamic_env.random_timestamp()
            _, (future_x, future_y) = dynamic_env(future_timestamp)
            past_timestamp = (
                future_timestamp - args.prev_time * dynamic_env.timestamp_interval
            )
            _, (past_x, past_y) = dynamic_env(past_timestamp)

            future_container = maml.adapt(TimeData(past_timestamp, past_x, past_y))
            future_y_hat = maml.predict(future_x, future_container)
            future_loss = maml.criterion(future_y_hat, future_y)
            meta_losses.append(future_loss)
        meta_loss = torch.stack(meta_losses).mean()
        meta_loss.backward()
        maml.step()

        logger.log(head_str + " meta-loss: {:.4f}".format(meta_loss.item()))
        success, best_score = maml.save_best(-meta_loss.item())
        if success:
            logger.log("Achieve the best with best_score = {:.3f}".format(best_score))
            save_checkpoint(maml.state_dict(), logger.path("model"), logger)
            last_success_epoch = iepoch
        if iepoch - last_success_epoch >= args.early_stop_thresh:
            logger.log("Early stop at {:}".format(iepoch))
            break

        per_epoch_time.update(time.time() - start_time)
        start_time = time.time()

    # meta-test
    maml.load_best()
    eval_env = env_info["dynamic_env"]
    assert eval_env.timestamp_interval == dynamic_env.timestamp_interval
    w_container_per_epoch = dict()
    for idx in range(args.prev_time, len(eval_env)):
        future_timestamp, (future_x, future_y) = eval_env[idx]
        past_timestamp = (
            future_timestamp.item() - args.prev_time * eval_env.timestamp_interval
        )
        _, (past_x, past_y) = eval_env(past_timestamp)
        future_container = maml.adapt(TimeData(past_timestamp, past_x, past_y))
        w_container_per_epoch[idx] = future_container.no_grad_clone()
        with torch.no_grad():
            future_y_hat = maml.predict(future_x, w_container_per_epoch[idx])
            future_loss = maml.criterion(future_y_hat, future_y)
        logger.log("meta-test: [{:03d}] -> loss={:.4f}".format(idx, future_loss.item()))
    save_checkpoint(
        {"w_container_per_epoch": w_container_per_epoch},
        logger.path(None) / "final-ckp.pth",
        logger,
    )

    logger.log("-" * 200 + "\n")
    logger.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser("Use the data in the past.")
    parser.add_argument(
        "--save_dir",
        type=str,
        default="./outputs/lfna-synthetic/use-maml",
        help="The checkpoint directory.",
    )
    parser.add_argument(
        "--env_version",
        type=str,
        required=True,
        help="The synthetic enviornment version.",
    )
    parser.add_argument(
        "--hidden_dim",
        type=int,
        default=16,
        help="The hidden dimension.",
    )
    parser.add_argument(
        "--meta_lr",
        type=float,
        default=0.01,
        help="The learning rate for the MAML optimizer (default is Adam)",
    )
    parser.add_argument(
        "--fail_thresh",
        type=float,
        default=1000,
        help="The threshold for the failure, which we reuse the previous best model",
    )
    parser.add_argument(
        "--inner_lr",
        type=float,
        default=0.005,
        help="The learning rate for the inner optimization",
    )
    parser.add_argument(
        "--inner_step", type=int, default=1, help="The inner loop steps for MAML."
    )
    parser.add_argument(
        "--prev_time",
        type=int,
        default=5,
        help="The gap between prev_time and current_timestamp",
    )
    parser.add_argument(
        "--meta_batch",
        type=int,
        default=64,
        help="The batch size for the meta-model",
    )
    parser.add_argument(
        "--epochs",
        type=int,
        default=2000,
        help="The total number of epochs.",
    )
    parser.add_argument(
        "--early_stop_thresh",
        type=int,
        default=50,
        help="The maximum epochs for early stop.",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=4,
        help="The number of data loading workers (default: 4)",
    )
    # Random Seed
    parser.add_argument("--rand_seed", type=int, default=-1, help="manual seed")
    args = parser.parse_args()
    if args.rand_seed is None or args.rand_seed < 0:
        args.rand_seed = random.randint(1, 100000)
    assert args.save_dir is not None, "The save dir argument can not be None"
    args.save_dir = "{:}-s{:}-mlr{:}-d{:}-prev{:}-e{:}-env{:}".format(
        args.save_dir,
        args.inner_step,
        args.meta_lr,
        args.hidden_dim,
        args.prev_time,
        args.epochs,
        args.env_version,
    )
    main(args)