import os import torch import argparse import copy from glob import glob from torch.utils.data import DataLoader, random_split from tensorboardX import SummaryWriter from data.dataloader_nba import NBADatasetMinMax as NBADatasetMinMax from data.dataloader_nba import seq_collate_nba, seq_collate_imle_train from utils.config import Config from utils.utils import back_up_code_git, set_random_seed, log_config_to_file from models.flow_matching import FlowMatcher from models.imle import IMLE from models.backbone import IMLETransformer from trainer.imle_trainers import IMLETrainer def parse_config(): """ Parse the command line arguments and return the configuration options. """ parser = argparse.ArgumentParser() # Basic configuration parser.add_argument('--ckpt_path', type=str, default=None, help='Path to the checkpoint to load the model from.') parser.add_argument('--cfg', default='auto', type=str, help="Config file path") parser.add_argument('--exp', default='', type=str, help='Experiment description for each run, name of the saving folder.') parser.add_argument('--save_samples', default=False, action='store_true', help='Save the samples during evaluation.') parser.add_argument('--eval_on_train', default=False, action='store_true', help='Evaluate the model on the training set.') # Data configuration parser.add_argument('--batch_size', default=None, type=int, help='Override the batch size in the config file.') parser.add_argument('--data_dir', type=str, default='./data/nba', help='Directory where the data is stored.') parser.add_argument('--n_train', type=int, default=32500, help='Number training scenes used.') parser.add_argument('--n_test', type=int, default=12500, help='Number testing scenes used.') parser.add_argument('--rotate', default=False, action='store_true', help='Whether to rotate the data to canonical x-axis or not.') parser.add_argument('--data_norm', default='min_max', choices=['min_max', 'sqrt'], help='Normalization method for the data.') # Reproducibility configuration parser.add_argument('--fix_random_seed', action='store_true', default=False, help='fix random seed for reproducibility') parser.add_argument('--seed', type=int, default=42, help='Set the random seed to split the testing set for training evaluation.') return parser.parse_args() def init_basics(args): """ Init the basic configurations for the experiment. """ """Load the config file""" result_dir = os.path.abspath(os.path.join(args.ckpt_path, '../../')) if args.cfg == 'auto': yml_ls = glob(result_dir+'/*.yml') assert len(yml_ls) >= 1, 'At least one config file should be found in the directory.' yml_path = [f for f in yml_ls if '_updated.yml' in os.path.basename(f)][0] args.cfg = yml_path cfg = Config(args.cfg, f'{args.exp}', train_mode=False) tag = '_' ### Update data configuration ### def _update_data_params(args, cfg, tag): if args.n_train != 32500: tag += f'_subset{args.n_train}' return cfg, tag cfg, tag = _update_data_params(args, cfg, tag) def _update_optimization_params(args, cfg, tag): if args.batch_size is not None: # override the batch size cfg.train_batch_size = args.batch_size cfg.test_batch_size = args.batch_size return cfg, tag cfg, tag = _update_optimization_params(args, cfg, tag) ### voila, create the saving directory ### tag += '_train_set' if args.eval_on_train else '_test_set' tag = tag.replace('__', '_') cfg.device = 'cuda' if torch.cuda.is_available() else 'cpu' logger = cfg.create_dirs(tag_suffix=tag) """fix random seed""" if args.fix_random_seed: set_random_seed(args.seed) """set up tensorboard and text log""" tb_dir = os.path.abspath(os.path.join(cfg.log_dir, '../tb_eval')) os.makedirs(tb_dir, exist_ok=True) tb_log = SummaryWriter(log_dir=tb_dir) """print the config file""" log_config_to_file(cfg.yml_dict, logger=logger) return cfg, logger, tb_log def build_data_loader(cfg, args): """ Build the data loader for the NBA dataset. """ train_dset = NBADatasetMinMax( data_dir=args.data_dir, obs_len=cfg.past_frames, pred_len=cfg.future_frames, training=True, num_scenes=args.n_train, overfit=False, cfg=cfg, rotate=args.rotate, data_norm=args.data_norm, imle=True) train_loader = DataLoader( train_dset, batch_size=cfg.train_batch_size, shuffle=True, num_workers=0, collate_fn=seq_collate_imle_train, pin_memory=True) test_dset = NBADatasetMinMax( data_dir=args.data_dir, obs_len=cfg.past_frames, pred_len=cfg.future_frames, training=False, overfit=False, test_scenes=args.n_test, cfg=cfg, rotate=args.rotate, data_norm=args.data_norm, imle=False) test_loader = DataLoader( test_dset, batch_size=cfg.test_batch_size, ### change it from 500 shuffle=False, num_workers=4, collate_fn=seq_collate_nba, pin_memory=True) return train_loader, test_loader def build_network(cfg, args, logger): """ Build the network for the denoising model. """ model = IMLETransformer( model_config=cfg.MODEL, logger=logger, config=cfg, ) imle_model = IMLE( cfg=cfg, model=model, logger=logger, ) return imle_model def main(): """ Main function to train the model. """ """Init everything""" args = parse_config() cfg, logger, tb_log = init_basics(args) train_loader, test_loader = build_data_loader(cfg, args) imle_model = build_network(cfg, args, logger) """Train or evaluate the model""" trainer = IMLETrainer( cfg, imle_model, train_loader, test_loader, tb_log=tb_log, logger=logger, gradient_accumulate_every=1, ema_decay = 0.995, ema_update_every = 1, save_samples=args.save_samples ) ### grid search trainer.test(mode='best', eval_on_train=args.eval_on_train) if __name__ == "__main__": main()