def accept_pruned_input(self, model: NNCFNetwork, graph: NNCFGraph,
node_module):
accept_pruned_input = True
if node_module.groups != 1:
if not is_depthwise_conv(node_module):
accept_pruned_input = False
return accept_pruned_input
def node_propagate_can_prune_attr(self, nncf_node: NNCFNode) -> bool:
"""
Whether node propagates can_prune attr through. That means a node can propagate pruning mask
(for example, activations propagate mask, but convolutions stop mask propagation)
:param nncf_node: node to work with
:return: bool: propagates this node can_prune throw or not
"""
node_module = self.model.get_module_by_scope(nncf_node.op_exec_context.scope_in_model)
node_type = nncf_node.op_exec_context.operator_name
is_conv = node_type in Convolution().get_all_op_aliases()
return not is_conv or (is_conv and is_depthwise_conv(node_module))
def _is_module_prunable(self, target_model: NNCFNetwork, module, module_scope: Scope):
"""
Check whether we should prune module according to algorithm parameters.
:param target_model: model to work with
:param module: module to check
:param module_scope: scope for module
:return: (prune: bool, msg: str)
prune: Whether we should prune module
msg: additional information why we should/shouldn't prune
"""
prune = True
msg = None
pruned_types = self.get_op_types_of_pruned_modules()
input_non_pruned_modules = get_first_pruned_modules(target_model, pruned_types + ['linear'])
output_non_pruned_modules = get_last_pruned_modules(target_model, pruned_types + ['linear'])
module_scope_str = str(module_scope)
if self.ignore_frozen_layers and not module.weight.requires_grad:
msg = "Ignored adding Weight Pruner in scope: {} because"\
" the layer appears to be frozen (requires_grad=False)".format(module_scope_str)
prune = False
elif not self._should_consider_scope(module_scope_str):
msg = "Ignored adding Weight Pruner in scope: {}".format(module_scope_str)
prune = False
elif not self.prune_first and module in input_non_pruned_modules:
msg = "Ignored adding Weight Pruner in scope: {} because"\
" this scope is one of the first convolutions".format(module_scope_str)
prune = False
elif not self.prune_last and module in output_non_pruned_modules:
msg = "Ignored adding Weight Pruner in scope: {} because"\
" this scope is one of the last convolutions".format(module_scope_str)
prune = False
elif is_grouped_conv(module):
if not is_depthwise_conv(module):
msg = "Ignored adding Weight Pruner in scope: {} because" \
" this scope is grouped convolution".format(module_scope_str)
prune = False
elif not self.prune_downsample_convs and is_conv_with_downsampling(module):
msg = "Ignored adding Weight Pruner in scope: {} because"\
" this scope is convolution with downsample".format(module_scope_str)
prune = False
return prune, msg
def _prune_weights(self, target_model: NNCFNetwork):
device = next(target_model.parameters()).device
modules_to_prune = target_model.get_nncf_modules()
insertion_commands = []
bn_for_depthwise = {}
input_non_pruned_modules = get_first_pruned_modules(
target_model,
self.get_types_of_pruned_modules() + ['linear'])
output_non_pruned_modules = get_last_pruned_modules(
target_model,
self.get_types_of_pruned_modules() + ['linear'])
for module_scope, module in modules_to_prune.items():
# Check that we need to prune weights in this op
if not self._is_pruned_module(module):
continue
module_scope_str = str(module_scope)
if self.ignore_frozen_layers and not module.weight.requires_grad:
nncf_logger.info(
"Ignored adding Weight Pruner in scope: {} because"
" the layer appears to be frozen (requires_grad=False)".
format(module_scope_str))
continue
if not self._should_consider_scope(module_scope_str):
nncf_logger.info(
"Ignored adding Weight Pruner in scope: {}".format(
module_scope_str))
continue
if not self.prune_first and module in input_non_pruned_modules:
nncf_logger.info(
"Ignored adding Weight Pruner in scope: {} because"
" this scope is one of the first convolutions".format(
module_scope_str))
continue
if not self.prune_last and module in output_non_pruned_modules:
nncf_logger.info(
"Ignored adding Weight Pruner in scope: {} because"
" this scope is one of the last convolutions".format(
module_scope_str))
continue
if is_grouped_conv(module):
if is_depthwise_conv(module):
previous_conv = get_previous_conv(target_model, module,
module_scope)
if previous_conv:
depthwise_bn = get_bn_for_module_scope(
target_model, module_scope)
bn_for_depthwise[str(previous_conv.op_exec_context.
scope_in_model)] = depthwise_bn
nncf_logger.info(
"Ignored adding Weight Pruner in scope: {} because"
" this scope is grouped convolution".format(
module_scope_str))
continue
if not self.prune_downsample_convs and is_conv_with_downsampling(
module):
nncf_logger.info(
"Ignored adding Weight Pruner in scope: {} because"
" this scope is convolution with downsample".format(
module_scope_str))
continue
nncf_logger.info(
"Adding Weight Pruner in scope: {}".format(module_scope_str))
operation = self.create_weight_pruning_operation(module)
hook = UpdateWeight(operation).to(device)
insertion_commands.append(
InsertionCommand(
InsertionPoint(
InputAgnosticOperationExecutionContext(
"", module_scope, 0),
InsertionType.NNCF_MODULE_PRE_OP), hook,
OperationPriority.PRUNING_PRIORITY))
related_modules = {}
if self.prune_batch_norms:
related_modules[
PrunedModuleInfo.BN_MODULE_NAME] = get_bn_for_module_scope(
target_model, module_scope)
self._pruned_module_info.append(
PrunedModuleInfo(module_scope_str, module, hook.operand,
related_modules))
if self.prune_batch_norms:
self.update_minfo_with_depthwise_bn(bn_for_depthwise)
return insertion_commands
def _create_pruning_groups(self, target_model: NNCFNetwork):
"""
This function groups ALL modules with pruning types to groups that should be pruned together.
1. Create clusters for special ops (eltwises) that should be pruned together
2. Create groups of nodes that should be pruned together (taking into account clusters of special ops)
3. Add remaining single nodes
4. Unite clusters for Conv + Depthwise conv (should be pruned together too)
5. Checks for groups (all nodes in group can prune or all group can't be pruned)
Return groups of modules that should be pruned together.
:param target_model: model to work with
:return: clusterisation of pruned nodes
"""
graph = target_model.get_original_graph()
pruned_types = self.get_op_types_of_pruned_modules()
all_modules_to_prune = target_model.get_nncf_modules_by_module_names(self.compressed_nncf_module_names)
all_nodes_to_prune = graph.get_nodes_by_types(pruned_types) # NNCFNodes here
assert len(all_nodes_to_prune) <= len(all_modules_to_prune)
# 1. Clusters for special ops
special_ops_types = self.get_types_of_grouping_ops()
identity_like_types = IdentityMaskForwardOps.get_all_op_aliases()
special_ops_clusterization = cluster_special_ops(target_model, special_ops_types,
identity_like_types)
pruned_nodes_clusterization = Clusterization("id")
# 2. Clusters for nodes that should be pruned together (taking into account clusters for special ops)
for i, cluster in enumerate(special_ops_clusterization.get_all_clusters()):
all_pruned_inputs = []
pruned_inputs_idxs = set()
for node in cluster.nodes:
sources = get_sources_of_node(node, graph, pruned_types)
for source_node in sources:
source_scope = source_node.op_exec_context.scope_in_model
source_module = target_model.get_module_by_scope(source_scope)
source_node_info = NodeInfo(source_node, source_module, source_scope)
if source_node.node_id not in pruned_inputs_idxs:
all_pruned_inputs.append(source_node_info)
pruned_inputs_idxs.add(source_node.node_id)
if all_pruned_inputs:
cluster = NodesCluster(i, list(all_pruned_inputs), [n.id for n in all_pruned_inputs])
pruned_nodes_clusterization.add_cluster(cluster)
last_cluster_idx = len(special_ops_clusterization.get_all_clusters())
# 3. Add remaining single nodes as separate clusters
for node in all_nodes_to_prune:
if not pruned_nodes_clusterization.is_node_in_clusterization(node.node_id):
scope = node.op_exec_context.scope_in_model
module = target_model.get_module_by_scope(scope)
node_info = NodeInfo(node, module, scope)
cluster = NodesCluster(last_cluster_idx, [node_info], [node.node_id])
pruned_nodes_clusterization.add_cluster(cluster)
last_cluster_idx += 1
# 4. Merge clusters for Conv + Depthwise conv (should be pruned together too)
for node in all_nodes_to_prune:
scope = node.op_exec_context.scope_in_model
module = target_model.get_module_by_scope(scope)
cluster_id = pruned_nodes_clusterization.get_cluster_by_node_id(node.node_id).id
if is_depthwise_conv(module):
previous_conv = get_previous_conv(target_model, module, scope)
if previous_conv:
previous_conv_cluster_id = pruned_nodes_clusterization.get_cluster_by_node_id(
previous_conv.node_id).id
pruned_nodes_clusterization.merge_clusters(cluster_id, previous_conv_cluster_id)
# 5. Checks for groups (all nodes in group can be pruned or all group can't be pruned).
model_analyser = ModelAnalyzer(target_model)
can_prune_analysis = model_analyser.analyse_model_before_pruning()
self._check_pruning_groups(target_model, pruned_nodes_clusterization, can_prune_analysis)
return pruned_nodes_clusterization