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))
