def __init__(self, desc=None, weight_file=None, pb_file=None): super(GraphGetter, self).__init__() if isinstance(desc, dict): src_model = ModelZoo().get_model(desc) else: src_model = desc weights = OrderedDict() if is_tf_backend(): import tensorflow.compat.v1 as tf from tensorflow.python.framework import tensor_util tf.reset_default_graph() data_shape = (1, 224, 224, 3) x = tf.ones(data_shape) if pb_file: with tf.io.gfile.GFile(pb_file, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) graph = tf.Graph() with graph.as_default(): tf.import_graph_def(graph_def, name='') weight_file = None wts = [n for n in graph_def.node if n.op == 'Const'] for n in wts: weights[n.name] = tensor_util.MakeNdarray( n.attr['value'].tensor) else: src_model(x, self.training) graph = tf.get_default_graph() desc = graph2desc(graph) tf.reset_default_graph() self.model = ModelZoo().get_model(desc, weight_file) if weights: self.model.load_checkpoint_from_numpy(weights)
def __init__(self, desc, from_graph=None, weight_file=None): super(Deformation, self).__init__() self._desc = {"props": deepcopy(desc.get('props')) or {}} if from_graph: self.model = GraphGetter(desc, weight_file).model else: self.model = ModelZoo().get_model(desc, weight_file) self._apply_names() self.get_search_space() self.decode() self.deform() self.props.clear()
def _init_model(self): """Load model desc from save path and parse to model.""" model = self.trainer.model if self.trainer.config.is_detection_trainer: model_desc = self.trainer.model_desc else: model_desc = self._get_model_desc() if model_desc: ModelConfig.model_desc = model_desc pretrained_model_file = self._get_pretrained_model_file() if not model: if not model_desc: raise Exception( "Failed to Init model, can not get model description.") model = ModelZoo.get_model(model_desc, pretrained_model_file) if model: if zeus.is_torch_backend(): import torch if self.trainer.use_cuda: model = model.cuda() if General._parallel and General.devices_per_trainer > 1: model = torch.nn.DataParallel(self.trainer.model) if zeus.is_tf_backend(): if pretrained_model_file: model_folder = os.path.dirname(pretrained_model_file) FileOps.copy_folder(model_folder, self.trainer.get_local_worker_path()) return model
def __init__(self, model_name=None, pretrained_model_file=None, downsamples=[4, 14, 27, 46], residauls=[7, 10, 17, 20, 23, 30, 33, 36, 39, 42, 49, 52], model=None, in_channels=None, out_channels=None, num_classes=None, **kwargs): super(PruneGetter, self).__init__() self.model = model if model_name is not None: model_desc = dict(type=model_name) model_desc.update(kwargs) self.model = ModelZoo().get_model(model_desc, pretrained_model_file) self.num_classes = num_classes convs = [ module for name, module in self.model.named_modules() if isinstance(module, torch.nn.Conv2d) ] self.in_channels_size = sum([conv.in_channels for conv in convs]) self.out_channels_size = sum([conv.out_channels for conv in convs]) if in_channels and len(in_channels) < self.in_channels_size: raise ValueError( "in channels mask length should be getter than {}".format( self.in_channels_size)) if out_channels and len(out_channels) < self.out_channels_size: raise ValueError( "out channels mask length should be getter than {}".format( self.out_channels_size)) in_channels_code = self.define_props('in_channels', in_channels) out_channels_code = self.define_props('out_channels', out_channels) if not in_channels_code and not out_channels_code: logging.info( "channels_code is null. use 1 as default, this model will not pruned." ) in_channels_code = [1] * self.in_channels_size out_channels_code = [1] * self.out_channels_size # TODO check downsample auto self.downsamples = downsamples self.residauls = residauls self.model = self._prune(self.model, in_channels_code, out_channels_code)
def _get_model(args): """Get model.""" from zeus.model_zoo import ModelZoo model = ModelZoo.get_model(args.model_desc, args.model) if vega.is_torch_backend(): if args.device == "GPU": model = model.cuda() model.eval() return model
def __init__(self, desc, weight_file=None): super(Deformation, self).__init__() self.search_space = OrderedDict() self.conditions = deque() self.desc = deepcopy(desc.get('desc')) or desc self.model = ModelZoo().get_model(desc, weight_file) self.get_search_space(self.model) self.decode() self.deform() self.props.clear()
def load_model(self): """Load model.""" if not self.model_desc and not self.weights_file: saved_folder = self.get_local_worker_path(self.step_name, self.worker_id) self.weights_file = FileOps.join_path( saved_folder, 'model_{}.pth'.format(self.worker_id)) self.model_desc = FileOps.join_path( saved_folder, 'desc_{}.json'.format(self.worker_id)) if 'modules' not in self.model_desc: self.model_desc = ModelConfig.model_desc self.model = ModelZoo.get_model(self.model_desc, self.weights_file)
class Deformation(Module): """Get output layer by layer names and connect into a OrderDict.""" def __init__(self, desc, from_graph=None, weight_file=None): super(Deformation, self).__init__() self._desc = {"props": deepcopy(desc.get('props')) or {}} if from_graph: self.model = GraphGetter(desc, weight_file).model else: self.model = ModelZoo().get_model(desc, weight_file) self._apply_names() self.get_search_space() self.decode() self.deform() self.props.clear() def deform(self): """Deform Modules.""" raise NotImplementedError def decode(self): """Decode Condition and search space.""" for k, v in self.conditions: if self.props.get(v): self.props[k] = self.props.get(v) def to_desc(self, recursion=True): """Convert to model desc.""" return dict(self.model.to_desc(), **self._desc) def state_dict(self): """Get state dict.""" return self.model.state_dict() def get_search_space(self): """Get Search Space from model.""" search_space = DeformationSearchSpace() search_space._traversal(self.model) self.search_space = search_space.search_space self.conditions = search_space.conditions
def load_model(self): """Load model.""" self.saved_folder = self.get_local_worker_path(self.step_name, self.worker_id) if not self.model_desc: self.model_desc = FileOps.join_path(self.saved_folder, 'desc_{}.json'.format(self.worker_id)) if not self.weights_file: if zeus.is_torch_backend(): self.weights_file = FileOps.join_path(self.saved_folder, 'model_{}.pth'.format(self.worker_id)) elif zeus.is_ms_backend(): for file in os.listdir(self.saved_folder): if file.startswith("CKP") and file.endswith(".ckpt"): self.weights_file = FileOps.join_path(self.saved_folder, file) if 'modules' not in self.model_desc: self.model_desc = ModelConfig.model_desc self.model = ModelZoo.get_model(self.model_desc, self.weights_file)
def get_space(self, desc): """Get model and input.""" model_desc = PipeStepConfig.model.model_desc model = ModelZoo().get_model( dict(type='PruneDeformation', desc=model_desc)) search_space = model.search_space params = [] for key, value in search_space.items(): hparam_name = 'network.props.{}'.format(key) params.append( dict(key=hparam_name, type="BINARY_CODE", range=[value])) params.append( dict(key='network.deformation', type="CATEGORY", range=['PruneDeformation'])) logging.info("Prune Search Space: {}".format(params)) return {"hyperparameters": params}
def get_space(self, desc): """Get model and input.""" model_desc = PipeStepConfig.model.model_desc model = ModelZoo().get_model( dict(type='BackboneDeformation', desc=model_desc)) search_space = model.search_space times = random.randint(3, 5) params = [ dict(key="network.props.doublechannel", type="BINARY_CODE", range=[len(search_space), times]), dict(key="network.props.downsample", type="BINARY_CODE", range=[len(search_space), times]) ] params.append( dict(key='network.deformation', type="CATEGORY", range=['BackboneDeformation'])) logging.info("Backbone Search Space: {}".format(params)) return {"hyperparameters": params}
def _init_model(self): """Load model desc from save path and parse to model.""" model = self.trainer.model if self.trainer.config.is_detection_trainer: model_desc = self.trainer.model_desc or self._get_model_desc() else: model_desc = self._get_model_desc() pretrained_model_file = self._get_pretrained_model_file() if not model: if not model_desc: raise Exception( "Failed to Init model, can not get model description.") model = ModelZoo.get_model(model_desc, pretrained_model_file, ModelConfig.head) if model: self.trainer.model_desc = model.desc if zeus.is_torch_backend(): import torch if self.trainer.use_cuda: model = model.cuda() if General._parallel and General.devices_per_trainer > 1: model = torch.nn.DataParallel(model) return model
def load_model(self): """Load model.""" self.saved_folder = self.get_local_worker_path(self.step_name, self.worker_id) if not self.model_desc: model_config = Config( FileOps.join_path(self.saved_folder, 'desc_{}.json'.format(self.worker_id))) if "type" not in model_config and "modules" not in model_config: model_config = ModelConfig.model_desc self.model_desc = model_config if not self.weights_file: if zeus.is_torch_backend(): self.weights_file = FileOps.join_path( self.saved_folder, 'model_{}.pth'.format(self.worker_id)) elif zeus.is_ms_backend(): for file in os.listdir(self.saved_folder): if file.endswith(".ckpt"): self.weights_file = FileOps.join_path( self.saved_folder, file) elif zeus.is_tf_backend(): self.weights_file = FileOps.join_path( self.saved_folder, 'model_{}'.format(self.worker_id)) self.model = ModelZoo.get_model(self.model_desc, self.weights_file)
class PruneGetter(Module): """Get output layer by layer names and connect into a OrderDict.""" def __init__(self, model_name=None, pretrained_model_file=None, downsamples=[4, 14, 27, 46], residauls=[7, 10, 17, 20, 23, 30, 33, 36, 39, 42, 49, 52], model=None, in_channels=None, out_channels=None, num_classes=None, **kwargs): super(PruneGetter, self).__init__() self.model = model if model_name is not None: model_desc = dict(type=model_name) model_desc.update(kwargs) self.model = ModelZoo().get_model(model_desc, pretrained_model_file) self.num_classes = num_classes convs = [ module for name, module in self.model.named_modules() if isinstance(module, torch.nn.Conv2d) ] self.in_channels_size = sum([conv.in_channels for conv in convs]) self.out_channels_size = sum([conv.out_channels for conv in convs]) if in_channels and len(in_channels) < self.in_channels_size: raise ValueError( "in channels mask length should be getter than {}".format( self.in_channels_size)) if out_channels and len(out_channels) < self.out_channels_size: raise ValueError( "out channels mask length should be getter than {}".format( self.out_channels_size)) in_channels_code = self.define_props('in_channels', in_channels) out_channels_code = self.define_props('out_channels', out_channels) if not in_channels_code and not out_channels_code: logging.info( "channels_code is null. use 1 as default, this model will not pruned." ) in_channels_code = [1] * self.in_channels_size out_channels_code = [1] * self.out_channels_size # TODO check downsample auto self.downsamples = downsamples self.residauls = residauls self.model = self._prune(self.model, in_channels_code, out_channels_code) def load_state_dict(self, state_dict, strict=True): """Call subclass load_state_dict function.""" return self.model.load_state_dict(state_dict, strict) def state_dict(self, destination=None, prefix='', keep_vars=False): """Call subclass state_dict function.""" return self.model.state_dict(destination, prefix, keep_vars) def _prune(self, model, in_channels_code, out_channels_code): logging.info("Start to Prune Model.") convs = [ module for name, module in model.named_modules() if isinstance(module, torch.nn.Conv2d) ] org_convs = copy.deepcopy(convs) batch_norms = [ module for name, module in model.named_modules() if isinstance(module, torch.nn.BatchNorm2d) ] pre_end_mask_code = None block_start_mask = None for idx, conv in enumerate(convs): end_mask = None if idx == 0: start_mask = None else: pre_out_channels = org_convs[idx - 1].out_channels if conv.in_channels == pre_out_channels and pre_end_mask_code: start_mask = pre_end_mask_code else: start_mask = in_channels_code[:conv.in_channels] # cache downsaple conv mask code if not block_start_mask: block_start_mask = start_mask if idx in self.downsamples: # downsaple conv start_mask = block_start_mask block_start_mask = None end_mask = pre_end_mask_code elif idx in self.residauls: # Identify jump node end_mask = block_start_mask block_start_mask = None else: in_channels_code = in_channels_code[conv.in_channels:] end_mask = end_mask or out_channels_code[:conv.out_channels] out_channels_code = out_channels_code[conv.out_channels:] # all code is 0, default 1 PruneConv2DFilter(conv).filter(end_mask, start_mask) batch_norm = batch_norms[idx] PruneBatchNormFilter(batch_norm).filter(end_mask) pre_end_mask_code = end_mask linear = [ module for name, module in model.named_modules() if isinstance(module, torch.nn.Linear) ][-1] if self.num_classes: linear.out_features = self.num_classes PruneLinearFilter(linear).filter(pre_end_mask_code) # remove the redundant code and calculate the length if in_channels_code: self.in_channels_size = self.in_channels_size - len( in_channels_code) if out_channels_code: self.out_channels_size = self.out_channels_size - len( out_channels_code) return model
def _get_model(args): from zeus.model_zoo import ModelZoo model = ModelZoo.get_model(args.model_desc_file, args.model_weights_file) model = model.cuda() model.eval() return model