Skip to content

Commit

Permalink
[2/N] Added automatic structured pruning for DistillBERT, CLIP. (open…
Browse files Browse the repository at this point in the history
…vinotoolkit#1680)

### Changes

Continue of the openvinotoolkit#1657 
Supported finding pruning groups for DistillBERT, CLIP models
Proper filtering groups - no opened/closed branches, invalidate when
reach "bad" operation (stop_mask, ignored_scope, output).
Introduced block hierarchy class for finding groups on the leaves and
for better visualization.
Refactored the functionality.

Block hierarchy with colored nodes: red - for invalid, green - for final
leaves, grey - intermediate, not active.

![image](https://user-images.githubusercontent.com/4014476/229438402-34eedce5-81cc-4fae-9aed-7d3bc03c9213.png)

added rules for DistillBERT case when pruning masks intersect with
attention masks:

![image](https://user-images.githubusercontent.com/4014476/229438941-35e82751-f166-4028-8b2f-d4c098c9b85a.png)

### Reason for changes

Expand scope of the supported models in JPDQ

### Related tickets

106556

### Tests

test_groups
  • Loading branch information
ljaljushkin authored Apr 14, 2023
1 parent 722e5bf commit c95f5a3
Show file tree
Hide file tree
Showing 22 changed files with 1,257 additions and 454 deletions.
104 changes: 104 additions & 0 deletions nncf/experimental/common/pruning/block_hierarchy.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
"""
Copyright (c) 2023 Intel Corporation
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from pathlib import Path
from typing import Dict
from typing import List

import networkx as nx
from nncf.common.utils.dot_file_rw import write_dot_graph
from nncf.experimental.common.pruning.nodes_grouping import PropagationGroup


class BlockHierarchy:
"""
Graph that represents the hierarchy of propagation blocks/groups
"""

PROPAGATION_GROUP = "propagation_group"

def __init__(self, root_groups: List[PropagationGroup]) -> None:
"""
Creates hierarchy of propagation blocks/groups by traversing children in the given root groups.
:param roots: list of the root groups
:return: networkx graph that represents the hierarchy of propagation blocks/groups.
"""

self._id_counter = 0
self._graph = nx.DiGraph()
self._visited_block_ids_map: Dict[int, int] = {}

for root_group in root_groups:
self._add_group_to_graph(root_group)
self._id_counter += 1

def get_groups_on_leaves(self) -> List[PropagationGroup]:
"""
Returns the list of all propagation groups on the leaves.
"""
groups = []
for node_id, data in self._graph.nodes(data=True):
is_leaf = self._graph.out_degree(node_id) == 0
if is_leaf:
groups.append(data[self.PROPAGATION_GROUP])
return groups


def visualize_graph(self, path: Path) -> None:
out_graph = self._get_graph_for_visualization()
write_dot_graph(out_graph, path)

def _add_group_to_graph(self, parent_group: PropagationGroup) -> None:
"""
Recursive helper to traverse children of the given PropagationBlock with adding them to the graph.
:param parent_group: current group for traversing.
"""
parent_graph_id = str(self._id_counter)
is_leaf = not parent_group.has_children()
attrs = {
self.PROPAGATION_GROUP: parent_group
}
self._graph.add_node(parent_graph_id, **attrs)
self._visited_block_ids_map[id(parent_group)] = parent_graph_id
if not is_leaf:
for child_group in parent_group.get_children():
child_id = id(child_group)
if child_id not in self._visited_block_ids_map:
self._id_counter += 1
child_graph_id = str(self._id_counter)
self._visited_block_ids_map[id(child_group)] = child_graph_id
self._graph.add_edge(parent_graph_id, child_graph_id)
self._add_group_to_graph(child_group)
else:
child_graph_id = self._visited_block_ids_map[child_id]
self._graph.add_edge(parent_graph_id, child_graph_id)

def _get_graph_for_visualization(self) -> nx.DiGraph:
"""
:return: A user-friendly graph .dot file, making it easier to debug the block hierarchy.
"""
out_graph = nx.DiGraph()
for node_id, node_data in self._graph.nodes(data=True):
group = node_data[self.PROPAGATION_GROUP]
is_leaf = not group.has_children()
color = 'grey'
if is_leaf:
color = 'red' if group.is_invalid else 'green'
out_graph.add_node(node_id, label=f"\"{str(group)}\"", color=color, style='filled')

for u, v in self._graph.edges:
out_graph.add_edge(u, v)

return out_graph
Loading

0 comments on commit c95f5a3

Please sign in to comment.