nncf/common/pruning/utils.py
def count_flops_and_weights_per_node(graph: NNCFGraph,
                                     input_shapes: Dict[NNCFNodeName, List[int]],
                                     output_shapes: Dict[NNCFNodeName, List[int]],
                                     conv_op_metatypes: List[Type[OperatorMetatype]],
                                     linear_op_metatypes: List[Type[OperatorMetatype]],
                                     input_channels: Dict[NNCFNodeName, int] = None,
                                     output_channels: Dict[NNCFNodeName, int] = None,
                                     kernel_sizes: Dict[NNCFNodeName, Tuple[int, int]] = None,
                                     op_addresses_to_skip: List[NNCFNodeName] = None) -> \
    Tuple[Dict[NNCFNodeName, int], Dict[NNCFNodeName, int]]:
    """
    Counts the number weights and FLOPs per node in the model for convolution and fully connected layers.

    :param graph: NNCFGraph.
    :param input_shapes: Dictionary of input dimension shapes for convolutions and
        fully connected layers. E.g {node_name: (height, width)}
    :param output_shapes: Dictionary of output dimension shapes for convolutions and
        fully connected layers. E.g {node_name: (height, width)}
    :param conv_op_metatypes: List of metatypes defining convolution operations.
    :param linear_op_metatypes: List of metatypes defining linear/fully connected operations.
    :param input_channels: Dictionary of input channels number in convolutions.
        If not specified, taken from the graph. {node_name: channels_num}
    :param output_channels: Dictionary of output channels number in convolutions.
        If not specified, taken from the graph. {node_name: channels_num}
    :param kernel_sizes: Dictionary of kernel sizes in convolutions.
        If not specified, taken from the graph. {node_name: kernel_size}. It's only supposed to be used in NAS in case
        of Elastic Kernel enabled.
    :param op_addresses_to_skip: List of operation addresses of layers that should be skipped from calculation.
        It's only supposed to be used in NAS in case of Elastic Depth enabled.
    :return Dictionary of FLOPs number {node_name: flops_num}
            Dictionary of weights number {node_name: weights_num}
    """
    flops = {}
    weights = {}
    input_channels = input_channels or {}
    output_channels = output_channels or {}
    kernel_sizes = kernel_sizes or {}
    op_addresses_to_skip = op_addresses_to_skip or []
    for node in graph.get_nodes_by_metatypes(conv_op_metatypes):
        name = node.node_name
        if name in op_addresses_to_skip:
            continue
        num_in_channels = input_channels.get(name, node.layer_attributes.in_channels)
        num_out_channels = output_channels.get(name, node.layer_attributes.out_channels)
        kernel_size = kernel_sizes.get(name, node.layer_attributes.kernel_size)
        if is_prunable_depthwise_conv(node):
            # Prunable depthwise conv processed in special way
            # because common way to calculate filters per
            # channel for such layer leads to zero in case
            # some of the output channels are pruned.
            filters_per_channel = 1
        else:
            filters_per_channel = num_out_channels // node.layer_attributes.groups

        flops_numpy = 2 * np.prod(kernel_size) * num_in_channels * filters_per_channel * np.prod(output_shapes[name])
        weights_numpy = np.prod(kernel_size) * num_in_channels * filters_per_channel
        flops[name] = flops_numpy.astype(int).item()
        weights[name] = weights_numpy.astype(int).item()

    for node in graph.get_nodes_by_metatypes(linear_op_metatypes):
        name = node.node_name
        if name in op_addresses_to_skip:
            continue
        flops_numpy = 2 * np.prod(input_shapes[name]) * np.prod(output_shapes[name])
        weights_numpy = np.prod(input_shapes[name]) * np.prod(output_shapes[name])
        flops[name] = flops_numpy.astype(int).item()
        weights[name] = weights_numpy.astype(int).item()

    return flops, weights

    Counts the number weights and FLOPs per node in the model for convolution and fully connected layers.

    :param graph: NNCFGraph.
    :param input_shapes: Dictionary of input dimension shapes for convolutions and
        fully connected layers. E.g {node_name: (height, width)}
    :param output_shapes: Dictionary of output dimension shapes for convolutions and
        fully connected layers. E.g {node_name: (height, width)}
    :param conv_op_metatypes: List of metatypes defining convolution operations.
    :param linear_op_metatypes: List of metatypes defining linear/fully connected operations.
    :param input_channels: Dictionary of input channels number in convolutions.
        If not specified, taken from the graph. {node_name: channels_num}
    :param output_channels: Dictionary of output channels number in convolutions.
        If not specified, taken from the graph. {node_name: channels_num}
    :param kernel_sizes: Dictionary of kernel sizes in convolutions.
        If not specified, taken from the graph. {node_name: kernel_size}. It's only supposed to be used in NAS in case
        of Elastic Kernel enabled.
    :param op_addresses_to_skip: List of operation addresses of layers that should be skipped from calculation.
        It's only supposed to be used in NAS in case of Elastic Depth enabled.
    :return Dictionary of FLOPs number {node_name: flops_num}
            Dictionary of weights number {node_name: weights_num}
    
Counts the number weights and FLOPs per node in the model for convolution and fully connected layers.
Counts the number weights and FLOPs per node in the model for convolution and fully connected layers .
def count_flops_and_weights_per_node(graph: NNCFGraph,
                                     input_shapes: Dict[NNCFNodeName, List[int]],
                                     output_shapes: Dict[NNCFNodeName, List[int]],
                                     conv_op_metatypes: List[Type[OperatorMetatype]],
                                     linear_op_metatypes: List[Type[OperatorMetatype]],
                                     input_channels: Dict[NNCFNodeName, int] = None,
                                     output_channels: Dict[NNCFNodeName, int] = None,
                                     kernel_sizes: Dict[NNCFNodeName, Tuple[int, int]] = None,
                                     op_addresses_to_skip: List[NNCFNodeName] = None) -> \
    Tuple[Dict[NNCFNodeName, int], Dict[NNCFNodeName, int]]:
    flops = {}
    weights = {}
    input_channels = input_channels or {}
    output_channels = output_channels or {}
    kernel_sizes = kernel_sizes or {}
    op_addresses_to_skip = op_addresses_to_skip or []
    for node in graph.get_nodes_by_metatypes(conv_op_metatypes):
        name = node.node_name
        if name in op_addresses_to_skip:
            continue
        num_in_channels = input_channels.get(name, node.layer_attributes.in_channels)
        num_out_channels = output_channels.get(name, node.layer_attributes.out_channels)
        kernel_size = kernel_sizes.get(name, node.layer_attributes.kernel_size)
        if is_prunable_depthwise_conv(node):
            filters_per_channel = 1
        else:
            filters_per_channel = num_out_channels // node.layer_attributes.groups
        flops_numpy = 2 * np.prod(kernel_size) * num_in_channels * filters_per_channel * np.prod(output_shapes[name])
        weights_numpy = np.prod(kernel_size) * num_in_channels * filters_per_channel
        flops[name] = flops_numpy.astype(int).item()
        weights[name] = weights_numpy.astype(int).item()
    for node in graph.get_nodes_by_metatypes(linear_op_metatypes):
        name = node.node_name
        if name in op_addresses_to_skip:
            continue
        flops_numpy = 2 * np.prod(input_shapes[name]) * np.prod(output_shapes[name])
        weights_numpy = np.prod(input_shapes[name]) * np.prod(output_shapes[name])
        flops[name] = flops_numpy.astype(int).item()
        weights[name] = weights_numpy.astype(int).item()
    return flops, weights
def count_flops_and_weights_per_node ( graph : NNCFGraph , input_shapes : Dict [ NNCFNodeName , List [ int ] ] , output_shapes : Dict [ NNCFNodeName , List [ int ] ] , conv_op_metatypes : List [ Type [ OperatorMetatype ] ] , linear_op_metatypes : List [ Type [ OperatorMetatype ] ] , input_channels : Dict [ NNCFNodeName , int ] = None , output_channels : Dict [ NNCFNodeName , int ] = None , kernel_sizes : Dict [ NNCFNodeName , Tuple [ int , int ] ] = None , op_addresses_to_skip : List [ NNCFNodeName ] = None ) -> Tuple [ Dict [ NNCFNodeName , int ] , Dict [ NNCFNodeName , int ] ] : flops = { } weights = { } input_channels = input_channels or { } output_channels = output_channels or { } kernel_sizes = kernel_sizes or { } op_addresses_to_skip = op_addresses_to_skip or [ ] for node in graph . get_nodes_by_metatypes ( conv_op_metatypes ) : name = node . node_name if name in op_addresses_to_skip : continue num_in_channels = input_channels . get ( name , node . layer_attributes . in_channels ) num_out_channels = output_channels . get ( name , node . layer_attributes . out_channels ) kernel_size = kernel_sizes . get ( name , node . layer_attributes . kernel_size ) if is_prunable_depthwise_conv ( node ) : filters_per_channel = 1 else : filters_per_channel = num_out_channels // node . layer_attributes . groups flops_numpy = 2 * np . prod ( kernel_size ) * num_in_channels * filters_per_channel * np . prod ( output_shapes [ name ] ) weights_numpy = np . prod ( kernel_size ) * num_in_channels * filters_per_channel flops [ name ] = flops_numpy . astype ( int ) . item ( ) weights [ name ] = weights_numpy . astype ( int ) . item ( ) for node in graph . get_nodes_by_metatypes ( linear_op_metatypes ) : name = node . node_name if name in op_addresses_to_skip : continue flops_numpy = 2 * np . prod ( input_shapes [ name ] ) * np . prod ( output_shapes [ name ] ) weights_numpy = np . prod ( input_shapes [ name ] ) * np . prod ( output_shapes [ name ] ) flops [ name ] = flops_numpy . astype ( int ) . item ( ) weights [ name ] = weights_numpy . astype ( int ) . item ( ) return flops , weights
Counts the number weights and FLOPs per node in the model for convolution and fully connected layers.
Counts the number weights and FLOPs per node in the model for convolution and fully connected layers .
"""
    Counts the number weights and FLOPs per node in the model for convolution and fully connected layers.

    :param graph: NNCFGraph.
    :param input_shapes: Dictionary of input dimension shapes for convolutions and
        fully connected layers. E.g {node_name: (height, width)}
    :param output_shapes: Dictionary of output dimension shapes for convolutions and
        fully connected layers. E.g {node_name: (height, width)}
    :param conv_op_metatypes: List of metatypes defining convolution operations.
    :param linear_op_metatypes: List of metatypes defining linear/fully connected operations.
    :param input_channels: Dictionary of input channels number in convolutions.
        If not specified, taken from the graph. {node_name: channels_num}
    :param output_channels: Dictionary of output channels number in convolutions.
        If not specified, taken from the graph. {node_name: channels_num}
    :param kernel_sizes: Dictionary of kernel sizes in convolutions.
        If not specified, taken from the graph. {node_name: kernel_size}. It's only supposed to be used in NAS in case
        of Elastic Kernel enabled.
    :param op_addresses_to_skip: List of operation addresses of layers that should be skipped from calculation.
        It's only supposed to be used in NAS in case of Elastic Depth enabled.
    :return Dictionary of FLOPs number {node_name: flops_num}
            Dictionary of weights number {node_name: weights_num}
    """ # Prunable depthwise conv processed in special way # because common way to calculate filters per # channel for such layer leads to zero in case # some of the output channels are pruned.
graph
NNCFGraph
input_shapes
Dict[NNCFNodeName, List[int]]
output_shapes
Dict[NNCFNodeName, List[int]]
conv_op_metatypes
List[Type[OperatorMetatype]]
linear_op_metatypes
List[Type[OperatorMetatype]]
input_channels
Dict[NNCFNodeName, int]
output_channels
Dict[NNCFNodeName, int]
kernel_sizes
Dict[NNCFNodeName, Tuple[int, int]]
op_addresses_to_skip
List[NNCFNodeName]