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()
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()
