[ASP] fix some bugs of asp (#43853)

python/paddle/fluid/contrib/sparsity/asp.py

@@ -722,6 +722,10 @@ class ASPHelper(object):
         if param_name in supported_layers_and_prune_func_map:
             return True
 
+        # The parameter's name is neither in *.* format nor added to supported_layers_and_prune_func_map, return False.
+        if len(param_name_list) == 1:
+            return False
+
         param_name_no_weight_suffix = param_name_list[0]
         param_type_suffix = param_name_list[1]
         layer_name = param_name_no_weight_suffix[:param_name_no_weight_suffix.

python/paddle/fluid/contrib/sparsity/supported_layer_list.py

@@ -15,14 +15,38 @@
 
 import numpy as np
 import paddle
+import copy
 from paddle.fluid.contrib import sparsity
 import threading
+import logging
+from ...log_helper import get_logger
 
 __all__ = ['add_supported_layer']
 
+_logger = get_logger(__name__,
+                     logging.INFO,
+                     fmt='%(asctime)s-%(levelname)s: %(message)s')
+
 
 def _default_pruning(weight_nparray, m, n, func_name, param_name):
 
+    # if the to-be-pruned dimension's size is smaller than m, we don't prune it. This strong assertion is required by the inference from cuSparseLT.
+    shape = weight_nparray.shape
+    weight_pruned_nparray = copy.deepcopy(weight_nparray)
+    weight_sparse_mask = np.ones_like(weight_pruned_nparray)
+    exlude_cond_shape2 = len(shape) == 2 and shape[0] < m
+    exlude_cond_shape4 = len(shape) == 4 and shape[1] < m
+    if exlude_cond_shape2:
+        _logger.warning(
+            '{} is not pruned because the first dimension of {} is smaller than {}'
+            .format(param_name, shape, m))
+        return weight_pruned_nparray, weight_sparse_mask
+    if exlude_cond_shape4:
+        _logger.warning(
+            '{} is not pruned because the second dimension of {} is smaller than {}'
+            .format(param_name, shape, m))
+        return weight_pruned_nparray, weight_sparse_mask
+
     checked_func_name = sparsity.CheckMethod.get_checking_method(func_name)
 
     # The double transpose ops here make sure pruning direction consistent with cuSparseLt.

python/paddle/fluid/tests/unittests/asp/test_asp_customized_pruning.py

@@ -223,6 +223,14 @@ class TestASPStaticCustomerizedPruneFunc(unittest.TestCase):
                 if (self.customer_prefix in param.name):
                     self.assertLessEqual(
                         np.sum(mat.flatten() - static_tensor.flatten()), 1e-4)
+                else:
+                    if (len(param.shape) == 4
+                            and param.shape[1] < 4) or (len(param.shape) == 2
+                                                        and param.shape[0] < 4):
+                        self.assertFalse(
+                            paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
+                                                                         n=2,
+                                                                         m=4))
                     else:
                         self.assertTrue(
                             sparsity.check_sparsity(
@@ -263,6 +271,14 @@ class TestASPStaticCustomerizedPruneFunc(unittest.TestCase):
                     self.assertLessEqual(
                         np.sum(mat_mask.flatten() -
                                static_tensor_mask.flatten()), 1e-4)
+                else:
+                    if (len(param.shape) == 4
+                            and param.shape[1] < 4) or (len(param.shape) == 2
+                                                        and param.shape[0] < 4):
+                        self.assertFalse(
+                            sparsity.check_sparsity(mat.T, n=2, m=4))
+                        self.assertFalse(
+                            sparsity.check_sparsity(mat_mask.T, n=2, m=4))
                     else:
                         self.assertTrue(
                             sparsity.check_sparsity(

python/paddle/fluid/tests/unittests/asp/test_asp_optimize_dynamic.py

@@ -131,6 +131,14 @@ class TestASPDynamicOptimize(unittest.TestCase):
             if ASPHelper._is_supported_layer(
                     paddle.static.default_main_program(), param.name):
                 mat = param.numpy()
+                if (len(param.shape) == 4
+                        and param.shape[1] < 4) or (len(param.shape) == 2
+                                                    and param.shape[0] < 4):
+                    self.assertFalse(
+                        paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
+                                                                     n=2,
+                                                                     m=4))
+                else:
                     self.assertTrue(
                         paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
                                                                      n=2,
@@ -165,6 +173,15 @@ class TestASPDynamicOptimize(unittest.TestCase):
             if ASPHelper._is_supported_layer(
                     paddle.static.default_main_program(), param.name):
                 mat = param.numpy()
+                if (len(param.shape) == 4
+                        and param.shape[1] < 4) or (len(param.shape) == 2
+                                                    and param.shape[0] < 4):
+                    self.assertFalse(
+                        paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
+                                                                     n=2,
+                                                                     m=4))
+                else:
+
                     self.assertTrue(
                         paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
                                                                      n=2,

python/paddle/fluid/tests/unittests/asp/test_asp_optimize_static.py

@@ -143,6 +143,14 @@ class TestASPStaticOptimize(unittest.TestCase):
             if ASPHelper._is_supported_layer(self.main_program, param.name):
                 mat = np.array(fluid.global_scope().find_var(
                     param.name).get_tensor())
+                if (len(param.shape) == 4
+                        and param.shape[1] < 4) or (len(param.shape) == 2
+                                                    and param.shape[0] < 4):
+                    self.assertFalse(
+                        paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
+                                                                     n=2,
+                                                                     m=4))
+                else:
                     self.assertTrue(
                         paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
                                                                      n=2,
@@ -172,6 +180,14 @@ class TestASPStaticOptimize(unittest.TestCase):
                 if ASPHelper._is_supported_layer(self.main_program, param.name):
                     mat = np.array(fluid.global_scope().find_var(
                         param.name).get_tensor())
+                    if (len(param.shape) == 4
+                            and param.shape[1] < 4) or (len(param.shape) == 2
+                                                        and param.shape[0] < 4):
+                        self.assertFalse(
+                            paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
+                                                                         n=2,
+                                                                         m=4))
+                    else:
                         self.assertTrue(
                             paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
                                                                          n=2,

python/paddle/fluid/tests/unittests/asp/test_asp_pruning_dynamic.py

@@ -85,6 +85,14 @@ class TestASPDynamicPruningBase(unittest.TestCase):
             if ASPHelper._is_supported_layer(
                     paddle.static.default_main_program(), param.name):
                 mat = param.numpy()
+                if (len(param.shape) == 4
+                        and param.shape[1] < 4) or (len(param.shape) == 2
+                                                    and param.shape[0] < 4):
+                    self.assertFalse(
+                        paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
+                                                                     n=2,
+                                                                     m=4))
+                else:
                     self.assertTrue(
                         paddle.fluid.contrib.sparsity.check_sparsity(
                             mat.T, func_name=self.mask_check_func, n=2, m=4))

python/paddle/fluid/tests/unittests/asp/test_asp_pruning_static.py

@@ -43,7 +43,8 @@ class TestASPStaticPruningBase(unittest.TestCase):
                                          padding=2,
                                          act="relu")
             hidden = fluid.layers.fc(input=hidden, size=32, act='softmax')
-            prediction = fluid.layers.fc(input=hidden, size=10, act='softmax')
+            hidden = fluid.layers.fc(input=hidden, size=3, act='softmax')
+            prediction = fluid.layers.fc(input=hidden, size=3, act='softmax')
             return img, label, prediction
 
         with fluid.program_guard(self.main_program, self.startup_program):
@@ -88,6 +89,14 @@ class TestASPStaticPruningBase(unittest.TestCase):
             if ASPHelper._is_supported_layer(self.main_program, param.name):
                 mat = np.array(fluid.global_scope().find_var(
                     param.name).get_tensor())
+                if (len(param.shape) == 4
+                        and param.shape[1] < 4) or (len(param.shape) == 2
+                                                    and param.shape[0] < 4):
+                    self.assertFalse(
+                        paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
+                                                                     n=2,
+                                                                     m=4))
+                else:
                     self.assertTrue(
                         paddle.fluid.contrib.sparsity.check_sparsity(
                             mat.T, func_name=self.mask_check_func, n=2, m=4))

python/paddle/fluid/tests/unittests/asp/test_asp_save_load.py

@@ -103,6 +103,14 @@ class TestASPDynamicOptimize(unittest.TestCase):
             if ASPHelper._is_supported_layer(
                     paddle.static.default_main_program(), param.name):
                 mat = param.numpy()
+                if (len(param.shape) == 4
+                        and param.shape[1] < 4) or (len(param.shape) == 2
+                                                    and param.shape[0] < 4):
+                    self.assertFalse(
+                        paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
+                                                                     n=2,
+                                                                     m=4))
+                else:
                     self.assertTrue(
                         paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
                                                                      n=2,
@@ -171,6 +179,14 @@ class TestASPStaticOptimize(unittest.TestCase):
             if ASPHelper._is_supported_layer(prog, param.name):
                 mat = np.array(fluid.global_scope().find_var(
                     param.name).get_tensor())
+                if (len(param.shape) == 4
+                        and param.shape[1] < 4) or (len(param.shape) == 2
+                                                    and param.shape[0] < 4):
+                    self.assertFalse(
+                        paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
+                                                                     n=2,
+                                                                     m=4))
+                else:
                     self.assertTrue(
                         paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
                                                                      n=2,

python/paddle/fluid/tests/unittests/asp/test_fleet_with_asp_dynamic.py

@@ -90,6 +90,14 @@ class TestFleetWithASPDynamic(unittest.TestCase):
             if ASPHelper._is_supported_layer(
                     paddle.static.default_main_program(), param.name):
                 mat = param.numpy()
+                if (len(param.shape) == 4
+                        and param.shape[1] < 4) or (len(param.shape) == 2
+                                                    and param.shape[0] < 4):
+                    self.assertFalse(
+                        paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
+                                                                     n=2,
+                                                                     m=4))
+                else:
                     self.assertTrue(
                         paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
                                                                      n=2,
@@ -146,6 +154,14 @@ class TestFleetWithASPAMPDynamic(unittest.TestCase):
             if ASPHelper._is_supported_layer(
                     paddle.static.default_main_program(), param.name):
                 mat = param.numpy()
+                if (len(param.shape) == 4
+                        and param.shape[1] < 4) or (len(param.shape) == 2
+                                                    and param.shape[0] < 4):
+                    self.assertFalse(
+                        paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
+                                                                     n=2,
+                                                                     m=4))
+                else:
                     self.assertTrue(
                         paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
                                                                      n=2,

python/paddle/fluid/tests/unittests/asp/test_fleet_with_asp_sharding.py

@@ -110,6 +110,14 @@ class TestFleetWithASPSharding(unittest.TestCase):
             if ASPHelper._is_supported_layer(train_prog, param.name):
                 mat = np.array(fluid.global_scope().find_var(
                     param.name).get_tensor())
+                if (len(param.shape) == 4
+                        and param.shape[1] < 4) or (len(param.shape) == 2
+                                                    and param.shape[0] < 4):
+                    self.assertFalse(
+                        paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
+                                                                     n=2,
+                                                                     m=4))
+                else:
                     self.assertTrue(
                         paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
                                                                      n=2,

python/paddle/fluid/tests/unittests/asp/test_fleet_with_asp_static.py

@@ -85,6 +85,14 @@ class TestFleetWithASPStatic(unittest.TestCase):
             if ASPHelper._is_supported_layer(train_prog, param.name):
                 mat = np.array(fluid.global_scope().find_var(
                     param.name).get_tensor())
+                if (len(param.shape) == 4
+                        and param.shape[1] < 4) or (len(param.shape) == 2
+                                                    and param.shape[0] < 4):
+                    self.assertFalse(
+                        paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
+                                                                     n=2,
+                                                                     m=4))
+                else:
                     self.assertTrue(
                         paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
                                                                      n=2,
@@ -146,6 +154,14 @@ class TestFleetWithASPAMPStatic(unittest.TestCase):
             if ASPHelper._is_supported_layer(train_prog, param.name):
                 mat = np.array(fluid.global_scope().find_var(
                     param.name).get_tensor())
+                if (len(param.shape) == 4
+                        and param.shape[1] < 4) or (len(param.shape) == 2
+                                                    and param.shape[0] < 4):
+                    self.assertFalse(
+                        paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
+                                                                     n=2,
+                                                                     m=4))
+                else:
                     self.assertTrue(
                         paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
                                                                      n=2,
@@ -187,6 +203,14 @@ class TestFleetWithASPAMPStatic(unittest.TestCase):
             if ASPHelper._is_supported_layer(train_prog, param.name):
                 mat = np.array(fluid.global_scope().find_var(
                     param.name).get_tensor())
+                if (len(param.shape) == 4
+                        and param.shape[1] < 4) or (len(param.shape) == 2
+                                                    and param.shape[0] < 4):
+                    self.assertFalse(
+                        paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
+                                                                     n=2,
+                                                                     m=4))
+                else:
                     self.assertTrue(
                         paddle.fluid.contrib.sparsity.check_sparsity(mat.T,
                                                                      n=2,

python/paddle/static/sparsity/__init__.py

@@ -17,6 +17,7 @@ from ...fluid.contrib.sparsity import calculate_density  #noqa: F401
 from ...fluid.contrib.sparsity import decorate  #noqa: F401
 from ...fluid.contrib.sparsity import prune_model  #noqa: F401
 from ...fluid.contrib.sparsity import reset_excluded_layers  #noqa: F401
+from ...fluid.contrib.sparsity import add_supported_layer  #noqa: F401
 from ...fluid.contrib import sparsity  #noqa: F401
 
 
@@ -27,5 +28,5 @@ def set_excluded_layers(main_program, param_names):
 
 __all__ = [  #noqa
     'calculate_density', 'decorate', 'prune_model', 'set_excluded_layers',
-    'reset_excluded_layers'
+    'reset_excluded_layers', 'add_supported_layer'
 ]