import os import torch import argparse import copy from torch.utils.data import DataLoader, random_split from tensorboardX import SummaryWriter from data.dataloader_eth_ucy import ETHDataset, seq_collate_eth 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.backbone_eth_ucy import ETHMotionTransformer from trainer.denoising_model_trainers import Trainer def parse_config(): """ Parse the command line arguments and return the configuration options. """ parser = argparse.ArgumentParser() # Basic configuration parser.add_argument('--cfg', default='cfg/eth_ucy/cor_fm.yml', type=str, help="Config file path") parser.add_argument('--exp', default='', type=str, help='Experiment description for each run, name of the saving folder.') # Data configuration parser.add_argument('--data_source', default='original', type=str, help='Data source for the experiment. Either be original or preprocessed ones from LED.') parser.add_argument('--epochs', default=None, type=int, help='Override the number of epochs in the config file.') parser.add_argument('--batch_size', default=32, type=int, help='Override the batch size in the config file.') parser.add_argument('--data_dir', type=str, default='./data/eth_ucy', help='Directory where the data is stored.') parser.add_argument('--n_train', type=int, default=None, help='Override the number training scenes used.') parser.add_argument('--n_test', type=int, default=None, help='Override the number testing scenes used.') parser.add_argument('--checkpt_freq', default=5, type=int, help='Override the checkpt_freq in the config file.') parser.add_argument('--max_num_ckpts', default=5, type=int, help='Override the max_num_ckpts in the config file.') parser.add_argument('--data_norm', default='min_max', choices=['min_max', 'original', 'sqrt'], help='Normalization method for the data.') parser.add_argument('--subset', type=str, required=True, default='eth', choices=['eth', 'hotel', 'univ', 'zara1', 'zara2'], help='Trajectory subset to run experiment') parser.add_argument('--rotate', default=False, action='store_true', help="Whether to rotate the trajectories in the dataset") parser.add_argument('--rotate_time_frame', type=int, default=0, help='Index of time frames to rotate the trajectories.') parser.add_argument('--rotate_aug', default=False, action='store_true', help='Whether to use rotation as data augmentation.') # 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.') ### FM parameters ### parser.add_argument('--sampling_steps', type=int, default=10, help='Number of sampling timesteps for the FlowMatcher.') # time scheduler during training parser.add_argument('--t_schedule', type=str, choices=['uniform', 'logit_normal'], default='logit_normal', help='Time schedule for the FlowMatcher.') parser.add_argument('--fm_skewed_t', default=None, type=str, help='Skewed time schedule for the FlowMatcher.') parser.add_argument('--logit_norm_mean', default=-0.5, type=float, help='Mean for the logit normal distribution.') parser.add_argument('--logit_norm_std', default=1.5, type=float, help='Standard deviation for the logit normal distribution.') parser.add_argument('--fm_wrapper', type=str, default='direct', choices=['direct', 'velocity', 'precond'], help='Wrapper for the FlowMatcher.') parser.add_argument('--fm_rew_sqrt', default=False, action='store_true', help='Whether to apply square root to the reweighting factor.') parser.add_argument('--fm_in_scaling', default=False, action='store_true', help='Whether to scale the input to the FlowMatcher.') parser.add_argument('--sigma_data', type=float, default=0.13, help='Standard deviation for the data.') # input dropout / masking rate parser.add_argument('--drop_method', default='emb', type=str, choices=['None', 'input', 'emb'], help='Dropout method for the FlowMatcher.') parser.add_argument('--drop_logi_k', default=20.0, type=float, help='Logistic growth rate for masking rate at different timesteps.') parser.add_argument('--drop_logi_m', default=0.5, type=float, help='Logistic midpoint for masking rate at different timesteps.') ### FM parameters ### ### Architecture configuration ### parser.add_argument('--use_pre_norm', default=False, action='store_true', help='Where to normalize the input trajectories in the Transformer Encoders.') parser.add_argument('--num_layers', type=int, default=None, help='Overwrite the number of layers in the config file.') parser.add_argument('--dropout', default=None, type=float, help='Overwrite the dropout rate in the config file.') ### Architecture configuration ### ### General denoising objective configuration ### parser.add_argument('--tied_noise', default=True, action='store_true', help='Whether to use tied noise for the denoiser.') ### General denoising objective configuration ### ### Regression loss configuration ### parser.add_argument('--loss_nn_mode', type=str, default='agent', choices=['agent', 'scene'], help='Whether to use the agent-wise or scene-wise NN loss.') parser.add_argument('--loss_reg_reduction', type=str, default='sum', choices=['mean', 'sum'], help='Reduction method for the regression loss.') parser.add_argument('--loss_reg_squared', default=False, action='store_true', help='Whether to use the squared regression loss.') parser.add_argument('--loss_velocity', default=False, action='store_true', help='Whether to use the regression loss for velocity.') ### Regression loss configuration ### ### Classification loss configuration ### parser.add_argument('--loss_cls_weight', type=float, default=1.0, help='Weight for the classification loss.') ### Classification loss configuration ### ### Optimization configuration ### parser.add_argument('--init_lr', type=float, default=None, help='Override the peak learning rate in the config file.') parser.add_argument('--weight_decay', type=float, default=None, help='Override the weight decay in the config file.') ### Optimization configuration ### return parser.parse_args() def init_basics(args): """ Init the basic configurations for the experiment. """ """Load the config file""" cfg = Config(args.cfg, f'{args.exp}') tag = '_' ### Update FM parameters ### def _update_fm_params(args, cfg, tag): if cfg.denoising_method == 'fm': cfg.sigma_data = args.sigma_data cfg.sampling_steps = args.sampling_steps if args.fm_skewed_t is not None: cfg.t_schedule = args.fm_skewed_t else: cfg.t_schedule = args.t_schedule if args.t_schedule == 'logit_normal': cfg.logit_norm_mean = args.logit_norm_mean cfg.logit_norm_std = args.logit_norm_std cfg.fm_wrapper = args.fm_wrapper cfg.fm_rew_sqrt = args.fm_rew_sqrt cfg.fm_in_scaling = args.fm_in_scaling if args.fm_skewed_t is not None: tag += f'FM_S{cfg.sampling_steps}_{cfg.t_schedule}_{cfg.fm_wrapper[:4]}' elif args.t_schedule == 'logit_normal': tag += f'FM_S{cfg.sampling_steps}_{cfg.t_schedule[:3]}_m{cfg.logit_norm_mean}_s{cfg.logit_norm_std}_{cfg.fm_wrapper[:4]}' elif args.t_schedule == 'uniform': tag += f'FM_S{cfg.sampling_steps}_{cfg.t_schedule[:3]}_{cfg.fm_wrapper[:4]}' if args.drop_method is not None and args.drop_logi_k is not None and args.drop_logi_m is not None: cfg.drop_method = args.drop_method cfg.drop_logi_k = args.drop_logi_k cfg.drop_logi_m = args.drop_logi_m tag += f'_drop_{cfg.drop_method}_m{cfg.drop_logi_m}_k{cfg.drop_logi_k}' if cfg.fm_rew_sqrt: tag += '_RESQ' if cfg.fm_in_scaling: tag += '_IS' return cfg, tag cfg, tag = _update_fm_params(args, cfg, tag) ### Architecture configuration ### def _update_architecture_params(args, cfg, tag): cfg.MODEL.USE_PRE_NORM = args.use_pre_norm cfg.MODEL.NUM_LAYERS = args.num_layers cfg.MODEL.DROPOUT = args.dropout if args.num_layers is not None: tag += f'_L{args.num_layers}' cfg.MODEL.CONTEXT_ENCODER.NUM_ATTN_LAYERS = args.num_layers cfg.MODEL.MOTION_DECODER.NUM_DECODER_BLOCKS = args.num_layers if args.dropout is not None: tag += f'_DO{args.dropout}' cfg.MODEL.CONTEXT_ENCODER.DROPOUT_OF_ATTN = args.dropout cfg.MODEL.MOTION_DECODER.DROPOUT_OF_ATTN = args.dropout return cfg, tag cfg, tag = _update_architecture_params(args, cfg, tag) ### General denoising objective configuration ### def _update_general_denoising_params(args, cfg, tag): cfg.tied_noise = args.tied_noise if args.tied_noise: tag += '_TN' return cfg, tag cfg, tag = _update_general_denoising_params(args, cfg, tag) ### Regression loss configuration ### def _update_regression_loss_params(args, cfg, tag): cfg.LOSS_NN_MODE = args.loss_nn_mode cfg.LOSS_REG_REDUCTION = args.loss_reg_reduction cfg.LOSS_REG_SQUARED = args.loss_reg_squared cfg.LOSS_VELOCITY = args.loss_velocity tag += f'_NN_{cfg.LOSS_NN_MODE[:1].upper()}' tag += f'_REG_{cfg.LOSS_REG_REDUCTION[:1].upper()}' if args.loss_reg_squared: tag += '_SQ' if args.loss_velocity: tag += '_VEL' cfg.MODEL.REGRESSION_MLPS[-1] += cfg.MODEL.MODEL_OUT_DIM return cfg, tag cfg, tag = _update_regression_loss_params(args, cfg, tag) ### Update data configuration ### def _update_data_params(args, cfg, tag): cfg.subset = args.subset tag += f'_{cfg.subset}' cfg.rotate = args.rotate if args.rotate: cfg.rotate_time_frame = args.rotate_time_frame cfg.rotate_aug = args.rotate_aug tag += f'_rot_{cfg.rotate_time_frame}' if cfg.rotate_aug: tag += '_aug' if args.n_train is not None: tag += f'_subset_train_{args.n_train}' if args.n_test is not None: tag += f'_test{args.n_test}' cfg.data_norm = args.data_norm tag += f'_{args.data_norm}' return cfg, tag cfg, tag = _update_data_params(args, cfg, tag) ### Update optimization configs ### def _update_optimization_params(args, cfg, tag): if args.init_lr is not None: cfg.OPTIMIZATION.LR = args.init_lr if args.weight_decay is not None: cfg.OPTIMIZATION.WEIGHT_DECAY = args.weight_decay cfg.OPTIMIZATION.LOSS_WEIGHTS['cls'] = args.loss_cls_weight tag += f'_LR{cfg.OPTIMIZATION.LR}_WD{cfg.OPTIMIZATION.WEIGHT_DECAY}_CLS_{args.loss_cls_weight}' if args.epochs is not None: # override the number of epochs cfg.OPTIMIZATION.NUM_EPOCHS = args.epochs 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 * 2 # larger BS for during-training evaluation if args.checkpt_freq is not None: # override the checkpt_freq cfg.checkpt_freq = args.checkpt_freq cfg.max_num_ckpts = args.max_num_ckpts tag += f'_BS{cfg.train_batch_size}_EP{cfg.OPTIMIZATION.NUM_EPOCHS}' return cfg, tag cfg, tag = _update_optimization_params(args, cfg, tag) ### viola, create the saving directory ### 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')) os.makedirs(tb_dir, exist_ok=True) tb_log = SummaryWriter(log_dir=tb_dir) """back up the code""" back_up_code_git(cfg, logger=logger) """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 ETH dataset. """ train_dset = ETHDataset( cfg=cfg, training=True, data_dir=args.data_dir, subset=args.subset, rotate_time_frame=args.rotate_time_frame, type = args.data_source) train_loader = DataLoader( train_dset, batch_size=cfg.train_batch_size, shuffle=True, num_workers=4, collate_fn=seq_collate_eth, pin_memory=True) test_dset = ETHDataset( cfg=cfg, training=False, data_dir=args.data_dir, subset=args.subset, rotate_time_frame=args.rotate_time_frame, type = args.data_source) test_loader = DataLoader( test_dset, batch_size=cfg.test_batch_size, ### change it from 500 shuffle=False, num_workers=4, collate_fn=seq_collate_eth, pin_memory=True) return train_loader, test_loader def build_network(cfg, args, logger): """ Build the network for the denoising model. """ model = ETHMotionTransformer( model_config=cfg.MODEL, logger=logger, config=cfg, ) if cfg.denoising_method == 'fm': denoiser = FlowMatcher( cfg, model, logger=logger, ) else: raise NotImplementedError(f'Denoising method [{cfg.denoising_method}] is not implemented yet.') return denoiser 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) denoiser = build_network(cfg, args, logger) """Train the model""" trainer = Trainer( cfg, denoiser, train_loader, test_loader, tb_log=tb_log, logger=logger, gradient_accumulate_every=1, ema_decay = 0.995, ema_update_every = 1, ) trainer.train() if __name__ == "__main__": main()