Class used to partition out unwanted graph nodes. e.g. - nodes are prevented from quantization annotation - nodes have been grouped together as a submodule Attributes ---------- fp_node_id_set : set a set contains nodes' name to be left in fp precision fp_n
| 76 | |
| 77 | |
| 78 | class _AnnotationSkipper(OperatorSupportBase): |
| 79 | """ |
| 80 | Class used to partition out unwanted graph nodes. |
| 81 | e.g. - nodes are prevented from quantization annotation |
| 82 | - nodes have been grouped together as a submodule |
| 83 | |
| 84 | Attributes |
| 85 | ---------- |
| 86 | fp_node_id_set : set |
| 87 | a set contains nodes' name to be left in fp precision |
| 88 | fp_node_op_set : set |
| 89 | a set contains nodes' target (aten dialect) to be left in fp precision |
| 90 | skip_annotated_submodule : bool |
| 91 | flag to skip annotated submodule or not |
| 92 | |
| 93 | Methods |
| 94 | ------- |
| 95 | should_delegate(n: torch.fx.Node) |
| 96 | identify the residual nodes haven't be lowered with fixed-precision |
| 97 | should_skip(n: torch.fx.Node) |
| 98 | identify the nodes should be kept out with fixed-precision or not |
| 99 | is_node_supported(_, node: torch.fx.Node) |
| 100 | overridden method for graph partitioning |
| 101 | """ |
| 102 | |
| 103 | def __init__( |
| 104 | self, |
| 105 | fp_node_id_set: set = None, |
| 106 | fp_node_op_set: set = None, |
| 107 | skip_annotated_submodule: bool = False, |
| 108 | ): |
| 109 | self.fp_node_id_set = fp_node_id_set |
| 110 | self.fp_node_op_set = fp_node_op_set |
| 111 | self.skip_annotated_submodule = skip_annotated_submodule |
| 112 | |
| 113 | def should_delegate(self, n: torch.fx.Node): |
| 114 | return n.op == "call_function" and n.target != operator.getitem |
| 115 | |
| 116 | def should_skip(self, n: torch.fx.Node): |
| 117 | return n.name in self.fp_node_id_set or n.target in self.fp_node_op_set |
| 118 | |
| 119 | def is_node_supported(self, _, node: torch.fx.Node) -> bool: |
| 120 | if self.skip_annotated_submodule: |
| 121 | if node.op == "get_attr": |
| 122 | return all(self.should_delegate(user) for user in node.users) |
| 123 | return self.should_delegate(node) |
| 124 | |
| 125 | if any( |
| 126 | [ |
| 127 | node.op in ("placeholder", "output"), |
| 128 | self.should_skip(node), |
| 129 | # check if parameters belong to fallbacked operator |
| 130 | ( |
| 131 | node.op == "get_attr" |
| 132 | and all(self.should_skip(user) for user in node.users) |
| 133 | ), |
| 134 | ] |
| 135 | ): |