| 7 | |
| 8 | |
| 9 | def parse(opt_path, is_train=True): |
| 10 | with open(opt_path, mode='r') as f: |
| 11 | opt = yaml.load(f, Loader=Loader) |
| 12 | # export CUDA_VISIBLE_DEVICES |
| 13 | gpu_list = ','.join(str(x) for x in opt['gpu_ids']) |
| 14 | os.environ['CUDA_VISIBLE_DEVICES'] = gpu_list |
| 15 | print('export CUDA_VISIBLE_DEVICES=' + gpu_list) |
| 16 | |
| 17 | opt['is_train'] = is_train |
| 18 | if opt['distortion'] == 'sr': |
| 19 | scale = opt['scale'] |
| 20 | |
| 21 | # datasets |
| 22 | for phase, dataset in opt['datasets'].items(): |
| 23 | phase = phase.split('_')[0] |
| 24 | dataset['phase'] = phase |
| 25 | if opt['distortion'] == 'sr': |
| 26 | dataset['scale'] = scale |
| 27 | is_lmdb = False |
| 28 | if dataset.get('dataroot_GT', None) is not None: |
| 29 | dataset['dataroot_GT'] = osp.expanduser(dataset['dataroot_GT']) |
| 30 | if dataset['dataroot_GT'].endswith('lmdb'): |
| 31 | is_lmdb = True |
| 32 | if dataset.get('dataroot_LQ', None) is not None: |
| 33 | dataset['dataroot_LQ'] = osp.expanduser(dataset['dataroot_LQ']) |
| 34 | if dataset['dataroot_LQ'].endswith('lmdb'): |
| 35 | is_lmdb = True |
| 36 | dataset['data_type'] = 'lmdb' if is_lmdb else 'img' |
| 37 | if dataset['mode'].endswith('mc'): # for memcached |
| 38 | dataset['data_type'] = 'mc' |
| 39 | dataset['mode'] = dataset['mode'].replace('_mc', '') |
| 40 | |
| 41 | # path |
| 42 | for key, path in opt['path'].items(): |
| 43 | if path and key in opt['path'] and key != 'strict_load': |
| 44 | opt['path'][key] = osp.expanduser(path) |
| 45 | # opt['path']['root'] = '/home/zhanghc/IEEEyellow/code_enhance/checkpoints' |
| 46 | if is_train: |
| 47 | experiments_root = osp.join(opt['path']['root'], 'experiments', opt['name']) |
| 48 | opt['path']['experiments_root'] = experiments_root |
| 49 | opt['path']['models'] = osp.join(experiments_root, 'models') |
| 50 | opt['path']['training_state'] = osp.join(experiments_root, 'training_state') |
| 51 | opt['path']['log'] = experiments_root |
| 52 | opt['path']['val_images'] = osp.join(experiments_root, 'val_images') |
| 53 | |
| 54 | # change some options for debug mode |
| 55 | if 'debug' in opt['name']: |
| 56 | opt['train']['val_freq'] = 8 |
| 57 | opt['logger']['print_freq'] = 1 |
| 58 | opt['logger']['save_checkpoint_freq'] = 8 |
| 59 | else: # test |
| 60 | results_root = osp.join(opt['path']['root'], 'results', opt['name']) |
| 61 | opt['path']['results_root'] = results_root |
| 62 | opt['path']['log'] = results_root |
| 63 | |
| 64 | # network |
| 65 | if opt['distortion'] == 'sr': |
| 66 | opt['network_G']['scale'] = scale |