def _traversal_add_connections(self, module): last_convs = [] last_bns = [] for child in module.children(): if isinstance(child, ops.Conv2d): add_convs = [child] elif isinstance(child, ops.Identity): continue else: add_convs = [ conv for name, conv in child.named_modules() if isinstance(conv, ops.Conv2d) ] add_bns = [ bn for name, bn in child.named_modules() if isinstance(bn, ops.BatchNorm2d) ] if add_convs: last_convs.append(add_convs[-1]) if vega.is_ms_backend(): last_bns.append(add_bns[-1]) tmp_pre_conv = self.pre_conv for child in module.children(): self.pre_conv = tmp_pre_conv self._traversal(child) if len(last_convs) > 1: self.pre_conv = last_convs[0] last_convs = last_convs[1:] else: self.pre_conv = tmp_pre_conv for conv in last_convs: self.add_condition(conv.name + '.out_channels', self.pre_conv.name + '.out_channels') # The out_channels value of the jump node is the same as that of the previous nodes # remove from the search space. if len(last_convs) == 1: self.add_forbidden(last_convs[0].name + '.out_channels') self.add_condition(last_convs[0].name + '.out_channels', self.pre_conv.name + '.out_channels') if vega.is_ms_backend(): self.add_condition(last_bns[-1].name + '.num_features', self.pre_conv.name + '.out_channels') else: for last_conv in last_convs: if self.pre_conv == last_conv: continue self.add_forbidden(last_conv.name + '.out_channels') self.add_condition(last_convs[0].name + '.out_channels', self.pre_conv.name + '.out_channels') for k, v in [(k, v) for k, v in self.conditions if v == last_conv.name + '.out_channels']: self.add_condition(k, self.pre_conv.name + '.out_channels') self.pre_conv = last_convs[0]
def __call__(self, model=None, distributed=False, **kwargs): """Call Optimizer class. :param model: model, used in torch case :param distributed: use distributed :return: optimizer """ params = self.map_config.get("params", {}) logging.debug("Call Optimizer. name={}, params={}".format( self.optim_cls.__name__, params)) optimizer = None try: if vega.is_torch_backend(): learnable_params = [ param for param in model.parameters() if param.requires_grad ] optimizer = self.optim_cls(learnable_params, **params) if distributed: optimizer = self.set_distributed(optimizer, model) elif vega.is_tf_backend(): optimizer = dynamic_optimizer(self.optim_cls, **params) elif vega.is_ms_backend(): if "dynamic_lr" in kwargs: params.update({"learning_rate": kwargs["dynamic_lr"]}) learnable_params = [ param for param in model.trainable_params() if param.requires_grad ] optimizer = self.optim_cls(learnable_params, **params) return optimizer except Exception as ex: logging.error("Failed to call Optimizer name={}, params={}".format( self.optim_cls.__name__, params)) raise ex
def _generate_init_model(self): """Generate init model by loading pretrained model. :return: initial model after loading pretrained model :rtype: torch.nn.Module """ model_init = self._new_model_init() chn_mask = self._init_chn_node_mask() if vega.is_torch_backend(): checkpoint = torch.load(self.config.init_model_file + '.pth') model_init.load_state_dict(checkpoint) model = PruneMobileNet(model_init).apply(chn_mask) model.to(self.device) elif vega.is_tf_backend(): model = model_init with tf.compat.v1.Session( config=self.trainer._init_session_config()) as sess: saver = tf.compat.v1.train.import_meta_graph("{}.meta".format( self.config.init_model_file)) saver.restore(sess, self.config.init_model_file) all_weight = tf.compat.v1.get_collection( tf.compat.v1.GraphKeys.VARIABLES) all_weight = [ t for t in all_weight if not t.name.endswith('Momentum:0') ] PruneMobileNet(all_weight).apply(chn_mask) save_file = FileOps.join_path( self.trainer.get_local_worker_path(), 'prune_model') saver.save(sess, save_file) elif vega.is_ms_backend(): parameter_dict = load_checkpoint(self.config.init_model_file) load_param_into_net(model_init, parameter_dict) model = PruneMobileNet(model_init).apply(chn_mask) return model
def eval_model_parameters(model): """Calculate number of parameters in million (M) for a model. :param model: A model :type model: nn.Module :return: The number of parameters :rtype: Float """ if vega.is_torch_backend(): return np.sum(v.numel() for name, v in model.named_parameters() if "auxiliary" not in name) / 1e6 elif vega.is_tf_backend(): import tensorflow as tf tf.compat.v1.reset_default_graph() dummy_input = tf.compat.v1.placeholder( dtype=tf.float32, shape=[1, 32, 32, 3] if model.data_format == 'channels_last' else [1, 3, 32, 32]) model.training = True model(dummy_input) all_weight = tf.compat.v1.get_collection( tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES) weight_op = [t for t in all_weight if "auxiliary" not in t.name] return np.sum([np.prod(t.get_shape().as_list()) for t in weight_op]) * 1e-6 elif vega.is_ms_backend(): return 0
def load_records_from_model_folder(cls, model_folder): """Transfer json_file to records.""" if not model_folder or not os.path.exists(model_folder): logging.error("Failed to load records from model folder, folder={}".format(model_folder)) return [] records = [] pattern = FileOps.join_path(model_folder, "desc_*.json") files = glob.glob(pattern) for _file in files: try: with open(_file) as f: worker_id = _file.split(".")[-2].split("_")[-1] weights_file = os.path.join(os.path.dirname(_file), "model_{}".format(worker_id)) if vega.is_torch_backend(): weights_file = '{}.pth'.format(weights_file) elif vega.is_ms_backend(): weights_file = '{}.ckpt'.format(weights_file) if not os.path.exists(weights_file): weights_file = None sample = dict(worker_id=worker_id, desc=json.load(f), weights_file=weights_file) record = ReportRecord().load_dict(sample) records.append(record) except Exception as ex: logging.info('Can not read records from json because {}'.format(ex)) return records
def before_train(self, logs=None): """Be called before the train process.""" self.config = self.trainer.config self.device = vega.is_gpu_device() if vega.is_gpu_device( ) is not True else 0 self.base_net_desc = self.trainer.model.desc sess_config = None if vega.is_torch_backend(): if vega.is_npu_device(): count_input = torch.FloatTensor(1, 3, 32, 32).npu() elif vega.is_gpu_device(): count_input = torch.FloatTensor(1, 3, 32, 32).to(self.device) elif vega.is_tf_backend(): count_input = tf.random.uniform([1, 3, 32, 32], dtype=tf.float32) sess_config = self.trainer._init_session_config() elif vega.is_ms_backend(): count_input = mindspore.Tensor( np.random.randn(1, 3, 32, 32).astype(np.float32)) self.flops_count, self.params_count = calc_model_flops_params( self.trainer.model, count_input) self.latency_count = calc_forward_latency(self.trainer.model, count_input, sess_config) logging.info("after prune model glops=%sM, params=%sK, latency=%sms", self.flops_count * 1e-6, self.params_count * 1e-3, self.latency_count * 1000) self.trainer.model = self._generate_init_model() if vega.is_torch_backend(): self.trainer.optimizer = Optimizer()( model=self.trainer.model, distributed=self.trainer.distributed) self.trainer.lr_scheduler = LrScheduler()(self.trainer.optimizer)
def _get_data_format(): if vega.is_torch_backend() or vega.is_ms_backend(): return 'channels_first' elif vega.is_tf_backend(): return 'channels_last' else: return None
def get_named_modules(layer): """Get named modules.""" if vega.is_tf_backend(): return [(op.name, op) for op in layer] elif vega.is_torch_backend(): return layer.named_modules() elif vega.is_ms_backend(): return layer._children_scope_recursive()
def _infer(args, loader, model=None): """Choose backend.""" if vega.is_torch_backend(): return _infer_pytorch(model, loader) elif vega.is_tf_backend(): return _infer_tf(args, model, loader) elif vega.is_ms_backend(): return _infer_ms(args, model, loader)
def get_shape(layer): """Get weight shape.""" if vega.is_tf_backend(): return layer.get_shape() elif vega.is_torch_backend(): return layer.weight.data.shape elif vega.is_ms_backend(): para_name = list(layer._params)[0] return getattr(layer, para_name).default_input.shape
def from_dict(cls, data, skip_check=True): """Restore config from a dictionary or a file.""" cls = super(LossConfig, cls).from_dict(data, skip_check) if vega.is_ms_backend(): if "params" not in data: cls.params = {'sparse': True} elif "sparse" not in data.params: cls.params.update({'sparse': True}) return cls
def Adapter(dataset): """Adapter of dataset.""" if vega.is_torch_backend(): from .pytorch.adapter import TorchAdapter as Adapter elif vega.is_tf_backend(): from .tensorflow.adapter import TfAdapter as Adapter elif vega.is_ms_backend(): from .mindspore.adapter import MsAdapter as Adapter else: raise ValueError return Adapter(dataset)
def _init_metrics(self, metrics=None): """Init metrics.""" if metrics is not None: return metrics else: if vega.is_torch_backend(): from vega.metrics.pytorch.metrics import Metrics elif vega.is_tf_backend(): from vega.metrics.tensorflow.metrics import Metrics elif vega.is_ms_backend(): from vega.metrics.mindspore.metrics import Metrics return Metrics()
def to_module(cls, model): """Build model desc before get model.""" if vega.is_ms_backend(): from vega.networks.mindspore.backbones.ms2vega import transform_model return transform_model(model) else: try: model_desc = cls.parse_desc_from_pretrained_model(model) except Exception as ex: logging.warn("Parse model desc failed: {}".format(ex)) return model return ModelZoo.get_model(model_desc)
def __init__(self, type_dict, params_dict): """Init config backend mapping.""" self.type_mapping_dict = copy.deepcopy(type_dict) self.params_mapping_dict = copy.deepcopy(params_dict) self.backend_type = None if vega.is_torch_backend(): self.backend_type = 'torch' elif vega.is_tf_backend(): self.backend_type = 'tf' elif vega.is_ms_backend(): self.backend_type = 'ms' else: raise ValueError('Backend type must be torch, tf or ms.')
def __init__(self, backbone_load_path, backbone_out_sizes, op_names, agg_size, aux_cell, sep_repeats, agg_concat, num_classes, config, method, code): """Construct the AdelaideFastNAS class. :param net_desc: config of the searched structure """ super(AdelaideFastNAS, self).__init__() self.encoder = MobileNetV2Tiny(backbone_load_path) self.decoder = MicroDecoder(backbone_out_sizes, op_names, num_classes, config, agg_size, aux_cell, sep_repeats, agg_concat) self.head = ops.InterpolateScale(mode='bilinear', align_corners=True) if vega.is_ms_backend(): self.permute = ops.Permute((0, 2, 3, 1))
def parse_module_name(name, module): """Parse the module name of mindspore.""" if vega.is_ms_backend(): while (list(module.cells()) != []): module = list(module.cells())[0] name_list = name.split("/")[1:] new_name = "" for name in name_list: name = "." + name.split("-")[0] new_name += name return new_name[1:], module else: return name, module
def calc_model_flops_params(model, input, custom_hooks=None, verbose=False): """Pytorch model flops and parameters calculation. :param model: pytorch model :type model: torch.nn.Module :param input: pytorch input tensor :type input: torch.Tensor :param custom_hooks: hooks defined by outside customer :type custom_hooks: dict or None :param verbose: whether to print op type which not in collection :type verbose: bool, default True :return: flops and params :rtype: float, float """ try: _model = deepcopy(model) except Exception: _model = model if vega.is_torch_backend(): from thop import profile if custom_hooks is None: custom_hooks = {} custom_hooks = add_new_hooks(custom_hooks) inputs = (input, ) flops, params = profile(_model, inputs, custom_hooks, verbose) del _model elif vega.is_tf_backend(): import tensorflow.compat.v1 as tf with tf.Graph().as_default() as graph: dummy_input = tf.placeholder(dtype=tf.float32, shape=input.shape.as_list()) _model.training = False _model(dummy_input) opts = tf.profiler.ProfileOptionBuilder.float_operation() flops = tf.profiler.profile(graph, cmd='op', options=opts).total_float_ops opts = tf.profiler.ProfileOptionBuilder.trainable_variables_parameter( ) params = tf.profiler.profile(graph, cmd='op', options=opts).total_parameters flops *= 0.5 del _model elif vega.is_ms_backend(): total_params = 0 for param in model.trainable_params(): total_params += np.prod(param.shape) params = total_params # TODO flops = 0 return flops, params
def before_train(self, logs=None): """Be called before the training process.""" self.config = self.trainer.config if vega.is_torch_backend(): count_input = torch.FloatTensor(1, 3, 192, 192).cuda() elif vega.is_tf_backend(): tf.compat.v1.reset_default_graph() count_input = tf.random.uniform([1, 192, 192, 3], dtype=tf.float32) elif vega.is_ms_backend(): count_input = mindspore.Tensor( np.random.randn(1, 3, 192, 192).astype(np.float32)) flops_count, params_count = calc_model_flops_params( self.trainer.model, count_input) self.flops_count, self.params_count = flops_count * 1e-9, params_count * 1e-3 logger.info("Flops: {:.2f} G, Params: {:.1f} K".format( self.flops_count, self.params_count))
def search_infer_step(self, alpha): """Infer in search stage. :param valid_queue: valid dataloader :type valid_queue: dataloader :param model: The model to be trained :type model: nn.Module :param alpha: encoding of a model :type alpha: array :return: Average top1 acc and loss :rtype: nn.Tensor """ if vega.is_torch_backend(): metrics = Metrics() alpha_tensor = torch.from_numpy(alpha).cuda() self.trainer.model.eval() with torch.no_grad(): for step, (input, target) in enumerate(self.trainer.valid_loader): input = input.cuda() target = target.cuda(non_blocking=True) logits = self.trainer.model(input, alpha=alpha_tensor) metrics(logits, target) elif vega.is_tf_backend(): # self.trainer.valid_alpha = tf.convert_to_tensor(alpha) metrics = self.trainer.valid_metrics setattr(self.trainer, 'valid_alpha', alpha) eval_results = self.trainer.estimator.evaluate( input_fn=self.trainer.valid_loader.input_fn, steps=len(self.trainer.valid_loader)) metrics.update(eval_results) elif vega.is_ms_backend(): metrics = self.trainer.valid_metrics setattr(self.trainer, 'valid_alpha', alpha) eval_metrics = self.trainer.ms_model.eval( valid_dataset=self.trainer.valid_loader, dataset_sink_mode=self.trainer.dataset_sink_mode) metrics.update(eval_metrics) performance = metrics.results objectives = metrics.objectives # support min for key, mode in objectives.items(): if mode == 'MIN': performance[key] = -1 * performance[key] performance.update({'params': self.eval_model_sizes(alpha)}) return performance
def _calc_forward_latency_gpu(model, input, sess_config=None, num=100): """Model forward latency calculation. :param model: network model :type model: torch or tf module :param input: input tensor :type input: Tensor of torch or tf :param num: forward number :type num: int :return: forward latency :rtype: float """ prepare_num = int(0.05 * num) if vega.is_torch_backend(): pre_mode = model.training model.train(False) for _ in range(prepare_num): model(input) start_time = time.time() for _ in range(num): model(input) latency = (time.time() - start_time) / num model.train(pre_mode) elif vega.is_tf_backend(): import tensorflow.compat.v1 as tf with tf.Graph().as_default(): input_holder = tf.placeholder(dtype=tf.float32, shape=input.shape.as_list()) pre_mode = model.training model.training = False output = model(input_holder) with tf.Session(config=sess_config) as sess: sess.run(tf.global_variables_initializer()) input = tf.random.uniform(input.shape.as_list(), dtype=input.dtype) input_numpy = input.eval(session=sess) for _ in range(prepare_num): sess.run(output, feed_dict={input_holder: input_numpy}) start_time = time.time() for _ in range(num): sess.run(output, feed_dict={input_holder: input_numpy}) latency = (time.time() - start_time) / num model.training = pre_mode elif vega.is_ms_backend(): latency = 0. return latency
def before_train(self, logs=None): """Be called before the train process.""" self.config = self.trainer.config self.device = self.trainer.config.device self.base_net_desc = self.trainer.model.desc if vega.is_torch_backend(): count_input = torch.FloatTensor(1, 3, 32, 32).to(self.device) elif vega.is_tf_backend(): count_input = tf.random.uniform([1, 32, 32, 3], dtype=tf.float32) elif vega.is_ms_backend(): count_input = mindspore.Tensor( np.random.randn(1, 3, 32, 32).astype(np.float32)) self.flops_count, self.params_count = calc_model_flops_params( self.trainer.model, count_input) print( f"flops:{self.flops_count}, model size:{self.params_count*4/1024**2} MB" ) self.trainer.model = self._generate_init_model()
def get_input_data(self): """Get input data.""" count_input = None if vega.is_torch_backend(): data_iter = iter(self.dataloader) input_data, _ = data_iter.next() count_input = input_data[:1] elif vega.is_tf_backend(): import tensorflow as tf datasets = self.dataloader.input_fn() data_iter = tf.compat.v1.data.make_one_shot_iterator(datasets) input_data, _ = data_iter.get_next() count_input = input_data[:1] elif vega.is_ms_backend(): data_iter = self.dataloader.create_dict_iterator() for batch in data_iter: count_input = batch['image'] break return count_input
def apply(self, mask_code): """Apply mask to batchNorm.""" end_mask = np.asarray(mask_code) idx = np.squeeze( np.argwhere(np.asarray(np.ones(end_mask.shape) - end_mask))).tolist() self._make_mask(idx) if vega.is_tf_backend(): import tensorflow as tf return tf.assign( self.layer, self.layer * tf.constant(self.mask, dtype=self.layer.dtype)) elif vega.is_torch_backend(): import torch self.layer.weight.data = self.layer.weight.data * torch.FloatTensor( self.mask) self.layer.bias.data = self.layer.bias.data * torch.FloatTensor( self.mask) self.layer.running_mean = self.layer.running_mean * torch.FloatTensor( self.mask) self.layer.running_var = self.layer.running_var * torch.FloatTensor( self.mask) self.layer.weight.data[idx].requires_grad = False self.layer.bias.data[idx].requires_grad = False self.layer.running_mean[idx].requires_grad = False self.layer.running_var[idx].requires_grad = False elif vega.is_ms_backend(): from mindspore import Tensor self.layer.moving_mean.default_input = self.layer.moving_mean.default_input * \ Tensor(self.mask, self.layer.moving_mean.default_input.dtype) self.layer.moving_variance.default_input = self.layer.moving_variance.default_input * \ Tensor(self.mask, self.layer.moving_variance.default_input.dtype) self.layer.gamma.default_input = self.layer.gamma.default_input * \ Tensor(self.mask, self.layer.gamma.default_input.dtype) self.layer.beta.default_input = self.layer.beta.default_input * \ Tensor(self.mask, self.layer.beta.default_input.dtype) for id in idx: self.layer.moving_mean.default_input[id].requires_grad = False self.layer.moving_variance.default_input[ id].requires_grad = False self.layer.gamma.default_input[id].requires_grad = False self.layer.beta.default_input[id].requires_grad = False
def apply(self, end_mask_code, start_mask_code=None): """Apply mask to weight.""" end_mask_code = np.array(end_mask_code) if start_mask_code is not None: start_mask_code = np.array(start_mask_code) start_channel_idx = None end_channel_idx = np.squeeze( np.argwhere( np.asarray(np.ones(end_mask_code.shape) - end_mask_code))).tolist() if start_mask_code is not None: start_channel_idx = np.squeeze( np.argwhere( np.asarray( np.ones(start_mask_code.shape) - start_mask_code))).tolist() self._make_mask(end_mask_code, start_mask_code) if vega.is_tf_backend(): import tensorflow as tf return tf.assign( self.layer, self.layer * tf.constant(self.mask, dtype=self.layer.dtype)) elif vega.is_torch_backend(): import torch self.layer.weight.data = self.layer.weight.data * torch.FloatTensor( self.mask) self.layer.weight.data[ end_channel_idx, :, :, :].requires_grad = False if start_channel_idx is not None: self.layer.weight.data[:, start_channel_idx, :, :].requires_grad = False elif vega.is_ms_backend(): from mindspore import Tensor self.layer.weight.default_input = self.layer.weight.default_input * \ Tensor(self.mask, self.layer.weight.default_input.dtype) for idx in end_channel_idx: self.layer.weight.default_input[ idx, :, :, :].requires_grad = False if start_channel_idx is not None: for idx in start_channel_idx: self.layer.weight.default_input[:, idx, :, :].requires_grad = False
def before_train(self, logs=None): """Fetch trainer info before train stage.""" self._fix_path = "_".join([self.trainer.step_name, str(self.trainer.worker_id)]) self.summary = SummaryBoard(self._archive_root, self._fix_path) # add graph only once. if vega.is_tf_backend(): import tensorflow as tf datasets = self.trainer.valid_input_fn() data_iter = tf.compat.v1.data.make_one_shot_iterator(datasets) input_data, _ = data_iter.get_next() self.input = input_data[:1] graph = self.trainer.graph _graph_name_list = [n.name for n in graph.as_graph_def().node] if len(_graph_name_list) < 2: graph = _fetch_tf_graph(self.trainer.model, self.input) self.summary.add_graph(graph=graph, backend="tf") elif vega.is_torch_backend(): model = self.trainer.model data_iter = iter(self.trainer.train_loader) input_batch, _ = data_iter.next() input_data = input_batch[:1] if not self.trainer.config.is_detection_trainer: if vega.is_gpu_device(): input_data = input_data.cuda() elif vega.is_npu_device(): input_data = input_data.npu() try: self.summary.add_graph(model=model, feed_data=input_data, backend="torch") except BaseException as err: logging.warning("Dump PyTorch model failed! with: \n{}".format(err)) elif vega.is_ms_backend(): logging.debug("Don't support mindspore model dump yet.") else: logging.warning("non-known backend.")
def _load_pretrained_model(cls, model, pretrained_model_file): pretrained_model_file = cls._get_abs_path(pretrained_model_file) logging.info("load model weights from file, weights file={}".format( pretrained_model_file)) if vega.is_torch_backend(): if not os.path.isfile(pretrained_model_file): raise Exception( f"Pretrained model is not existed, model={pretrained_model_file}" ) import torch checkpoint = torch.load(pretrained_model_file) model.load_state_dict(checkpoint) # del checkpoint if vega.is_tf_backend(): if pretrained_model_file.endswith('.pth'): checkpoint = convert_checkpoint_from_pytorch( pretrained_model_file, model) model.load_checkpoint_from_numpy(checkpoint) else: pretrained_model_file = cls._get_tf_model_file( pretrained_model_file) model.load_checkpoint(pretrained_model_file) elif vega.is_ms_backend(): from mindspore.train.serialization import load_checkpoint if hasattr(model, "pretrained"): pretrained_weight = model.pretrained(pretrained_model_file) else: if os.path.isfile(pretrained_model_file): pretrained_weight = pretrained_model_file else: for file in os.listdir(pretrained_model_file): if file.endswith(".ckpt"): pretrained_weight = os.path.join( pretrained_model_file, file) break load_checkpoint(pretrained_weight, net=model) # os.remove(pretrained_weight) return model
def apply(self, mask_code): """Apply mask to linear.""" mask_code = np.asarray(mask_code) idx = np.squeeze( np.argwhere(np.asarray(np.ones(mask_code.shape) - mask_code))).tolist() self._make_mask(idx) if vega.is_tf_backend(): import tensorflow as tf return tf.assign( self.layer, self.layer * tf.constant(self.mask, dtype=self.layer.dtype)) elif vega.is_torch_backend(): import torch self.layer.weight.data = self.layer.weight.data * torch.FloatTensor( self.mask) self.layer.weight.data[:, idx].requires_grad = False elif vega.is_ms_backend(): self.layer.weight.default_input = self.layer.weight.default_input * \ torch.FloatTensor(self.mask, self.layer.weight.default_input.dtype) for id in idx: self.layer.weight.default_input[:, id].requires_grad = False
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 = self._get_model_desc() if not self.weights_file: if vega.is_torch_backend(): self.weights_file = FileOps.join_path( self.saved_folder, 'model_{}.pth'.format(self.worker_id)) elif vega.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 vega.is_tf_backend(): self.weights_file = FileOps.join_path( self.saved_folder, 'model_{}'.format(self.worker_id)) if self.weights_file is not None and os.path.exists(self.weights_file): self.model = ModelZoo.get_model(self.model_desc, self.weights_file) else: logger.info("evalaute model without loading weights file") self.model = ModelZoo.get_model(self.model_desc)
def before_train(self, logs=None): """Be called before the training process.""" self.config = self.trainer.config input_shape = [ 1, 3, 192, 192 ] if General.data_format == 'channels_first' else [1, 192, 192, 3] if vega.is_torch_backend(): if vega.is_gpu_device(): count_input = torch.FloatTensor(*input_shape).cuda() elif vega.is_npu_device(): input_shape = [1, 3, 192, 192] count_input = torch.FloatTensor(*input_shape).npu() elif vega.is_tf_backend(): tf.compat.v1.reset_default_graph() count_input = tf.random.uniform(input_shape, dtype=tf.float32) elif vega.is_ms_backend(): count_input = mindspore.Tensor( np.random.randn(*input_shape).astype(np.float32)) flops_count, params_count = calc_model_flops_params( self.trainer.model, count_input) self.flops_count, self.params_count = flops_count * 1e-9, params_count * 1e-3 logger.info("Flops: {:.2f} G, Params: {:.1f} K".format( self.flops_count, self.params_count))