No public description

PiperOrigin-RevId: 615892230 Change-Id: I5c1bd1a4332d6925158985ba468e2422aeebe4f4
google · Mar 14, 2024 · d389659 · d389659
1 parent dc0bce9
commit d389659
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 9 deletions.
diff --git a/qkeras/qtools/generate_layer_data_type_map.py b/qkeras/qtools/generate_layer_data_type_map.py
@@ -57,13 +57,16 @@ class TagMissingError(ValueError):
     "QConv1D",
     "QConv2D",
     "QDepthwiseConv2D",
+    "QConv2D",
+    "QConv2DTranspose",
 ]
 
 KERAS_LAYERS = [
     "Dense",
     "Conv1D",
     "Conv2D",
     "DepthwiseConv2D",
+    "Conv2DTranspose",
 ]
 
 
@@ -507,10 +510,10 @@ def generate_layer_data_type_map(
         output_quantizer = update_output_quantizer_in_graph(
             graph, node_id, quantizer_factory, input_quantizer, for_reference)
         layer_data_type_map[layer] = {
-          "input_quantizer_list": input_quantizer_list,
-          "output_quantizer": output_quantizer,
-          "output_shapes": input_shape,
-          "operation_count": operation_count
+            "input_quantizer_list": input_quantizer_list,
+            "output_quantizer": output_quantizer,
+            "output_shapes": input_shape,
+            "operation_count": operation_count
         }
       else:
         (gamma_quantizer, beta_quantizer, mean_quantizer, variance_quantizer,
@@ -552,19 +555,23 @@ def set_output(op, output):
           if keras_accumulator:
             set_output(
                 qbn.internal_divide_quantizer,
-                quantizer_factory.make_default_quantizer(mode=keras_accumulator))
+                quantizer_factory.make_default_quantizer(
+                    mode=keras_accumulator))
 
             set_output(
                 qbn.internal_multiplier,
-                quantizer_factory.make_default_quantizer(mode=keras_accumulator))
+                quantizer_factory.make_default_quantizer(
+                    mode=keras_accumulator))
 
             set_output(
                 qbn.internal_accumulator,
-                quantizer_factory.make_default_quantizer(mode=keras_accumulator))
+                quantizer_factory.make_default_quantizer(
+                    mode=keras_accumulator))
 
             set_output(
                 qbn.internal_output.output,
-                quantizer_factory.make_default_quantizer(mode=keras_accumulator))
+                quantizer_factory.make_default_quantizer(
+                    mode=keras_accumulator))
 
         gamma_range = None
         if hasattr(layer, "gamma_range"):

diff --git a/qkeras/qtools/qtools_util.py b/qkeras/qtools/qtools_util.py
@@ -150,7 +150,9 @@ def get_operation_count(layer, input_shape):
         "BatchNormalization" in layer.__class__.__name__):
     operation_count = np.prod(input_shape[1:])
 
-  elif layer.__class__.__name__ in ["QConv2D", "Conv2D", "QConv2DBatchnorm"]:
+  elif layer.__class__.__name__ in [
+      "QConv2D", "Conv2D", "QConv2DBatchnorm",
+      "QConv2DTranspose", "Conv2DTranspose"]:
 
     output_shape = layer.compute_output_shape(input_shape)
     _, _, _, channels_i = input_shape