def get_all_node_op_addresses_in_block(
compressed_model: NNCFNetwork,
block: BuildingBlock) -> Set[OperationAddress]:
"""
Returns set of operation addresses of all layers included in the block.
:param compressed_model: Target model.
:param block: Building blocks.
:return: Set of operation addresses for building block.
"""
graph = compressed_model.get_original_graph()
nx_graph = graph.get_nx_graph_copy()
start_node, end_node = block
start_node_key, end_node_key = None, None
#pylint: disable=protected-access
for node in nx_graph._node.values():
if start_node == str(node['node_name']):
start_node_key = node['key']
if end_node == str(node['node_name']):
end_node_key = node['key']
simple_paths = nx.all_simple_paths(nx_graph, start_node_key, end_node_key)
op_adresses = set()
for node_keys_in_path in simple_paths:
for node_key in node_keys_in_path:
op_adresses.add(
OperationAddress.from_str(
nx_graph._node[node_key]['node_name']))
start_op_address = OperationAddress.from_str(
nx_graph._node[start_node_key]['node_name'])
op_adresses.remove(start_op_address)
return op_adresses
def test_operator_metatype_marking(self):
from nncf.torch.graph.operator_metatypes import PTConv2dMetatype, PTBatchNormMetatype, PTRELUMetatype, \
PTMaxPool2dMetatype, PTTransposeMetatype, \
PTConvTranspose2dMetatype, PTDepthwiseConv2dSubtype, PTAddMetatype, PTAvgPool2dMetatype, PTLinearMetatype
ref_scope_vs_metatype_dict = {
"/" + MODEL_INPUT_OP_NAME + "_0": PTInputNoopMetatype,
"ModelForMetatypeTesting/NNCFConv2d[conv_regular]/conv2d_0": PTConv2dMetatype,
"ModelForMetatypeTesting/NNCFBatchNorm[bn]/batch_norm_0": PTBatchNormMetatype,
"ModelForMetatypeTesting/relu_0": PTRELUMetatype,
"ModelForMetatypeTesting/transpose__0": PTTransposeMetatype,
"ModelForMetatypeTesting/MaxPool2d[max_pool2d]/max_pool2d_0": PTMaxPool2dMetatype,
"ModelForMetatypeTesting/NNCFConvTranspose2d[conv_transpose]/conv_transpose2d_0": PTConvTranspose2dMetatype,
"ModelForMetatypeTesting/NNCFConv2d[conv_depthwise]/conv2d_0": PTDepthwiseConv2dSubtype,
"ModelForMetatypeTesting/__iadd___0": PTAddMetatype,
"ModelForMetatypeTesting/AdaptiveAvgPool2d[adaptive_avg_pool]/adaptive_avg_pool2d_0": PTAvgPool2dMetatype,
"ModelForMetatypeTesting/NNCFLinear[linear]/linear_0": PTLinearMetatype,
'ModelForMetatypeTesting/flatten_0': PTReshapeMetatype,
"/" + MODEL_OUTPUT_OP_NAME + "_0": PTOutputNoopMetatype,
}
class ModelForMetatypeTesting(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv_regular = torch.nn.Conv2d(in_channels=3,
out_channels=16,
kernel_size=3)
self.bn = torch.nn.BatchNorm2d(num_features=16)
self.max_pool2d = torch.nn.MaxPool2d(kernel_size=2)
self.conv_transpose = torch.nn.ConvTranspose2d(in_channels=16,
out_channels=8,
kernel_size=3)
self.conv_depthwise = torch.nn.Conv2d(in_channels=8, out_channels=8,
kernel_size=5, groups=8)
self.adaptive_avg_pool = torch.nn.AdaptiveAvgPool2d(output_size=1)
self.linear = torch.nn.Linear(in_features=8, out_features=1)
def forward(self, input_):
x = self.conv_regular(input_)
x = self.bn(x)
x = torch.nn.functional.relu(x)
x.transpose_(2, 3)
x = self.max_pool2d(x)
x = self.conv_transpose(x)
x = self.conv_depthwise(x)
x += torch.ones_like(x)
x = self.adaptive_avg_pool(x)
x = self.linear(x.flatten())
return x
model = ModelForMetatypeTesting()
nncf_network = NNCFNetwork(model, [ModelInputInfo([1, 3, 300, 300])])
nncf_graph = nncf_network.get_original_graph()
for nncf_node in nncf_graph.get_all_nodes(): # type: NNCFNode
assert nncf_node.node_name in ref_scope_vs_metatype_dict
ref_metatype = ref_scope_vs_metatype_dict[nncf_node.node_name]
assert nncf_node.metatype == ref_metatype
def test_compressed_graph_models_hw(desc, hw_config_type):
model = desc.model_builder()
config = get_basic_quantization_config_with_hw_config_type(hw_config_type.value,
input_sample_size=desc.input_sample_sizes)
input_info_list = create_input_infos(config)
compressed_model = NNCFNetwork(model, input_infos=input_info_list)
# pylint:disable=protected-access
quantization_builder = QuantizationBuilder(config, should_init=False)
single_config_quantizer_setup = quantization_builder._get_quantizer_setup(compressed_model)
sketch_graph = compressed_model.get_original_graph()
potential_quantizer_graph = prepare_potential_quantizer_graph(sketch_graph, single_config_quantizer_setup)
check_nx_graph(potential_quantizer_graph, desc.dot_filename, _case_dir(hw_config_type.value), sort_dot_graph=False)
def get_bn_for_conv_node_by_name(
target_model: NNCFNetwork,
conv_node_name: NNCFNodeName) -> Optional[torch.nn.Module]:
"""
Returns a batch norm module in target_model that corresponds immediately following a given
convolution node in the model's NNCFGraph representation.
:param target_model: NNCFNetwork to work with
:param module_scope:
:return: batch norm module
"""
graph = target_model.get_original_graph()
conv_node = graph.get_node_by_name(conv_node_name)
bn_node = get_bn_node_for_conv(graph, conv_node)
if bn_node is None:
return None
bn_module = target_model.get_containing_module(bn_node.node_name)
return bn_module
def __init__(self, model: NNCFNetwork,
weight_quantizers: Dict[WeightQuantizerId,
WeightQuantizerInfo],
constraints: HardwareQuantizationConstraints):
self._wq_affected_module_node_name_vs_qid_dict = {
k.target_node_name: k
for k in weight_quantizers.keys()
}
self._quantizer_module_scope_vs_qid_dict = {
} # type: Dict[Scope, WeightQuantizerId]
self._skipped_quantized_weight_node_names = []
self._skipped_weight_quantizers = {
} # type: Dict[WeightQuantizerId, BaseQuantizer]
self._weight_quantizers_in_execution_order_per_scope = OrderedDict(
) # type: Dict[Scope, BaseQuantizer]
self._weight_quantizers_in_execution_order = OrderedDict(
) # type: Dict[WeightQuantizerId, BaseQuantizer]
quantization_types = [
class_type.__name__
for class_type in QUANTIZATION_MODULES.registry_dict.values()
]
weight_module_dict = model.get_nncf_wrapped_model()
quantizers_in_execution_order_per_scope = get_all_modules_by_type(
weight_module_dict, quantization_types)
for scope, quantizer in quantizers_in_execution_order_per_scope.items(
):
if self.is_wq_scope(scope):
affected_module_scope = self.get_owning_module_scope_from_wq_scope(
scope)
affected_module_node = model.get_original_graph(
).get_op_nodes_in_scope(affected_module_scope)[0]
if affected_module_node.node_name in self._wq_affected_module_node_name_vs_qid_dict:
qid = self._wq_affected_module_node_name_vs_qid_dict[
affected_module_node.node_name]
if len(constraints.get_all_unique_bitwidths(qid)) != 1:
self._weight_quantizers_in_execution_order_per_scope[
scope] = quantizer
self._weight_quantizers_in_execution_order[
qid] = quantizer
else:
self._skipped_quantized_weight_node_names.append(
affected_module_node.node_name)
self._skipped_weight_quantizers[qid] = quantizer
def test_pruning_node_selector(
test_input_info_struct_: GroupPruningModulesTestStruct):
model = test_input_info_struct_.model
non_pruned_module_nodes = test_input_info_struct_.non_pruned_module_nodes
pruned_groups_by_node_id = test_input_info_struct_.pruned_groups_by_node_id
prune_first, prune_downsample = test_input_info_struct_.prune_params
pruning_operations = [v.op_func_name for v in NNCF_PRUNING_MODULES_DICT]
grouping_operations = PTElementwisePruningOp.get_all_op_aliases()
from nncf.common.pruning.node_selector import PruningNodeSelector
pruning_node_selector = PruningNodeSelector(PT_PRUNING_OPERATOR_METATYPES,
pruning_operations,
grouping_operations, None,
None, prune_first,
prune_downsample)
model = model()
model.eval()
nncf_network = NNCFNetwork(model,
input_infos=[ModelInputInfo([1, 1, 8, 8])])
graph = nncf_network.get_original_graph()
pruning_groups = pruning_node_selector.create_pruning_groups(graph)
# 1. Check all not pruned modules
all_pruned_nodes = pruning_groups.get_all_nodes()
all_pruned_modules = [
nncf_network.get_containing_module(node.node_name)
for node in all_pruned_nodes
]
for node_name in non_pruned_module_nodes:
module = nncf_network.get_containing_module(node_name)
assert module is not None and module not in all_pruned_modules
# 2. Check that all pruned groups are valid
for group_by_id in pruned_groups_by_node_id:
first_node_id = group_by_id[0]
cluster = pruning_groups.get_cluster_containing_element(first_node_id)
cluster_node_ids = [n.node_id for n in cluster.elements]
cluster_node_ids.sort()
assert Counter(cluster_node_ids) == Counter(group_by_id)
def test_get_op_nodes_in_scope():
model = TwoConvTestModel()
nncf_model = NNCFNetwork(deepcopy(model), input_infos=[ModelInputInfo([1, 1, 4, 4])]) # type: NNCFNetwork
nncf_graph = nncf_model.get_original_graph()
# Valid scopes should be successfully found
valid_nncf_modules = nncf_model.get_nncf_modules()
nodes_list = list(nncf_graph.get_all_node_ids())
for module_scope, _ in valid_nncf_modules.items():
matching_nncf_nodes = nncf_graph.get_op_nodes_in_scope(module_scope)
assert len(matching_nncf_nodes) == 1
node = matching_nncf_nodes[0]
assert isinstance(node, NNCFNode)
assert node.node_id in nodes_list
fake_model = BasicConvTestModel()
fake_nncf_model = NNCFNetwork(deepcopy(fake_model), input_infos=[ModelInputInfo([1, 1, 4, 4])])
# Not valid scopes shouldn't be found
fake_nncf_modules = fake_nncf_model.get_nncf_modules()
for module_scope, _ in fake_nncf_modules.items():
matching_nncf_nodes = nncf_graph.get_op_nodes_in_scope(module_scope)
assert not matching_nncf_nodes
def _prune_weights(self, target_model: NNCFNetwork):
target_model_graph = target_model.get_original_graph()
groups_of_nodes_to_prune = self.pruning_node_selector.create_pruning_groups(
target_model_graph)
device = next(target_model.parameters()).device
insertion_commands = []
self.pruned_module_groups_info = Clusterization[PrunedModuleInfo](
lambda x: x.node_name)
for i, group in enumerate(groups_of_nodes_to_prune.get_all_clusters()):
group_minfos = []
for node in group.elements:
node_name = node.node_name
module = target_model.get_containing_module(node_name)
module_scope = target_model_graph.get_scope_by_node_name(
node_name)
# Check that we need to prune weights in this op
assert self._is_pruned_module(module)
nncf_logger.info(
"Adding Weight Pruner in scope: {}".format(node_name))
pruning_block = self.create_weight_pruning_operation(
module, node_name)
# Hook for weights and bias
hook = UpdateWeightAndBias(pruning_block).to(device)
insertion_commands.append(
PTInsertionCommand(
PTTargetPoint(TargetType.PRE_LAYER_OPERATION,
target_node_name=node_name), hook,
TransformationPriority.PRUNING_PRIORITY))
group_minfos.append(
PrunedModuleInfo(
node_name=node_name,
module_scope=module_scope,
module=module,
operand=pruning_block,
node_id=node.node_id,
is_depthwise=is_prunable_depthwise_conv(node)))
cluster = Cluster[PrunedModuleInfo](
i, group_minfos, [n.node_id for n in group.elements])
self.pruned_module_groups_info.add_cluster(cluster)
# Propagate masks to find norm layers to prune
init_output_masks_in_graph(
target_model_graph, self.pruned_module_groups_info.get_all_nodes())
MaskPropagationAlgorithm(
target_model_graph, PT_PRUNING_OPERATOR_METATYPES,
PTNNCFPruningTensorProcessor).mask_propagation()
# Adding binary masks also for Batch/Group Norms to allow applying masks after propagation
types_to_apply_mask = ['group_norm']
if self.prune_batch_norms:
types_to_apply_mask.append('batch_norm')
all_norm_layers = target_model_graph.get_nodes_by_types(
types_to_apply_mask)
for node in all_norm_layers:
if node.data['output_mask'] is None:
# Skip elements that will not be pruned
continue
node_name = node.node_name
module = target_model.get_containing_module(node_name)
pruning_block = self.create_weight_pruning_operation(
module, node_name)
# Hook for weights and bias
hook = UpdateWeightAndBias(pruning_block).to(device)
insertion_commands.append(
PTInsertionCommand(
PTTargetPoint(TargetType.PRE_LAYER_OPERATION,
target_node_name=node_name), hook,
TransformationPriority.PRUNING_PRIORITY))
self._pruned_norms_operators.append((node, pruning_block, module))
return insertion_commands
def get_building_blocks(
compressed_model: NNCFNetwork,
max_block_size: int = 50,
allow_nested_blocks: bool = True,
allow_linear_combination: bool = False) -> List[BuildingBlock]:
"""
This algorithm finds building blocks based on the analysis of the transformed graph.
A building block is a block that satisfies the following rules:
- has one input and one output tensors
- input and output tensors shape are the same
- removing a block from the graph (that is, the layers included in the block are not executed)
does not lead to duplication of edges along which the same tensor flows
- removing a block from the graph (that is, the layers included in the block are not executed)
does not lead to dangling edges
- removing a block from the graph (that is, the layers included in the block are not executed)
does not lead to duplicate activation layers
"""
orig_graph = compressed_model.get_original_graph() # PTNNCFGraph
sgraph = get_search_graph(orig_graph)
fn_rules = [
check_graph_has_no_duplicate_edges_after_block_removal,
check_graph_has_no_act_layer_duplication_after_block_removal,
check_graph_has_no_hanging_edges_after_block_removal
]
blocks = []
act_input_shape, act_output_shape = get_potential_candidate_for_block(
sgraph)
for shape, start_nodes in act_input_shape.items():
for start_node in start_nodes:
pred_start_node = sgraph.get_prev_nodes(start_node.node_key)
if start_node.node_type == IGNORED_NAME_OPERATORS or len(
pred_start_node) != 1:
continue
for end_node in act_output_shape[shape]:
if end_node.main_id - start_node.main_id > max_block_size:
continue
if end_node.node_type in IGNORED_NAME_OPERATORS:
continue
if end_node.main_id <= start_node.main_id:
continue
is_one_edge = (end_node.main_id - start_node.bottom_id) == 1
# CHECK RULES
all_rules_is_true = True
for rule_fn in fn_rules:
if not rule_fn(sgraph, start_node, end_node):
all_rules_is_true = False
break
if all_rules_is_true and not is_one_edge:
blocks.append(BuildingBlock(start_node, end_node))
sorted_blocks = sorted(blocks, key=cmp_to_key(compare_for_building_block))
if not allow_linear_combination:
sorted_blocks = remove_linear_combination(sorted_blocks)
if not allow_nested_blocks:
sorted_blocks = remove_nested_blocks(sorted_blocks)
building_blocks_in_orig_graph = restore_node_name_in_orig_graph(
sorted_blocks, orig_graph)
return building_blocks_in_orig_graph
