##################################################### # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.04 # ##################################################### # python exps/LFNA/basic.py ##################################################### import sys, time, torch, random, argparse 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 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 from models.xcore import get_model def main(args): torch.set_num_threads(args.workers) prepare_seed(args.rand_seed) logger = prepare_logger(args) dynamic_env = get_synthetic_env() historical_x, historical_y = None, None for idx, (timestamp, (allx, ally)) in enumerate(dynamic_env): if historical_x is not None: mean, std = historical_x.mean().item(), historical_x.std().item() else: mean, std = 0, 1 model_kwargs = dict(input_dim=1, output_dim=1, mean=mean, std=std) model = get_model(dict(model_type="simple_mlp"), **model_kwargs) # create the current data loader if historical_x is not None: train_dataset = torch.utils.data.TensorDataset(historical_x, historical_y) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, ) optimizer = torch.optim.Adam( model.parameters(), lr=args.init_lr, amsgrad=True ) criterion = torch.nn.MSELoss() lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones=[ int(args.epochs * 0.25), int(args.epochs * 0.5), int(args.epochs * 0.75), ], gamma=0.3, ) for _iepoch in range(args.epochs): results = basic_train_fn( train_loader, model, criterion, optimizer, MSEMetric(), logger ) lr_scheduler.step() if _iepoch % args.log_per_epoch == 0: log_str = ( "[{:}]".format(time_string()) + " [{:04d}/{:04d}][{:04d}/{:04d}]".format( idx, len(dynamic_env), _iepoch, args.epochs ) + " mse: {:.5f}, lr: {:.4f}".format( results["mse"], min(lr_scheduler.get_last_lr()) ) ) logger.log(log_str) results = basic_eval_fn(train_loader, model, MSEMetric(), logger) logger.log( "[{:}] [{:04d}/{:04d}] train-mse: {:.5f}".format( time_string(), idx, len(dynamic_env), results["mse"] ) ) metric = ComposeMetric(MSEMetric(), SaveMetric()) eval_dataset = torch.utils.data.TensorDataset(allx, ally) eval_loader = torch.utils.data.DataLoader( eval_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, ) results = basic_eval_fn(eval_loader, model, metric, logger) log_str = ( "[{:}]".format(time_string()) + " [{:04d}/{:04d}]".format(idx, len(dynamic_env)) + " eval-mse: {:.5f}".format(results["mse"]) ) logger.log(log_str) save_path = logger.path(None) / "{:04d}-{:04d}.pth".format( idx, len(dynamic_env) ) save_checkpoint( {"model": model.state_dict(), "index": idx, "timestamp": timestamp}, save_path, logger, ) # Update historical data if historical_x is None: historical_x, historical_y = allx, ally else: historical_x, historical_y = torch.cat((historical_x, allx)), torch.cat( (historical_y, ally) ) logger.log("") logger.log("-" * 200 + "\n") logger.close() if __name__ == "__main__": parser = argparse.ArgumentParser("Use all the past data to train.") parser.add_argument( "--save_dir", type=str, default="./outputs/lfna-synthetic/use-all-past-data", help="The checkpoint directory.", ) parser.add_argument( "--init_lr", type=float, default=0.1, help="The initial learning rate for the optimizer (default is Adam)", ) parser.add_argument( "--batch_size", type=int, default=256, help="The batch size", ) parser.add_argument( "--epochs", type=int, default=2000, help="The total number of epochs.", ) parser.add_argument( "--log_per_epoch", type=int, default=200, help="Log the training information per __ epochs.", ) 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" main(args)