speedup sparse block (PaddlePaddle#1022)

* speedup_sparse_block

* speedup_sparse_block

* speedup_sparse_block

Co-authored-by: minghaoBD <[email protected]>

Author: RachelXu7

paddleslim/prune/unstructured_pruner_utils.py

@@ -8,11 +8,25 @@
 
 def cal_mxn_avg_matrix(mat, m=1, n=1):
     if m == 1 and n == 1: return copy.deepcopy(mat)
+
+    ori_row, ori_col = mat.shape[0], mat.shape[1]
+    if len(mat.shape) == 4:
+        assert mat.shape[2:] == (1, 1), "Only support for (n, n, 1, 1) for now."
+        mat = mat.reshape(ori_row, ori_col)
+
+    res_col = n - len(mat[0]) % n
+    res_row = m - len(mat) % m
+
+    mat = np.pad(mat, ((0, res_col), (0, res_col)), 'reflect')
     avg_mat = np.zeros_like(mat)
-    rows = len(mat) // m + 1
-    cols = len(mat[0]) // n + 1
-    for row in range(rows):
-        for col in range(cols):
-            avg_mat[m * row:m * row + m, n * col:n * col + n] = np.mean(mat[
-                m * row:m * row + m, n * col:n * col + n])
+    new_shape = [len(mat) // m, len(mat[0]) // n, m, n]
+    strides = mat.itemsize * np.array([len(mat) * m, n, len(mat), 1])
+    mat = np.lib.stride_tricks.as_strided(mat, shape=new_shape, strides=strides)
+    mat = mat.mean((2, 3), keepdims=True)
+    mat = np.tile(mat, (1, 1, m, n))
+    for i in range(len(mat)):
+        avg_mat[i * m:i * m + m] = np.concatenate(list(mat[i]), axis=1)
+    avg_mat = avg_mat[:ori_row, :ori_col]
+    if len(mat.shape) == 4:
+        avg_mat = avg_mat.reshape(ori_row, ori_col, 1, 1)
     return avg_mat