Add more ops to calculate output scales (#32036)

PaddlePaddle · Apr 2, 2021 · cd74b20 · cd74b20
1 parent ed49b41
commit cd74b20
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 23 deletions.
diff --git a/python/paddle/fluid/contrib/slim/quantization/imperative/qat.py b/python/paddle/fluid/contrib/slim/quantization/imperative/qat.py
@@ -468,7 +468,7 @@ def _is_target_layer(self, layer):
         """
         Whether the layer needs to calculate output scales.
         """
-        return isinstance(layer, tuple(utils.quant_output_layers_map.values())) \
+        return isinstance(layer, utils.quant_output_layers) \
             or ('quantized' in layer.full_name() and \
                 'quantized_noweight' not in layer.full_name())
 

diff --git a/python/paddle/fluid/contrib/slim/quantization/imperative/utils.py b/python/paddle/fluid/contrib/slim/quantization/imperative/utils.py
@@ -43,28 +43,18 @@
     "fake_quantize_dequantize_moving_average_abs_max"
 ]
 
-quant_output_layers_map = {
-    'Conv2D': paddle.nn.Conv2D,
-    'Conv2DTranspose': paddle.nn.Conv2DTranspose,
-    'Linear': paddle.nn.Linear,
-    'AdaptiveAvgPool2D': paddle.nn.AdaptiveAvgPool2D,
-    'AdaptiveMaxPool2D': paddle.nn.AdaptiveMaxPool2D,
-    'AvgPool2D': paddle.nn.AvgPool2D,
-    'MaxPool2D': paddle.nn.MaxPool2D,
-    'BatchNorm': paddle.nn.BatchNorm,
-    'BatchNorm2D': paddle.nn.BatchNorm2D,
-    'SyncBatchNorm': paddle.nn.SyncBatchNorm,
-    'ELU': paddle.nn.ELU,
-    'GELU': paddle.nn.GELU,
-    'LeakyReLU': paddle.nn.LeakyReLU,
-    'PReLU': paddle.nn.PReLU,
-    'ReLU': paddle.nn.ReLU,
-    'ReLU6': paddle.nn.ReLU6,
-    'Sigmoid': paddle.nn.Sigmoid,
-    'Softmax': paddle.nn.Softmax,
-    'Tanh': paddle.nn.Tanh,
-    'Swish': paddle.nn.Swish,
-}
+quant_output_layers = (
+    paddle.nn.Conv2D, paddle.nn.Conv2DTranspose, paddle.nn.Linear,
+    paddle.nn.AdaptiveAvgPool2D, paddle.nn.AdaptiveMaxPool2D,
+    paddle.nn.AvgPool2D, paddle.nn.MaxPool2D, paddle.nn.BatchNorm,
+    paddle.nn.BatchNorm2D, paddle.nn.LayerNorm, paddle.nn.SyncBatchNorm,
+    paddle.nn.ELU, paddle.nn.GELU, paddle.nn.Hardshrink, paddle.nn.Hardsigmoid,
+    paddle.nn.Hardswish, paddle.nn.Hardtanh, paddle.nn.LeakyReLU,
+    paddle.nn.LogSigmoid, paddle.nn.LogSoftmax, paddle.nn.Maxout,
+    paddle.nn.PReLU, paddle.nn.ReLU, paddle.nn.ReLU6, paddle.nn.SELU,
+    paddle.nn.Sigmoid, paddle.nn.Softmax, paddle.nn.Softplus,
+    paddle.nn.Softshrink, paddle.nn.Softsign, paddle.nn.Swish, paddle.nn.Tanh,
+    paddle.nn.Tanhshrink, paddle.nn.ThresholdedReLU, paddle.nn.Upsample)
 
 weight_op_types = [
     "conv2d", "depthwise_conv2d", "matmul", "conv2d_transpose",