From 5de6f35559ab200fa109d27225d2a1ad290c8952 Mon Sep 17 00:00:00 2001
From: Yaman Umuroglu <maltanar@gmail.com>
Date: Thu, 30 Apr 2020 14:48:59 +0100
Subject: [PATCH] [Transform] add quantization annotation for TopK indices
---
src/finn/transformation/insert_topk.py | 6 ++++++
1 file changed, 6 insertions(+)
diff --git a/src/finn/transformation/insert_topk.py b/src/finn/transformation/insert_topk.py
index 84906c7c2..213d2cedf 100644
--- a/src/finn/transformation/insert_topk.py
+++ b/src/finn/transformation/insert_topk.py
@@ -32,6 +32,7 @@ from onnx import TensorProto
from onnx import helper as oh
from finn.transformation import Transformation
+from finn.core.datatype import DataType
class InsertTopK(Transformation):
@@ -87,4 +88,9 @@ class InsertTopK(Transformation):
# replace the existing output definition with topk indices
model.graph.output.insert(0, topk_indices)
model.graph.output.pop(1)
+ # set quantization annotation for indices
+ # minimal output dtype for TopK indices dependens on num. classes
+ # assuming UINT32 is large enough for now (FINN has currently no
+ # DataType.INT64)
+ model.set_tensor_datatype(topk_indices.name, DataType.UINT32)
return (model, True)
--
GitLab