ONNX (finn): add QuantIdentity handler

Signed-off-by: Alessandro Pappalardo <[email protected]>
Xilinx · Sep 8, 2020 · 59f1fdd · 59f1fdd
1 parent 7f12633
commit 59f1fdd
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/brevitas/onnx/finn/handler/act.py b/brevitas/onnx/finn/handler/act.py
@@ -3,7 +3,7 @@
 import torch
 from torch import Tensor
 
-from brevitas.nn import QuantReLU, QuantHardTanh
+from brevitas.nn import QuantReLU, QuantHardTanh, QuantIdentity
 from .base import FINNQuantInputHandler, FINNQuantIOHandler
 from ..function.act import QuantReLUPlaceholderFunction, QuantHardTanhPlaceholderFunction
 
@@ -133,3 +133,8 @@ def prepare_for_symbolic_execution(self, module: QuantHardTanh):
     def symbolic_execution(self, inp: Tensor):
         ret = QuantHardTanhPlaceholderFunction.apply(inp, *self.symbolic_kwargs.values())
         return ret
+
+
+class FINNQuantIdentityHandler(FINNQuantHardTanhHandler):
+    handled_layer = QuantIdentity
+
diff --git a/brevitas/onnx/finn/manager.py b/brevitas/onnx/finn/manager.py
@@ -4,7 +4,7 @@
 from .transform import move_quant_attributes_into_annotations
 from .handler.parameter import FINNQuantConv2dHandler, FINNQuantLinearHandler
 from .handler.parameter import FINNQuantConv1dHandler
-from .handler.act import FINNQuantHardTanhHandler, FINNQuantReLUHandler
+from .handler.act import FINNQuantHardTanhHandler, FINNQuantReLUHandler, FINNQuantIdentityHandler
 from .handler.acc import FINNQuantAvgPool2dHandler
 
 
@@ -15,6 +15,7 @@ class FINNManager(BaseManager):
         FINNQuantConv1dHandler,
         FINNQuantConv2dHandler,
         FINNQuantReLUHandler,
+        FINNQuantIdentityHandler,
         FINNQuantHardTanhHandler,
         FINNQuantAvgPool2dHandler]