def load_net(runner, net, model_path=None): if model_path is not None or runner.configer.get('network', 'resume') is not None: resume_path = runner.configer.get('network', 'resume') resume_path = model_path if model_path is not None else resume_path if not os.path.exists(resume_path): Log.warn('Resume path: {} not exists...'.format(resume_path)) return net Log.info('Resuming from {}'.format(resume_path)) resume_dict = torch.load(resume_path, map_location="cpu") if 'state_dict' in resume_dict: checkpoint_dict = resume_dict['state_dict'] elif 'model' in resume_dict: checkpoint_dict = resume_dict['model'] elif isinstance(resume_dict, OrderedDict): checkpoint_dict = resume_dict else: raise RuntimeError( 'No state_dict found in checkpoint file {}'.format(runner.configer.get('network', 'resume'))) # load state_dict if hasattr(net, 'module'): RunnerHelper.load_state_dict(net.module, checkpoint_dict, runner.configer.get('network', 'resume_strict')) else: RunnerHelper.load_state_dict(net, checkpoint_dict, runner.configer.get('network', 'resume_strict')) if runner.configer.get('network', 'resume_continue'): runner.runner_state = resume_dict['runner_state'] return net
def __getitem__(self, index): img = None valid = True while img is None: try: img = ImageHelper.read_image( self.img_list[index], tool=self.configer.get('data', 'image_tool'), mode=self.configer.get('data', 'input_mode')) assert isinstance(img, np.ndarray) or isinstance( img, Image.Image) except: Log.warn('Invalid image path: {}'.format(self.img_list[index])) img = None valid = False index = (index + 1) % len(self.img_list) label = torch.from_numpy(np.array(self.label_list[index])) if self.aug_transform is not None: img = self.aug_transform(img) if self.img_transform is not None: img = self.img_transform(img) return dict(valid=valid, img=DataContainer(img, stack=True), label=DataContainer(label, stack=True))
def __read_file(self, data_dir, dataset): img_list = list() mlabel_list = list() img_dict = dict() all_img_list = [] with open(self.configer.get('data.{}_label_path'.format(dataset)), 'r') as file_stream: all_img_list += file_stream.readlines() if dataset == 'train' and self.configer.get('data.include_val', default=False): with open(self.configer.get('data.val_label_path'), 'r') as file_stream: all_img_list += file_stream.readlines() for line_cnt in range(len(all_img_list)): line_items = all_img_list[line_cnt].strip().split() if len(line_items) == 0: continue path = line_items[0] if not os.path.exists(os.path.join( data_dir, path)) or not ImageHelper.is_img(path): Log.warn('Invalid Image Path: {}'.format( os.path.join(data_dir, path))) continue img_list.append(os.path.join(data_dir, path)) mlabel_list.append([int(item) for item in line_items[1:]]) assert len(img_list) > 0 Log.info('Length of {} imgs is {}...'.format(dataset, len(img_list))) return img_list, mlabel_list
def __getitem__(self, index): img = None valid = True while img is None: try: img = ImageHelper.read_image(self.item_list[index][0], tool=self.configer.get('data', 'image_tool'), mode=self.configer.get('data', 'input_mode')) assert isinstance(img, np.ndarray) or isinstance(img, Image.Image) except: Log.warn('Invalid image path: {}'.format(self.item_list[index][0])) img = None valid = False index = (index + 1) % len(self.item_list) ori_img_size = ImageHelper.get_size(img) if self.aug_transform is not None: img = self.aug_transform(img) border_hw = ImageHelper.get_size(img)[::-1] if self.img_transform is not None: img = self.img_transform(img) meta = dict( valid=valid, ori_img_size=ori_img_size, border_hw=border_hw, img_path=self.item_list[index][0], filename=self.item_list[index][1], label=self.item_list[index][2] ) return dict( img=DataContainer(img, stack=True), meta=DataContainer(meta, stack=False, cpu_only=True) )
def __init__(self, max_degree, rotate_ratio=0.5, mean=(104, 117, 123)): assert isinstance(max_degree, int) self.max_degree = max_degree self.ratio = rotate_ratio self.mean = mean Log.warn( 'Currently `RandomRotate` is only implemented for `img`, `labelmap` and `maskmap`.' )
def _set_value(key, value): """ We directly operate on `params_root`. """ remained_parts = key.split('.') consumed_parts = [] parent_dict = self.params_root while len(remained_parts) > 1: cur_key = remained_parts.pop(0) consumed_parts.append(cur_key) if cur_key not in parent_dict: parent_dict[cur_key] = dict() Log.info('{} not exists, set as `dict()`.'.format( '.'.join(consumed_parts))) elif not isinstance(parent_dict[cur_key], dict): Log.error( 'Cannot set {child_name} on {root_name}, as {root_name} is `{root_type}`.' .format(root_name='.'.join(consumed_parts), child_name='.'.join(remained_parts), root_type=type(parent_dict[cur_key]))) sys.exit(1) parent_dict = parent_dict[cur_key] cur_key = remained_parts.pop(0) consumed_parts.append(cur_key) if cur_key.endswith('+'): cur_key = cur_key[:-1] target = parent_dict.get(cur_key) if not isinstance(target, list): Log.error( 'Cannot append to {key}, as its type is {target_type}.' .format(key=key[:-1], target_type=type(target))) sys.exit(1) target.append(value) Log.info('Append {value} to {key}. Current: {target}.'.format( key=key[:-1], value=value, target=target, )) return existing_value = parent_dict.get(cur_key) if existing_value is not None: Log.warn( 'Override {key} using {value}. Previous value: {old_value}.' .format(key=key, value=value, old_value=existing_value)) else: Log.info('Set {key} as {value}.'.format(key=key, value=value)) parent_dict[cur_key] = value
def get(self, *key, **kwargs): key = '.'.join(key) if key not in self.params_root: Log.warn('Key: {} not exists'.format(key)) if key in self.params_root or 'default' in kwargs: return self.params_root.get(key, **kwargs) else: Log.error('{} KeyError: {}.'.format(self._get_caller(), key)) exit(1)
def load_state_dict(module, state_dict, strict=False): """Load state_dict to a module. This method is modified from :meth:`torch.nn.Module.load_state_dict`. Default value for ``strict`` is set to ``False`` and the message for param mismatch will be shown even if strict is False. Args: module (Module): Module that receives the state_dict. state_dict (OrderedDict): Weights. strict (bool): whether to strictly enforce that the keys in :attr:`state_dict` match the keys returned by this module's :meth:`~torch.nn.Module.state_dict` function. Default: ``False``. """ if list(state_dict.keys())[0].startswith('module.'): state_dict = {k[7:]: v for k, v in state_dict.items()} unexpected_keys = [] unmatched_keys = [] own_state = module.state_dict() for name, param in state_dict.items(): if name not in own_state: unexpected_keys.append(name) continue if isinstance(param, torch.nn.Parameter): # backwards compatibility for serialized parameters param = param.data try: own_state[name].copy_(param) except Exception: if strict: raise RuntimeError('While copying the parameter named {}, ' 'whose dimensions in the model are {} and ' 'whose dimensions in the checkpoint are {}.' .format(name, own_state[name].size(), param.size())) else: unmatched_keys.append(name) missing_keys = set(own_state.keys()) - set(state_dict.keys()) err_msg = [] if unexpected_keys: err_msg.append('unexpected key in source state_dict: {}'.format(', '.join(unexpected_keys))) if missing_keys: err_msg.append('missing keys in source state_dict: {}'.format(', '.join(missing_keys))) if unmatched_keys: err_msg.append('unmatched keys in source state_dict: {}'.format(', '.join(unmatched_keys))) err_msg = '\n'.join(err_msg) if err_msg: if strict: raise RuntimeError(err_msg) else: Log.warn(err_msg)
def update_performance(self): try: rs = self.running_scores[self.save_net_main_key] if self.save_net_metric == 'miou': perf = rs.get_mean_iou() elif self.save_net_metric == 'acc': perf = rs.get_pixel_acc() max_perf = self.configer.get('max_performance') self.configer.update(['performance'], perf) if perf > max_perf: Log.info('Performance {} -> {}'.format(max_perf, perf)) except Exception as e: Log.warn(e)
def __read_file(self, root_dir, dataset, label_path): img_list = list() mlabel_list = list() with open(label_path, 'r') as file_stream: for line in file_stream.readlines(): line_items = line.rstrip().split() path = line_items[0] if not os.path.exists(os.path.join(root_dir, path)) or not ImageHelper.is_img(path): Log.warn('Invalid Image Path: {}'.format(os.path.join(root_dir, path))) continue img_list.append(os.path.join(root_dir, path)) mlabel_list.append([int(item) for item in line_items[1:]]) assert len(img_list) > 0 Log.info('Length of {} imgs is {}...'.format(dataset, len(img_list))) return img_list, mlabel_list
def __init__(self, args_parser=None, config_file=None, config_dict=None, valid_flag=None): self.params_root = None if config_dict is not None: assert config_file is None self.params_root = ConfigFactory.from_dict(config_dict) elif config_file is not None: if not os.path.exists(config_file): Log.error('Json Path:{} not exists!'.format(config_file)) exit(1) self.params_root = ConfigFactory.parse_file(config_file) elif 'config_file' in args_parser and args_parser.config_file is not None: if not os.path.exists(args_parser.config_file): Log.error('Json Path:{} not exists!'.format( args_parser.config_file)) exit(1) self.params_root = ConfigFactory.parse_file( args_parser.config_file) else: Log.warn('Base settings not set!') self.params_root = ConfigFactory.from_dict({}) if args_parser is not None: for key, value in args_parser.__dict__.items(): if valid_flag is not None and key.split('.')[0] != valid_flag: continue if key not in self.params_root: self.add(key, value) elif value is not None: self.update(key, value)
def __read_and_split_file(self, root_dir, dataset, label_path): img_list = list() mlabel_list = list() select_interval = int(1 / self.configer.get('data', 'val_ratio')) img_dict = dict() with open(label_path, 'r') as file_stream: for line in file_stream.readlines(): label = line.strip().split()[1] if int(label) in img_dict: img_dict[int(label)].append(line) else: img_dict[int(label)] = [line] all_img_list = [] for i in sorted(img_dict.keys()): all_img_list += img_dict[i] for line_cnt in range(len(all_img_list)): if line_cnt % select_interval == 0 and dataset == 'train' and not self.configer.get('data', 'include_val'): continue if line_cnt % select_interval != 0 and dataset == 'val': continue line_items = all_img_list[line_cnt].strip().split() path = line_items[0] if not os.path.exists(os.path.join(root_dir, path)) or not ImageHelper.is_img(path): Log.warn('Invalid Image Path: {}'.format(os.path.join(root_dir, path))) continue img_list.append(os.path.join(root_dir, path)) mlabel_list.append([int(item) for item in line_items[1:]]) assert len(img_list) > 0 Log.info('Length of {} imgs is {} after split trainval...'.format(dataset, len(img_list))) return img_list, mlabel_list
def load_tf_efficientnet_model(model, pretrained=None, strict=False): if pretrained is None: return model if len(glob.glob(pretrained + '*')) == 0 or not os.path.exists( glob.glob(pretrained + '*')[0]): Log.info('{} not exists.'.format(pretrained)) return model import pdb from tensorflow.python import pywrap_tensorflow Log.info('Loading pretrained model:{}'.format(pretrained)) tf_reader = pywrap_tensorflow.NewCheckpointReader(pretrained) for tf_key in tf_reader.get_variable_to_shape_map(): if tf_key.startswith('efficientnet'): model_name = tf_key.split('/')[0] break model_dict = model.state_dict() load_dict = dict() bn_list = [['running_mean', 'moving_mean'], ['running_var', 'moving_variance'], ['weight', 'gamma'], ['bias', 'beta']] # stem and head block for block_name in ['stem', 'head']: key = '_conv_{}.weight'.format(block_name) tf_key = '{}/{}/conv2d/kernel'.format(model_name, block_name) assert (key in model_dict) tf_value = tf_reader.get_tensor(tf_key) load_dict[key] = torch.from_numpy(tf_value.transpose(3, 2, 0, 1)) for bn_name in bn_list: key = '_bn_{}.{}'.format(block_name, bn_name[0]) tf_key = '{}/{}/tpu_batch_normalization/{}'.format( model_name, block_name, bn_name[1]) assert (key in model_dict) tf_value = tf_reader.get_tensor(tf_key) load_dict[key] = torch.from_numpy(tf_value) # MBConvBlocks module_list = [[['_expand_conv', '_project_conv'], 'conv2d'], [['_depthwise_conv'], 'depthwise_conv2d'], [['_se_reduce', '_se_expand'], 'se/conv2d'], [['_expand_bn', '_depthwise_bn', '_project_bn'], 'tpu_batch_normalization']] mb_block_id = 0 mb_block_flag = True while mb_block_flag: for module in module_list: key_id = 0 tf_key_part = module[1] for key_part in module[0]: if key_part.startswith('_depthwise_conv'): key = '_blocks.{}.{}.weight'.format( mb_block_id, key_part) tf_key = '{}/blocks_{}/{}/depthwise_kernel'.format( model_name, mb_block_id, tf_key_part) if key not in model_dict: mb_block_flag = False Log.info('Ignore parameter: {} <---> {}'.format( key, tf_key)) continue tf_value = tf_reader.get_tensor(tf_key) load_dict[key] = torch.from_numpy( tf_value.transpose(2, 3, 0, 1)) elif key_part.endswith('_conv'): key = '_blocks.{}.{}.weight'.format( mb_block_id, key_part) if key_id > 0: tf_key = '{}_{}'.format(tf_key_part, key_id) else: tf_key = '{}'.format(tf_key_part) tf_key = '{}/blocks_{}/{}/kernel'.format( model_name, mb_block_id, tf_key) if key not in model_dict: Log.info('Ignore parameter: {} <---> {}'.format( key, tf_key)) continue tf_value = tf_reader.get_tensor(tf_key) load_dict[key] = torch.from_numpy( tf_value.transpose(3, 2, 0, 1)) key_id += 1 elif key_part.startswith('_se'): key = '_blocks.{}.{}.weight'.format( mb_block_id, key_part) if key_id > 0: tf_key = '{}_{}'.format(tf_key_part, key_id) else: tf_key = '{}'.format(tf_key_part) tf_key = '{}/blocks_{}/{}/kernel'.format( model_name, mb_block_id, tf_key) if key not in model_dict: Log.info('Ignore parameter: {} <---> {}'.format( key, tf_key)) continue tf_value = tf_reader.get_tensor(tf_key) load_dict[key] = torch.from_numpy( tf_value.transpose(3, 2, 0, 1)) key = '_blocks.{}.{}.bias'.format( mb_block_id, key_part) if key_id > 0: tf_key = '{}_{}'.format(tf_key_part, key_id) else: tf_key = '{}'.format(tf_key_part) tf_key = '{}/blocks_{}/{}/bias'.format( model_name, mb_block_id, tf_key) tf_value = tf_reader.get_tensor(tf_key) load_dict[key] = torch.from_numpy(tf_value) key_id += 1 elif key_part.endswith('_bn'): bn_flag = False for bn_name in bn_list: key = '_blocks.{}.{}.{}'.format( mb_block_id, key_part, bn_name[0]) if key_id > 0: tf_key = '{}_{}'.format(tf_key_part, key_id) else: tf_key = '{}'.format(tf_key_part) tf_key = '{}/blocks_{}/{}/{}'.format( model_name, mb_block_id, tf_key, bn_name[1]) if key not in model_dict: Log.info( 'Ignore parameter: {} <---> {}'.format( key, tf_key)) continue bn_flag = True tf_value = tf_reader.get_tensor(tf_key) load_dict[key] = torch.from_numpy(tf_value) if bn_flag: key_id += 1 mb_block_id += 1 unexpected_keys = [] unmatched_keys = [] own_state = model.state_dict() for name, param in load_dict.items(): if name not in own_state: unexpected_keys.append(name) continue if isinstance(param, torch.nn.Parameter): # backwards compatibility for serialized parameters param = param.data try: own_state[name].copy_(param) except Exception: if strict: raise RuntimeError( 'While copying the parameter named {}, ' 'whose dimensions in the model are {} and ' 'whose dimensions in the checkpoint are {}.'.format( name, own_state[name].size(), param.size())) else: unmatched_keys.append(name) missing_keys = set(own_state.keys()) - set(load_dict.keys()) err_msg = [] if unexpected_keys: err_msg.append('unexpected key in source state_dict: {}'.format( ', '.join(unexpected_keys))) if missing_keys: err_msg.append('missing keys in source state_dict: {}'.format( ', '.join(missing_keys))) if unmatched_keys: err_msg.append('unmatched keys in source state_dict: {}'.format( ', '.join(unmatched_keys))) err_msg = '\n'.join(err_msg) if err_msg: if strict: raise RuntimeError(err_msg) else: Log.warn(err_msg) return model