fix(speedup): make the tensor contiguous before randomizing (#5141)

nni/compression/pytorch/speedup/compressor.py

@@ -197,7 +197,10 @@ class ModelSpeedup:
                 continue
             # The detach operation here is for the in-place operation. We cannot
             # directly can the backward on the output tensor of an in-place operator.
-            dummy_input.append(self.internal_result[_input].detach())
+            if isinstance(self.internal_result[_input], torch.Tensor):
+                dummy_input.append(self.internal_result[_input].detach())
+            else:
+                dummy_input.append(self.internal_result[_input])
 
             debugnames.append(_input)
 
@@ -369,7 +372,7 @@ class ModelSpeedup:
         for node in self.torch_graph.nodes_py.nodes_op:
             successors = self.torch_graph.find_successors(node.unique_name)
             out_degree[node.unique_name] = len(successors)
-            predecessors = self.torch_graph.find_predecessors(node.unique_name)
+            predecessors = set(self.torch_graph.find_predecessors(node.unique_name))
             in_degree[node.unique_name] = len(predecessors)
             if in_degree[node.unique_name] == 0:
                 visit_queue.put(node)
@@ -390,8 +393,8 @@ class ModelSpeedup:
         while not visit_queue.empty():
             curnode = visit_queue.get()
             self.update_indirect_sparsity(curnode)
-            predecessors = self.torch_graph.find_predecessors(
-                curnode.unique_name)
+            predecessors = set(self.torch_graph.find_predecessors(
+                curnode.unique_name))
             for predecessor in predecessors:
                 out_degree[predecessor] -= 1
                 if out_degree[predecessor] == 0:

nni/compression/pytorch/speedup/infer_mask.py

@@ -127,7 +127,7 @@ class AutoMaskInference:
         # when the confidence is low. In the future, we will add the mask inference
         # rules for ReLU6 to break this range constraint.
         with torch.no_grad():
-            for tensor in self.dummy_input:
+            for index, tensor in enumerate(self.dummy_input):
                 if isinstance(tensor, torch.Tensor) and len(tensor.size()) > self.batch_dim\
                     and tensor.size(self.batch_dim) == self.batch_size:
                     # if the input tensor only has one dimension, which means
@@ -135,6 +135,9 @@ class AutoMaskInference:
                     # this tensor, because our tensor scrambling is on the batch
                     # dimention. For example, if the tensor is a scalar(returned
                     # by the size operator), then we will skip this tensor
+                    if not tensor.is_contiguous():
+                        tensor = tensor.contiguous()
+                        self.dummy_input[index] = tensor
                     randomize_tensor(tensor, start, end)
             for para in self.weights:
                 randomize_tensor(self.weights[para].data, start, end)

test/algo/compression/v2/test_auto_conv.py

@@ -138,12 +138,21 @@ class TorchModel1(torch.nn.Module):
         x = torch.cat((y1, y2), dim=1)
 
         x = x.type_as(x)
-        x = x.expand_as(x)
+        dim = x.dim()
+        y = torch.sum(x, dim=dim-1, keepdim=True)
+        z = y.expand_as(x)
+        x = x / z
         x = torch.matmul(x, x.t())
-        x = torch.split(x, 1, dim=1)
-        x = torch.cat(x, dim=1)
+        
+        # TODO: should be available in #5107
+        # x = torch.split(x, 1, dim=1)
+        # x = torch.cat(x, dim=1)
+        
         # x = self.cond(x) # condition is not support now
-        x = self.asub(x)
+        # TODO: the asub execution is bad.
+        # reason: the graph_utils assumes modules with no sub_module are leaf_modules.
+        #         so the asub will be treated as a leaf_module.
+        # x = self.asub(x)
         x = torch.constant_pad_nd(x, (1,1,1,1), 3.14159)
 
         return x