Speedup sparse init (#6899)

mttk · fmassa · commit c96f2624a280 · 2018-05-03T14:29:12.000+01:00
* Sparse initialization speedup

* +empty line

* simplify indexing

* Can't reproduce locally...

* Can't reproduce locally...+

* Can't reproduce locally...+

* Fix test, cleanup
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -5227,7 +5227,7 @@ def test_sparse_default_std(self):
 
                 for col_idx in range(input_tensor.size(1)):
                     column = input_tensor[:, col_idx]
-                    assert column[column == 0].nelement() >= math.ceil(sparsity * cols)
+                    assert column[column == 0].nelement() >= math.ceil(sparsity * rows)
 
                 assert self._is_normal(input_tensor[input_tensor != 0], 0, std)
 
diff --git a/torch/nn/init.py b/torch/nn/init.py
@@ -360,17 +360,14 @@ def sparse_(tensor, sparsity, std=0.01):
         raise ValueError("Only tensors with 2 dimensions are supported")
 
     rows, cols = tensor.shape
-    num_zeros = int(math.ceil(rows * sparsity))
+    num_zeros = int(math.ceil(sparsity * rows))
 
     with torch.no_grad():
         tensor.normal_(0, std)
         for col_idx in range(cols):
-            row_indices = list(range(rows))
-            random.shuffle(row_indices)
+            row_indices = torch.randperm(rows)
             zero_indices = row_indices[:num_zeros]
-            for row_idx in zero_indices:
-                tensor[row_idx, col_idx] = 0
-
+            tensor[zero_indices, col_idx] = 0
     return tensor
 
 
