Beispiel #1
0
def main():
    parser = argparse.ArgumentParser(
        description="PyTorch MaskPyramid Training for Cityscapes")
    parser.add_argument(  # args.config
        "--config",
        default="maskpyramid/config/gallery/MaskPyramidV1.yml",
        metavar="FILE",
        help="path to config file",
        type=str,
    )
    parser.add_argument(  # args.opts
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )
    args = parser.parse_args()

    cfg.merge_from_file(args.config)
    # cfg.merge_from_list(['DATALOADER.BATCH_SIZE_TRAIN', 4])
    cfg.merge_from_list(['DATALOADER.NUM_WORKERS', 0])
    # cfg.merge_from_list(['MODEL.WEIGHT', ''])
    cfg.merge_from_list(args.opts)
    cfg.freeze()

    experiment_name = os.path.basename(__file__).split('.py')[0]
    if cfg.MULTI_RUN:
        run_names = glob.glob(cfg.OUTPUT_DIR + '/*')
        matched_num = ''
        for name in run_names:
            matched_num = name.split(
                os.path.join(cfg.OUTPUT_DIR, experiment_name) + '_')[-1]
        experiment_name += '_' + str(int(matched_num) +
                                     1) if matched_num else '_1'
    output_dir = os.path.join(cfg.OUTPUT_DIR, experiment_name)

    mkdir(output_dir)
    with open(os.path.join(output_dir, experiment_name + '_config.yml'),
              'w') as f:
        f.write(cfg.dump())
    logger = setup_logger("MaskPyramid", output_dir, 0,
                          experiment_name + "_log.txt")  # get_rank())
    logger.info(args)
    logger.info("Collecting env info (might take some time)")
    logger.info("\n" + collect_env_info())

    logger.info("Loaded and archive configuration file at '{}'".format(
        os.path.join(output_dir, experiment_name + '_config.yml')))

    trainer = Trainer(cfg, output_dir)
    logger.info('Starting Epoch:{} / Total Epoches:{}'.format(
        trainer.start_epoch, cfg.SOLVER.EPOCHES))

    for epoch in range(trainer.start_epoch, cfg.SOLVER.EPOCHES):
        trainer.training(epoch)
        trainer.validation(epoch)

    import pdb
    pdb.set_trace()
    trainer.writer.close()
Beispiel #2
0
def main():
    parser = argparse.ArgumentParser(description="PyTorch MaskPyramid Training for Cityscapes")
    parser.add_argument(        # args.config
        "--config",
        default="hrnet/hrnet.yml",
        metavar="FILE",
        help="path to config file",
        type=str,
    )
    parser.add_argument(        # args.opts
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )
    args = parser.parse_args()

    cfg.merge_from_file(args.config)
    # # cfg.merge_from_list(['SOLVER.SHOW_IMAGE', False])
    # cfg.merge_from_list(['DATALOADER.BATCH_SIZE_TRAIN', 2])
    # # cfg.merge_from_list(['DATALOADER.NUM_WORKERS', 0])
    # cfg.merge_from_list(['SOLVER.SEMATIC_ONLY', True])
    # cfg.merge_from_list(['SOLVER.BASE_LR', 1e-7])
    # cfg.merge_from_list(['TEST_ONLY', True])
    cfg.merge_from_list(['MULTI_RUN', True])
    # cfg.merge_from_list(['MODEL.WEIGHT', 'run/gm_v4_sematic_21/Every_5_model_Epoch_5.pth'])
    cfg.merge_from_list(args.opts)
    cfg.freeze()

    extra_exp_name = 'sematic' if cfg.SOLVER.SEMATIC_ONLY else ''
    experiment_name = os.path.basename(__file__).split('.py')[0]
    experiment_name += '_'+extra_exp_name if extra_exp_name else ''
    if cfg.MULTI_RUN:
        run_names = glob.glob(cfg.OUTPUT_DIR + '/*')
        matched_nums = []
        for name in run_names:
            matched_nums += re.findall(r'%s_(\d+)'%experiment_name, name)
        experiment_name += '_' + str(max([int(n) for n in matched_nums])+1) if matched_nums else '_1'
    output_dir = os.path.join(cfg.OUTPUT_DIR, experiment_name)
    
    mkdir(output_dir)
    with open(os.path.join(output_dir, experiment_name+'_config.yml'), 'w') as f:
        f.write(cfg.dump())
    logger = setup_logger("MaskPyramid", output_dir, 0, experiment_name+"_log.txt")  # get_rank())
    logger.info(args)
    logger.info("Collecting env info (might take some time)")
    logger.info("\n" + collect_env_info())
    logger.info("Loaded and archive configuration file at '{}'".format(os.path.join(output_dir, experiment_name+'_config.yml')))
    log_gpu_stat(logger)


    # cudnn related setting
    cudnn.benchmark = cfg.CUDNN.BENCHMARK
    cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
    cudnn.enabled = cfg.CUDNN.ENABLED

    trainer = Trainer(cfg, output_dir)
    logger.info('Starting Epoch:{} / Total Epoches:{}'.format(trainer.start_epoch, cfg.SOLVER.EPOCHES))
    if cfg.MODEL.WEIGHT:
        trainer.load_weights(cfg.MODEL.WEIGHT, subdict='hrnet')

    if cfg.TEST_ONLY:
        assert cfg.MODEL.WEIGHT, "You MUST provide a weight file for the testing mode!"
        trainer.validation(0)
    else:
        for epoch in range(trainer.start_epoch, cfg.SOLVER.EPOCHES):
            trainer.training(epoch)
            trainer.validation(epoch)
            log_gpu_stat(logger)
            if epoch == 5 or epoch % 20 == 0:
                save_data = {}
                save_data["epoch"] = epoch + 1
                save_data["best_pred"] = -1
                save_data["model"] = trainer.model.state_dict()
                if trainer.optimizer is not None:
                    save_data["optimizer"] = trainer.optimizer.state_dict()
                if trainer.scheduler is not None:
                    save_data["scheduler"] = trainer.scheduler.__dict__ 
                torch.save(save_data, trainer.output_dir+"/Every_5_model_Epoch_{}.pth".format(epoch))



    trainer.writer.close()