def _sparsify_weights( self, target_model: NNCFNetwork) -> List[PTInsertionCommand]: device = next(target_model.parameters()).device sparsified_module_nodes = target_model.get_weighted_original_graph_nodes( nncf_module_names=self.compressed_nncf_module_names) insertion_commands = [] for module_node in sparsified_module_nodes: node_name = module_node.node_name if not self._should_consider_scope(node_name): nncf_logger.info( "Ignored adding Weight Sparsifier in scope: {}".format( node_name)) continue nncf_logger.info( "Adding Weight Sparsifier in scope: {}".format(node_name)) compression_lr_multiplier = \ self.config.get_redefinable_global_param_value_for_algo('compression_lr_multiplier', self.name) operation = self.create_weight_sparsifying_operation( module_node, compression_lr_multiplier) hook = operation.to(device) insertion_commands.append( PTInsertionCommand( PTTargetPoint(TargetType.OPERATION_WITH_WEIGHTS, target_node_name=node_name), hook, TransformationPriority.SPARSIFICATION_PRIORITY)) sparsified_module = target_model.get_containing_module(node_name) self._sparsified_module_info.append( SparseModuleInfo(node_name, sparsified_module, hook)) return insertion_commands
def add_adjust_padding_nodes(bitwidth_graph: nx.DiGraph, model: NNCFNetwork) -> nx.DiGraph(): # pylint:disable=protected-access NewNodeArgs = namedtuple('NewNodeArgs', ('node_key', 'attr', 'parent_node_key')) nncf_graph = model.get_graph() args = [] for node_key in bitwidth_graph.nodes: node = nncf_graph.get_node_by_key(node_key) module = model.get_containing_module(node.node_name) if isinstance(module, NNCFConv2d): adjust_padding_ops = filter( lambda x: isinstance(x, UpdatePaddingValue), module.pre_ops.values()) for _ in adjust_padding_ops: new_node_key = f'{node_key}_apad' attr = dict(type='', label='adjust_padding_value', style='filled', color='yellow') args.append(NewNodeArgs(new_node_key, attr, node_key)) for arg in args: bitwidth_graph.add_node(arg.node_key, **arg.attr) bitwidth_graph.add_edge(arg.node_key, arg.parent_node_key) return bitwidth_graph
def input_prune(cls, model: NNCFNetwork, node: NNCFNode, graph: NNCFGraph) -> None: input_mask = node.data['input_masks'][0] if input_mask is None: return bool_mask = torch.tensor(input_mask, dtype=torch.bool) new_num_channels = int(torch.sum(input_mask)) is_depthwise = is_prunable_depthwise_conv(node) node_module = model.get_containing_module(node.node_name) old_num_channels = int(node_module.weight.size(1)) if is_depthwise: # In depthwise case prune output channels by input mask, here only fix for new number of input channels node_module.groups = new_num_channels node_module.in_channels = new_num_channels old_num_channels = int(node_module.weight.size(0)) else: out_channels = node_module.weight.size(0) broadcasted_mask = bool_mask.repeat(out_channels).view( out_channels, bool_mask.size(0)) new_weight_shape = list(node_module.weight.shape) new_weight_shape[1] = new_num_channels node_module.in_channels = new_num_channels node_module.weight = torch.nn.Parameter( node_module.weight[broadcasted_mask].view(new_weight_shape)) nncf_logger.info( 'Pruned Convolution {} by input mask. Old input filters number: {}, new filters number:' ' {}.'.format(node.data['key'], old_num_channels, new_num_channels))
def output_prune(cls, model: NNCFNetwork, node: NNCFNode, graph: NNCFGraph) -> None: output_mask = node.data['output_mask'] if output_mask is None: return bool_mask = torch.tensor(output_mask, dtype=torch.bool) new_num_channels = int(torch.sum(bool_mask)) node_module = model.get_containing_module(node.node_name) old_num_clannels = int(node_module.weight.size(1)) in_channels = node_module.weight.size(0) broadcasted_mask = bool_mask.repeat(in_channels).view( in_channels, bool_mask.size(0)) new_weight_shape = list(node_module.weight.shape) new_weight_shape[1] = new_num_channels node_module.out_channels = new_num_channels node_module.weight = torch.nn.Parameter( node_module.weight[broadcasted_mask].view(new_weight_shape)) if node_module.bias is not None: node_module.bias = torch.nn.Parameter(node_module.bias[bool_mask]) nncf_logger.info( 'Pruned ConvTranspose {} by pruning mask. Old output filters number: {}, new filters number:' ' {}.'.format(node.data['key'], old_num_clannels, node_module.out_channels))
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 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 input_prune(cls, model: NNCFNetwork, node: NNCFNode, graph: NNCFGraph) -> None: input_mask = node.data['input_masks'][0] if input_mask is None: return bool_mask = torch.tensor(input_mask, dtype=torch.bool) node_module = model.get_containing_module(node.node_name) old_num_clannels = int(node_module.weight.size(0)) node_module.in_channels = int(torch.sum(bool_mask)) node_module.weight = torch.nn.Parameter(node_module.weight[bool_mask]) nncf_logger.info( 'Pruned ConvTranspose {} by input mask. Old input filters number: {}, new filters number:' ' {}.'.format(node.data['key'], old_num_clannels, node_module.in_channels))
def output_prune(cls, model: NNCFNetwork, node: NNCFNode, graph: NNCFGraph) -> None: mask = node.data['output_mask'] if mask is None: return bool_mask = torch.tensor(mask, dtype=torch.bool) node_module = model.get_containing_module(node.node_name) old_num_clannels = int(node_module.weight.size(0)) node_module.out_channels = int(torch.sum(mask)) node_module.weight = torch.nn.Parameter(node_module.weight[bool_mask]) if node_module.bias is not None: node_module.bias = torch.nn.Parameter(node_module.bias[bool_mask]) nncf_logger.info( 'Pruned Convolution {} by pruning mask. Old output filters number: {}, new filters number:' ' {}.'.format(node.data['key'], old_num_clannels, node_module.out_channels))
def input_prune(cls, model: NNCFNetwork, node: NNCFNode, graph: NNCFGraph) -> None: input_mask = node.data['input_masks'][0] if input_mask is None: return bool_mask = torch.tensor(input_mask, dtype=torch.bool) node_module = model.get_containing_module(node.node_name) if isinstance(node_module, tuple(NNCF_WRAPPED_USER_MODULES_DICT)): assert node_module.target_weight_dim_for_compression == 0, \ "Implemented only for target_weight_dim_for_compression == 0" old_num_clannels = int(node_module.weight.size(0)) new_num_channels = int(torch.sum(input_mask)) node_module.weight = torch.nn.Parameter( node_module.weight[bool_mask]) node_module.n_channels = new_num_channels nncf_logger.info( 'Pruned Elementwise {} by input mask. Old num features: {}, new num features:' ' {}.'.format(node.data['key'], old_num_clannels, new_num_channels))
def input_prune(cls, model: NNCFNetwork, node: NNCFNode, graph: NNCFGraph) -> None: input_mask = node.data['input_masks'][0] if input_mask is None: return node_module = model.get_containing_module(node.node_name) bool_mask = torch.tensor(input_mask, dtype=torch.bool) old_num_clannels = int(node_module.weight.size(0)) new_num_channels = int(torch.sum(input_mask)) node_module.num_channels = new_num_channels node_module.num_groups = new_num_channels node_module.weight = torch.nn.Parameter(node_module.weight[bool_mask]) node_module.bias = torch.nn.Parameter(node_module.bias[bool_mask]) nncf_logger.info( 'Pruned GroupNorm {} by input mask. Old num features: {}, new num features:' ' {}.'.format(node.data['key'], old_num_clannels, new_num_channels))
def __init__(self, algo_ctrl: QuantizationController, model: NNCFNetwork, groups_of_adjacent_quantizers: GroupsOfAdjacentQuantizers, add_flops=False): # pylint:disable=too-many-branches # pylint:disable=too-many-statements nncf_graph = model.get_graph() self._nx_graph = nncf_graph.get_graph_for_structure_analysis() if add_flops: flops_per_module = model.get_flops_per_module() flops_vs_node_group = defaultdict(set) # type: Dict[int, Tuple[int, Set[NNCFNode]]] for idx, module_node_name_and_flops in enumerate(flops_per_module.items()): module_node_name, flops = module_node_name_and_flops node_set = set(nncf_graph.get_op_nodes_in_scope(nncf_graph.get_scope_by_node_name(module_node_name))) flops_vs_node_group[idx] = (flops, node_set) grouped_mode = bool(groups_of_adjacent_quantizers) for node_key in nncf_graph.get_all_node_keys(): node = nncf_graph.get_node_by_key(node_key) color = '' operator_name = node.node_type module = model.get_containing_module(node.node_name) if isinstance(module, NNCFConv2d): color = 'lightblue' if module.groups == module.in_channels and module.in_channels > 1: operator_name = 'DW_Conv2d' color = 'purple' kernel_size = 'x'.join(map(str, module.kernel_size)) operator_name += f'_k{kernel_size}' padding_values = set(module.padding) padding_enabled = len(padding_values) >= 1 and padding_values.pop() if padding_enabled: operator_name += '_PAD' if add_flops: matches = [f_nodes_tpl for idx, f_nodes_tpl in flops_vs_node_group.items() if node in f_nodes_tpl[1]] assert len(matches) == 1 flops, affected_nodes = next(iter(matches)) operator_name += f'_FLOPS:{str(flops)}' if len(affected_nodes) > 1: node_ids = sorted([n.node_id for n in affected_nodes]) operator_name += "(shared among nodes {})".format(",".join( [str(node_id) for node_id in node_ids])) operator_name += '_#{}'.format(node.node_id) target_node_to_draw = self._nx_graph.nodes[node_key] target_node_to_draw['label'] = operator_name target_node_to_draw['style'] = 'filled' if color: target_node_to_draw['color'] = color non_weight_quantizers = algo_ctrl.non_weight_quantizers bitwidth_color_map = {2: 'purple', 4: 'red', 8: 'green', 6: 'orange'} for quantizer_id, quantizer_info in non_weight_quantizers.items(): self._paint_activation_quantizer_node(nncf_graph, quantizer_id, quantizer_info, bitwidth_color_map, groups_of_adjacent_quantizers) for wq_id, wq_info in algo_ctrl.weight_quantizers.items(): nodes = [nncf_graph.get_node_by_name(tp.target_node_name) for tp in wq_info.affected_insertions] if not nodes: raise AttributeError('Failed to get affected nodes for quantized module node: {}'.format( wq_id.target_node_name)) preds = [nncf_graph.get_previous_nodes(node) for node in nodes] wq_nodes = [] for pred_list in preds: for pred_node in pred_list: if 'UpdateWeight' in pred_node.node_name: wq_nodes.append(pred_node) assert len(wq_nodes) == 1 node = wq_nodes[0] node_id = node.node_id key = nncf_graph.get_node_key_by_id(node_id) nx_node_to_draw_upon = self._nx_graph.nodes[key] quantizer = wq_info.quantizer_module_ref bitwidths = quantizer.num_bits nx_node_to_draw_upon['label'] = 'WFQ_[{}]_#{}'.format(quantizer.get_quantizer_config(), str(node_id)) if grouped_mode: group_id_str = 'UNDEFINED' group_id = groups_of_adjacent_quantizers.get_group_id_for_quantizer(wq_id) if group_id is None: nncf_logger.error('No group for weight quantizer for: {}'.format(wq_id)) else: group_id_str = str(group_id) nx_node_to_draw_upon['label'] += '_G' + group_id_str nx_node_to_draw_upon['color'] = bitwidth_color_map[bitwidths] nx_node_to_draw_upon['style'] = 'filled'
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