Merge pull request #885 from a-maci/2dmasking_sparsity

2d masking and sparsity

apex/contrib/sparsity/asp.py

@@ -6,7 +6,7 @@ torchvision_imported=True
 try:
     import torchvision
 except ImportError:
-    print("[ASP][Warning] torchvision cannot be imported, may infuence functionality of MaskRCNN/KeypointRCNN network from torchvision.")
+    print("[ASP][Warning] torchvision cannot be imported.")
     torchvision_imported=False
 
 def eligible_modules(model, whitelist_layer_types, allowed_layer_names, disallowed_layer_names):
@@ -78,7 +78,7 @@ class ASP:
         # function to extract variables that will be sparsified. 
         # idea is that you will add one of these functions for each module type that can be sparsified.
         if torchvision_imported:
-            print("[ASP] torchvision is imported, can work smoothly with the MaskRCNN/KeypointRCNN from torchvision.")
+            print("[ASP] torchvision is imported, can work with the MaskRCNN/KeypointRCNN from torchvision.")
             sparse_parameter_list = {torch.nn.Linear: ['weight'], torch.nn.Conv1d: ['weight'], torch.nn.Conv2d: ['weight'], torch.nn.Conv3d: ['weight'], torchvision.ops.misc.Conv2d: ['weight']}
         else:
             sparse_parameter_list = {torch.nn.Linear: ['weight'], torch.nn.Conv1d: ['weight'], torch.nn.Conv2d: ['weight'], torch.nn.Conv3d: ['weight']}

apex/contrib/sparsity/sparse_masklib.py

@@ -5,7 +5,7 @@ import collections
 from itertools import permutations
 
 
-""" compute density (helper fn to compute % NNZs in a tensor)"""
+""" compute density (helper fn to compute % NNZs in a tensor) """
 def fill(x):
     return float(x.nonzero().size(0))/torch.numel(x)
 
@@ -20,7 +20,7 @@ def reshape_1d(matrix, m):
     else:
         return matrix.view(-1,m), matrix.shape
 
-""" return all possible m:n patterns in a 1d vector. """
+""" return all possible m:n patterns in a 1d vector """
 valid_m4n2_1d_patterns = None
 def compute_valid_1d_patterns(m,n):
     # Early exit if patterns was already created.
@@ -49,8 +49,21 @@ def mn_1d_best(matrix, m, n):
 def m4n2_1d(mat, density):
     return mn_1d_best(mat, 4, 2)
 
-""" Comment: Following 2d masking related code (for training) can be removed or marked experimental (78 LOC) """
-""" m:n 2d structured greedy """
+"""
+  Below 2d-masking related code is targeted more for training (from scratch).
+  2d-pruning of a weight tensor is done to accelerate DGRAD step during backprop
+  phase of training algorithm. Acceleration comes from using SpMMA instructions in
+  Tensor Cores of NVIDIA Ampere GPU Architecture 
+  (note: this code does not do the acceleration, GPU kernels are required for this).
+  1d pruning of weight tensor helps speed up FPROP step by pruning in 2:4 pattern
+  along the horizontal (logical) direction.
+  During DGRAD step, weight tensor is transposed. 2d pruning functions below, mask
+  weight tensor such that their transposed versions are also 2:4 sparse along the
+  horizontal (logical) direction. Thus, with 2d pruning, weight tensors are 
+  2:4 sparse along row and column directions.
+ """
+
+""" m:n 2d structured pruning: greedy method to select mask """
 def mn_2d_greedy(matrix, m, n):
     # Convert to numpy
     mat = matrix.cpu().detach().numpy()
@@ -105,7 +118,7 @@ def compute_valid_2d_patterns(m,n):
     if m == 4  and n == 2: valid_m4n2_2d_patterns  = valid_patterns
     return valid_patterns
 
-""" m:n 2d structured best """
+""" m:n 2d structured pruning: exhaustive method to select best mask """
 def mn_2d_best(matrix, m, n):
     # Find all possible patterns.
     patterns = compute_valid_2d_patterns(m,n).cuda()
@@ -127,6 +140,7 @@ def mn_2d_best(matrix, m, n):
 def m4n2_2d_best(mat, density):
     return mn_2d_best(mat, 4, 2)
 
+
 """ returns a sparse mask """
 def create_mask(tensor, pattern="m4n2_1d", density=0.5):
     # Reshape tensor and mask.