datnguyentien204/CoMoGAN_Modified

1import time2import os3os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # Disables tensorflow loggings4 5from options import get_options6from data import create_dataset7from networks import create_model, get_model_options8from argparse import ArgumentParser as AP9 10import pytorch_lightning as pl11from pytorch_lightning.loggers import TensorBoardLogger12 13from util.callbacks import LogAndCheckpointEveryNSteps14from human_id import generate_id15 16def start(cmdline):17 18    pl.trainer.seed_everything(cmdline.seed)19    opt = get_options(cmdline)20 21    dataset = create_dataset(opt)  # create a dataset given opt.dataset_mode and other options22    model = create_model(opt)      # create a model given opt.model and other options23 24    callbacks = []25 26    logger = None27    if not cmdline.debug:28        root_dir = os.path.join('logs/', generate_id()) if cmdline.id == None else os.path.join('logs/', cmdline.id)29        logger = TensorBoardLogger(save_dir=os.path.join(root_dir, 'tensorboard'))30        logger.log_hyperparams(opt)31        callbacks.append(LogAndCheckpointEveryNSteps(save_step_frequency=opt.save_latest_freq,32                                                     viz_frequency=opt.display_freq,33                                                     log_frequency=opt.print_freq))34    else:35        root_dir = os.path.join('/tmp', generate_id())36 37    precision = 16 if cmdline.mixed_precision else 3238 39    trainer = pl.Trainer(default_root_dir=os.path.join(root_dir, 'checkpoints'), callbacks=callbacks,40                         gpus=cmdline.gpus, logger=logger, precision=precision, amp_level='01')41    trainer.fit(model, dataset)42 43 44if __name__ == '__main__':45    ap = AP()46    ap.add_argument('--id', default=None, type=str, help='Set an existing uuid to resume a training')47    ap.add_argument('--debug', default=False, action='store_true', help='Disables experiment saving')48    ap.add_argument('--gpus', default=[0], type=int, nargs='+', help='gpus to train on')49    ap.add_argument('--model', default='comomunit', type=str, help='Choose model for training')50    ap.add_argument('--data_importer', default='day2timelapse', type=str, help='Module name of the dataset importer')51    ap.add_argument('--path_data', default='/datasets/waymo_comogan/train/', type=str, help='Path to the dataset')52    ap.add_argument('--learning_rate', default=0.0001, type=float, help='Learning rate')53    ap.add_argument('--scheduler_policy', default='step', type=str, help='Scheduler policy')54    ap.add_argument('--decay_iters_step', default=200000, type=int, help='Decay iterations step')55    ap.add_argument('--decay_step_gamma', default=0.5, type=float, help='Decay step gamma')56    ap.add_argument('--seed', default=1, type=int, help='Random seed')57    ap.add_argument('--mixed_precision', default=False, action='store_true', help='Use mixed precision to reduce memory usage')58    start(ap.parse_args())59 60