def get_space(self, desc):
"""Get model and input."""
model = NetworkDesc(PipeStepConfig.model.model_desc).to_model()
arch_params_key = 'network._arch_params.Prune.{}.out_channels'
search_space = [dict(key=arch_params_key.format(module.name), type="BINARY_CODE", range=[module.out_channels])
for name, module in model.named_modules() if isinstance(module, ops.Conv2d)]
logging.info("Prune Nas Search Space: {}".format(search_space))
return {"hyperparameters": search_space}
def new_model(self):
"""Build new model."""
net_desc = NetworkDesc(self.search_space)
model_new = net_desc.to_model().cuda()
for x, y in zip(model_new.arch_parameters(),
self.model.arch_parameters()):
x.detach().copy_(y.detach())
return model_new
def transform_model(model):
"""Transform the torch model to Vega model."""
new_model_dict = OrderedDict()
for name, module in model.name_cells().items():
if isinstance(module, atom_op):
if isinstance(module, nn.Flatten):
new_model_dict["mean"] = ops.AdaptiveAvgPool2d()
new_model_dict[name] = _transform_op(module)
sub_modules = OrderedDict()
if isinstance(module, nn.SequentialCell):
for sub_name, sub_module in module.name_cells().items():
if isinstance(sub_module, atom_block):
sub_modules[sub_name] = _transform_block(sub_module)
sub_modules[sub_name].update_parameters_name(name + "." +
sub_name +
".")
if isinstance(sub_module, atom_op):
sub_modules[sub_name] = _transform_op(sub_module)
sub_modules[sub_name].update_parameters_name(name + "." +
sub_name +
".")
new_model_dict[name] = Sequential(sub_modules)
model = Sequential(new_model_dict)
desc = model.to_desc()
vega_model = NetworkDesc(desc).to_model()
return vega_model
def _init_model(self):
"""Initialize the model architecture for full train step.
:return: train model
:rtype: class
"""
logging.info('Initializing model')
if self.cfg.model_desc:
logging.debug("model_desc: {}".format(self.cfg.model_desc))
_file = FileOps.join_path(
self.worker_path, "model_desc_{}.json".format(self._worker_id))
with open(_file, "w") as f:
json.dump(self.cfg.model_desc, f)
if self.cfg.distributed:
hvd.join()
model_desc = self.cfg.model_desc
net_desc = NetworkDesc(model_desc)
model = net_desc.to_model()
return model
else:
return None
def get_model(cls,
model_desc=None,
pretrained_model_file=None,
exclude_weight_prefix=None):
"""Get model from model zoo.
:param network_name: the name of network, eg. ResNetVariant.
:type network_name: str or None.
:param network_desc: the description of network.
:type network_desc: str or None.
:param pretrained_model_file: path of model.
:type pretrained_model_file: str.
:return: model.
:rtype: model.
"""
model = None
if model_desc is not None:
try:
network = NetworkDesc(model_desc)
model = network.to_model()
except Exception as e:
logging.error(
"Failed to get model, model_desc={}, msg={}".format(
model_desc, str(e)))
raise e
logging.info("Model was created.")
if not isinstance(model, Module):
model = cls.to_module(model)
if pretrained_model_file is not None:
if exclude_weight_prefix:
model.exclude_weight_prefix = exclude_weight_prefix
model = cls._load_pretrained_model(model, pretrained_model_file)
model = transform_architecture(model, pretrained_model_file)
if model is None:
raise ValueError("Failed to get mode, model is None.")
return model
def _new_model_init(self):
"""Init new model.
:return: initial model after loading pretrained model
:rtype: torch.nn.Module
"""
init_model_file = self.config.init_model_file
if ":" in init_model_file:
local_path = FileOps.join_path(
self.trainer.get_local_worker_path(),
os.path.basename(init_model_file))
FileOps.copy_file(init_model_file, local_path)
self.config.init_model_file = local_path
network_desc = copy.deepcopy(self.base_net_desc)
network_desc.backbone.cfgs = network_desc.backbone.base_cfgs
model_init = NetworkDesc(network_desc).to_model()
return model_init
def transform_architecture(model, pretrained_model_file=None):
"""Transform architecture."""
if not hasattr(model, "_arch_params") or not model._arch_params or \
PipeStepConfig.pipe_step.get("type") == "TrainPipeStep":
return model
model._apply_names()
logging.info(
"Start to transform architecture, model arch params type: {}".format(
model._arch_params_type))
ConnectionsArchParamsCombiner().combine(model)
if vega.is_ms_backend():
from mindspore.train.serialization import load_checkpoint
changed_name_list = []
mask_weight_list = []
for name, module in model.named_modules():
if not ClassFactory.is_exists(model._arch_params_type,
module.model_name):
continue
changed_name_list, mask_weight_list = decode_fn_ms(
module, changed_name_list, mask_weight_list)
assert len(changed_name_list) == len(mask_weight_list)
# change model and rebuild
model_desc = model.desc
root_name = [
name for name in list(model_desc.keys())
if name not in ('type', '_arch_params')
]
for changed_name, mask in zip(changed_name_list, mask_weight_list):
name = changed_name.split('.')
name[0] = root_name[int(name[0])]
assert len(name) <= 6
if len(name) == 6:
model_desc[name[0]][name[1]][name[2]][name[3]][name[4]][
name[5]] = sum(mask)
if len(name) == 5:
model_desc[name[0]][name[1]][name[2]][name[3]][name[4]] = sum(
mask)
if len(name) == 4:
model_desc[name[0]][name[1]][name[2]][name[3]] = sum(mask)
if len(name) == 3:
model_desc[name[0]][name[1]][name[2]] = sum(mask)
if len(name) == 2:
model_desc[name[0]][name[1]] = sum(mask)
network = NetworkDesc(model_desc)
model = network.to_model()
model_desc.pop(
'_arch_params') if '_arch_params' in model_desc else model_desc
model.desc = model_desc
# change weight
if hasattr(model, "pretrained"):
pretrained_weight = model.pretrained(pretrained_model_file)
load_checkpoint(pretrained_weight, net=model)
os.remove(pretrained_weight)
else:
for name, module in model.named_modules():
if not ClassFactory.is_exists(model._arch_params_type,
module.model_name):
continue
arch_cls = ClassFactory.get_cls(model._arch_params_type,
module.model_name)
decode_fn(module, arch_cls)
module.register_forward_pre_hook(arch_cls.fit_weights)
module.register_forward_hook(module.clear_module_arch_params)
return model
def __init__(self, desc):
"""Init faster rcnn.
:param desc: config dict
"""
super(FasterRCNN, self).__init__()
self.num_classes = int(desc.num_classes)
self.number_of_stages = int(desc.number_of_stages)
# Backbone for feature extractor
self.feature_extractor = NetworkDesc(desc.backbone).to_model()
# First stage anchor generator
self.first_stage_anchor_generator = NetworkDesc(
desc["first_stage_anchor_generator"]).to_model()
# First stage target assigner
self.use_matmul_gather_in_matcher = False # Default
self.first_stage_target_assigner = target_assigner.create_target_assigner(
'FasterRCNN',
'proposal',
use_matmul_gather=self.use_matmul_gather_in_matcher)
# First stage box predictor
self.first_stage_box_predictor_arg_scope_fn = scope_generator.get_hyper_params_scope(
desc.first_stage_box_predictor_conv_hyperparams)
self.first_stage_atrous_rate = 1 # Default: 1
self.first_stage_box_predictor_kernel_size = 3 # Default
self.first_stage_box_predictor_depth = 512 # Default
self.first_stage_minibatch_size = 256 # Default
# First stage sampler
self.first_stage_positive_balance_fraction = 0.5 # Default
self.use_static_balanced_label_sampler = False # Default
self.use_static_shapes = False # Default
self.first_stage_sampler = sampler.BalancedPositiveNegativeSampler(
positive_fraction=self.first_stage_positive_balance_fraction,
is_static=(self.use_static_balanced_label_sampler
and self.use_static_shapes))
# First stage NMS
self.first_stage_nms_score_threshold = 0.0
self.first_stage_nms_iou_threshold = 0.7
self.first_stage_max_proposals = 300
self.use_partitioned_nms_in_first_stage = True # Default
self.use_combined_nms_in_first_stage = False # Default
self.first_stage_non_max_suppression_fn = functools.partial(
post_processing.batch_multiclass_non_max_suppression,
score_thresh=self.first_stage_nms_score_threshold,
iou_thresh=self.first_stage_nms_iou_threshold,
max_size_per_class=self.first_stage_max_proposals,
max_total_size=self.first_stage_max_proposals,
use_static_shapes=self.use_static_shapes,
use_partitioned_nms=self.use_partitioned_nms_in_first_stage,
use_combined_nms=self.use_combined_nms_in_first_stage)
# First stage localization loss weight
self.first_stage_localization_loss_weight = 2.0
# First stage objectness loss weight
self.first_stage_objectness_loss_weight = 1.0
# Second stage target assigner
self.second_stage_target_assigner = target_assigner.create_target_assigner(
'FasterRCNN',
'detection',
use_matmul_gather=self.use_matmul_gather_in_matcher)
# Second stage sampler
self.second_stage_batch_size = 64 # Default
self.second_stage_balance_fraction = 0.25 # Default
self.second_stage_sampler = sampler.BalancedPositiveNegativeSampler(
positive_fraction=self.second_stage_balance_fraction,
is_static=(self.use_static_balanced_label_sampler
and self.use_static_shapes))
# Second stage box predictor
self.second_stage_box_predictor = NetworkDesc(
desc.mask_rcnn_box).to_model()
# Second stage NMS function
self.second_stage_non_max_suppression_fn, self.second_stage_score_conversion_fn = \
post_processing_util.get_post_processing_fn(desc.second_stage_post_processing)
# Second stage mask prediction loss weight
self.second_stage_mask_prediction_loss_weight = 1.0 # default
# Second stage localization loss weight
self.second_stage_localization_loss_weight = 2.0
# Second stage classification loss weight
self.second_stage_classification_loss_weight = 1.0
# Second stage classification loss
self.logit_scale = 1.0 # Default
self.second_stage_classification_loss = losses.WeightedSoftmaxClassificationLoss(
logit_scale=self.logit_scale)
self.hard_example_miner = None
self.add_summaries = True
# Crop and resize function
self.use_matmul_crop_and_resize = False # Default
self.crop_and_resize_fn = (
spatial_ops.multilevel_matmul_crop_and_resize
if self.use_matmul_crop_and_resize else
spatial_ops.native_crop_and_resize)
self.clip_anchors_to_image = False # Default
self.resize_masks = True # Default
self.return_raw_detections_during_predict = False # Default
self.output_final_box_features = False # Default
# Image resizer function
self.image_resizer_fn = image_resizer_util.get_image_resizer(
desc.image_resizer)
self.initial_crop_size = 14
self.maxpool_kernel_size = 2
self.maxpool_stride = 2
# Real model to be called
self.model = None
class FasterRCNN(object):
"""Faster RCNN."""
def __init__(self, desc):
"""Init faster rcnn.
:param desc: config dict
"""
super(FasterRCNN, self).__init__()
self.num_classes = int(desc.num_classes)
self.number_of_stages = int(desc.number_of_stages)
# Backbone for feature extractor
self.feature_extractor = NetworkDesc(desc.backbone).to_model()
# First stage anchor generator
self.first_stage_anchor_generator = NetworkDesc(
desc["first_stage_anchor_generator"]).to_model()
# First stage target assigner
self.use_matmul_gather_in_matcher = False # Default
self.first_stage_target_assigner = target_assigner.create_target_assigner(
'FasterRCNN',
'proposal',
use_matmul_gather=self.use_matmul_gather_in_matcher)
# First stage box predictor
self.first_stage_box_predictor_arg_scope_fn = scope_generator.get_hyper_params_scope(
desc.first_stage_box_predictor_conv_hyperparams)
self.first_stage_atrous_rate = 1 # Default: 1
self.first_stage_box_predictor_kernel_size = 3 # Default
self.first_stage_box_predictor_depth = 512 # Default
self.first_stage_minibatch_size = 256 # Default
# First stage sampler
self.first_stage_positive_balance_fraction = 0.5 # Default
self.use_static_balanced_label_sampler = False # Default
self.use_static_shapes = False # Default
self.first_stage_sampler = sampler.BalancedPositiveNegativeSampler(
positive_fraction=self.first_stage_positive_balance_fraction,
is_static=(self.use_static_balanced_label_sampler
and self.use_static_shapes))
# First stage NMS
self.first_stage_nms_score_threshold = 0.0
self.first_stage_nms_iou_threshold = 0.7
self.first_stage_max_proposals = 300
self.use_partitioned_nms_in_first_stage = True # Default
self.use_combined_nms_in_first_stage = False # Default
self.first_stage_non_max_suppression_fn = functools.partial(
post_processing.batch_multiclass_non_max_suppression,
score_thresh=self.first_stage_nms_score_threshold,
iou_thresh=self.first_stage_nms_iou_threshold,
max_size_per_class=self.first_stage_max_proposals,
max_total_size=self.first_stage_max_proposals,
use_static_shapes=self.use_static_shapes,
use_partitioned_nms=self.use_partitioned_nms_in_first_stage,
use_combined_nms=self.use_combined_nms_in_first_stage)
# First stage localization loss weight
self.first_stage_localization_loss_weight = 2.0
# First stage objectness loss weight
self.first_stage_objectness_loss_weight = 1.0
# Second stage target assigner
self.second_stage_target_assigner = target_assigner.create_target_assigner(
'FasterRCNN',
'detection',
use_matmul_gather=self.use_matmul_gather_in_matcher)
# Second stage sampler
self.second_stage_batch_size = 64 # Default
self.second_stage_balance_fraction = 0.25 # Default
self.second_stage_sampler = sampler.BalancedPositiveNegativeSampler(
positive_fraction=self.second_stage_balance_fraction,
is_static=(self.use_static_balanced_label_sampler
and self.use_static_shapes))
# Second stage box predictor
self.second_stage_box_predictor = NetworkDesc(
desc.mask_rcnn_box).to_model()
# Second stage NMS function
self.second_stage_non_max_suppression_fn, self.second_stage_score_conversion_fn = \
post_processing_util.get_post_processing_fn(desc.second_stage_post_processing)
# Second stage mask prediction loss weight
self.second_stage_mask_prediction_loss_weight = 1.0 # default
# Second stage localization loss weight
self.second_stage_localization_loss_weight = 2.0
# Second stage classification loss weight
self.second_stage_classification_loss_weight = 1.0
# Second stage classification loss
self.logit_scale = 1.0 # Default
self.second_stage_classification_loss = losses.WeightedSoftmaxClassificationLoss(
logit_scale=self.logit_scale)
self.hard_example_miner = None
self.add_summaries = True
# Crop and resize function
self.use_matmul_crop_and_resize = False # Default
self.crop_and_resize_fn = (
spatial_ops.multilevel_matmul_crop_and_resize
if self.use_matmul_crop_and_resize else
spatial_ops.native_crop_and_resize)
self.clip_anchors_to_image = False # Default
self.resize_masks = True # Default
self.return_raw_detections_during_predict = False # Default
self.output_final_box_features = False # Default
# Image resizer function
self.image_resizer_fn = image_resizer_util.get_image_resizer(
desc.image_resizer)
self.initial_crop_size = 14
self.maxpool_kernel_size = 2
self.maxpool_stride = 2
# Real model to be called
self.model = None
def _init_model(self, training):
# Init FasterRCNNMetaArch
common_kwargs = {
'is_training':
training,
'num_classes':
self.num_classes,
'image_resizer_fn':
self.image_resizer_fn,
'feature_extractor':
self.feature_extractor.get_real_model(training),
'number_of_stages':
self.number_of_stages,
'first_stage_anchor_generator':
self.first_stage_anchor_generator.get_real_model(training),
'first_stage_target_assigner':
self.first_stage_target_assigner,
'first_stage_atrous_rate':
self.first_stage_atrous_rate,
'first_stage_box_predictor_arg_scope_fn':
self.first_stage_box_predictor_arg_scope_fn,
'first_stage_box_predictor_kernel_size':
self.first_stage_box_predictor_kernel_size,
'first_stage_box_predictor_depth':
self.first_stage_box_predictor_depth,
'first_stage_minibatch_size':
self.first_stage_minibatch_size,
'first_stage_sampler':
self.first_stage_sampler,
'first_stage_non_max_suppression_fn':
self.first_stage_non_max_suppression_fn,
'first_stage_max_proposals':
self.first_stage_max_proposals,
'first_stage_localization_loss_weight':
self.first_stage_localization_loss_weight,
'first_stage_objectness_loss_weight':
self.first_stage_objectness_loss_weight,
'second_stage_target_assigner':
self.second_stage_target_assigner,
'second_stage_batch_size':
self.second_stage_batch_size,
'second_stage_sampler':
self.second_stage_sampler,
'second_stage_non_max_suppression_fn':
self.second_stage_non_max_suppression_fn,
'second_stage_score_conversion_fn':
self.second_stage_score_conversion_fn,
'second_stage_localization_loss_weight':
self.second_stage_localization_loss_weight,
'second_stage_classification_loss':
self.second_stage_classification_loss,
'second_stage_classification_loss_weight':
self.second_stage_classification_loss_weight,
'hard_example_miner':
self.hard_example_miner,
'add_summaries':
self.add_summaries,
'crop_and_resize_fn':
self.crop_and_resize_fn,
'clip_anchors_to_image':
self.clip_anchors_to_image,
'use_static_shapes':
self.use_static_shapes,
'resize_masks':
self.resize_masks,
'return_raw_detections_during_predict':
self.return_raw_detections_during_predict,
'output_final_box_features':
self.output_final_box_features
}
self.model = faster_rcnn_meta_arch.FasterRCNNMetaArch(
initial_crop_size=self.initial_crop_size,
maxpool_kernel_size=self.maxpool_kernel_size,
maxpool_stride=self.maxpool_stride,
second_stage_mask_rcnn_box_predictor=self.
second_stage_box_predictor.get_real_model(training),
second_stage_mask_prediction_loss_weight=(
self.second_stage_mask_prediction_loss_weight),
**common_kwargs)
def get_real_model(self, training):
"""Get or init real model."""
if self.model:
return self.model
else:
self._init_model(training)
return self.model
def __call__(self, features, labels, training):
"""Forward function of faster-rcnn."""
if training:
self.get_real_model(training).provide_groundtruth(
groundtruth_boxes_list=tf.unstack(
labels[fields.InputDataFields.groundtruth_boxes]),
groundtruth_classes_list=tf.unstack(
labels[fields.InputDataFields.groundtruth_classes]),
groundtruth_weights_list=tf.unstack(
labels[fields.InputDataFields.groundtruth_weights]))
predict_results = self.get_real_model(training).predict(
features[fields.InputDataFields.image],
features[fields.InputDataFields.true_image_shape])
return predict_results
def loss(self, predict_results, true_image_shapes):
"""Get loss function of faster-rcnn."""
return self.get_real_model(True).loss(predict_results,
true_image_shapes)
def updates(self):
"""Update faster-rcnn model."""
return self.get_real_model(True).updates()
def regularization_losses(self):
"""Get regularization loss of faster-rcnn."""
return self.get_real_model(True).regularization_losses()
def restore_map(self, fine_tune_checkpoint_type,
load_all_detection_checkpoint_vars):
"""Restore map of faster-rcnn."""
return self.get_real_model(True).restore_map(
fine_tune_checkpoint_type=fine_tune_checkpoint_type,
load_all_detection_checkpoint_vars=(
load_all_detection_checkpoint_vars))
def get_model_input(self, desc):
"""Get model and input."""
from vega.networks.network_desc import NetworkDesc
model = NetworkDesc(desc).to_model()
count_input = self.get_input_data()
return model, count_input