Skip to content

Commit 653d89e

Browse files
roman-janik-nxpStrycekSimon
authored andcommitted
NXP backend: Improve quantization annotation process
- fixes multiple subsequent nodes with SharedSpecPattern problem
1 parent 00491fd commit 653d89e

File tree

1 file changed

+26
-1
lines changed

1 file changed

+26
-1
lines changed

backends/nxp/quantizer/neutron_quantizer.py

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import (
1111
NeutronAtenPassManager,
1212
)
13-
1413
from executorch.backends.nxp.quantizer.patterns import (
1514
AddmmPattern,
1615
AvgPoolPattern,
@@ -24,6 +23,7 @@
2423
ReluInPlacePattern,
2524
ReluPattern,
2625
ReshapePattern,
26+
SharedSpecPattern,
2727
SoftMaxPattern,
2828
ViewPattern,
2929
)
@@ -204,9 +204,34 @@ def __init__(self):
204204
NeutronAtenQuantizer(ViewPattern(), static_qconfig),
205205
]
206206
)
207+
self.op_to_quantizer = {
208+
pt: q for q in self.quantizers for pt in q.pattern.partition_types()
209+
}
210+
self.op_to_applied_quantizer = {
211+
pt: False for q in self.quantizers for pt in q.pattern.partition_types()
212+
}
207213

208214
def transform_for_annotation(
209215
self, model: torch.fx.GraphModule
210216
) -> torch.fx.GraphModule:
211217
pass_runner = NeutronAtenPassManager()
212218
return pass_runner(model).graph_module
219+
220+
def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
221+
nodes = list(model.graph.nodes)
222+
for node in nodes:
223+
if (
224+
node.target not in self.op_to_quantizer
225+
or self.op_to_applied_quantizer[node.target]
226+
):
227+
continue
228+
else:
229+
quantizer = self.op_to_quantizer[node.target]
230+
quantizer.annotate(model)
231+
if not isinstance(quantizer.pattern, SharedSpecPattern):
232+
self.op_to_applied_quantizer[node.target] = True
233+
234+
return model
235+
236+
def validate(self, model: torch.fx.GraphModule) -> None:
237+
return super().validate(model)

0 commit comments

Comments
 (0)