()
| 18 | import cv2 |
| 19 | |
| 20 | def main(): |
| 21 | #### options |
| 22 | parser = argparse.ArgumentParser() |
| 23 | parser.add_argument('--opt', type=str, default='./options/train/test_Enhance_MSEC.yml', |
| 24 | help='Path to option YAML file.') |
| 25 | parser.add_argument('--launcher', choices=['none', 'pytorch'], default='pytorch', help='job launcher') |
| 26 | parser.add_argument('--local_rank', type=int, default=0) |
| 27 | args = parser.parse_args() |
| 28 | opt = option.parse(args.opt, is_train=True) |
| 29 | |
| 30 | #### distributed training settings |
| 31 | |
| 32 | opt['dist'] = False |
| 33 | rank = -1 |
| 34 | print('Disabled distributed training.') |
| 35 | |
| 36 | |
| 37 | #### loading resume state if exists |
| 38 | if opt['path'].get('resume_state', None): |
| 39 | # distributed resuming: all load into default GPU |
| 40 | device_id = torch.cuda.current_device() |
| 41 | resume_state = torch.load(opt['path']['resume_state'], |
| 42 | map_location=lambda storage, loc: storage.cuda(device_id)) |
| 43 | option.check_resume(opt, resume_state['iter']) # check resume options |
| 44 | else: |
| 45 | resume_state = None |
| 46 | |
| 47 | #### mkdir and loggers |
| 48 | if rank <= 0: # normal training (rank -1) OR distributed training (rank 0) |
| 49 | if resume_state is None: |
| 50 | util.mkdir_and_rename( |
| 51 | opt['path']['experiments_root']) # rename experiment folder if exists |
| 52 | util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' |
| 53 | and 'pretrain_model' not in key and 'resume' not in key)) |
| 54 | |
| 55 | # config loggers. Before it, the log will not work |
| 56 | util.setup_logger('base', opt['path']['log'], 'train_' + opt['name'], level=logging.INFO, |
| 57 | screen=True, tofile=True) |
| 58 | logger = logging.getLogger('base') |
| 59 | logger.info(option.dict2str(opt)) |
| 60 | # tensorboard logger |
| 61 | if opt['use_tb_logger'] and 'debug' not in opt['name']: |
| 62 | version = float(torch.__version__[0:3]) |
| 63 | if version >= 1.1: # PyTorch 1.1 |
| 64 | from torch.utils.tensorboard import SummaryWriter |
| 65 | else: |
| 66 | logger.info( |
| 67 | 'You are using PyTorch {}. Tensorboard will use [tensorboardX]'.format(version)) |
| 68 | from tensorboardX import SummaryWriter |
| 69 | tb_logger = SummaryWriter(log_dir=(os.path.join(opt['path']['root'],'tb_logger',opt['name']))) |
| 70 | |
| 71 | else: |
| 72 | util.setup_logger('base', opt['path']['log'], 'train', level=logging.INFO, screen=True) |
| 73 | logger = logging.getLogger('base') |
| 74 | |
| 75 | |
| 76 | |
| 77 | # convert to NoneDict, which returns None for missing keys |
no test coverage detected