remove ndim referneces to support older pytorch versions (#415)

jaybdub · web-flow · commit ea835dc36890 · 2020-09-14T19:35:57.000-07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,4 +4,5 @@
 
 ### Added 
 
-- Added reduced precision documentation page
+- Replaced Tensor.ndim references with len(tensor.shape) to support older pytorch versions
+- Added reduced precision documentation page
diff --git a/torch2trt/converters/ReLU6.py b/torch2trt/converters/ReLU6.py
@@ -8,7 +8,7 @@ def convert_ReLU6(ctx):
     output = ctx.method_return
     
     input_a_trt, input_b_trt = add_missing_trt_tensors(ctx.network, [input, 6])
-    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], output.ndim - 1)
+    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], len(output.shape) - 1)
 
     layer = ctx.network.add_activation(
         input=input_a_trt, type=trt.ActivationType.RELU)
@@ -20,4 +20,4 @@ def convert_ReLU6(ctx):
     
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 4, 5)])
 def test_relu6_basic():
-    return torch.nn.ReLU6()
+    return torch.nn.ReLU6()
diff --git a/torch2trt/converters/add.py b/torch2trt/converters/add.py
@@ -11,7 +11,7 @@ def convert_add(ctx):
     input_b = ctx.method_args[1]
     output = ctx.method_return
     input_a_trt, input_b_trt = add_missing_trt_tensors(ctx.network, [input_a, input_b])
-    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], output.ndim - 1)
+    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], len(output.shape) - 1)
     layer = ctx.network.add_elementwise(input_a_trt, input_b_trt, trt.ElementWiseOperation.SUM)
     output._trt = layer.get_output(0)
     
@@ -106,4 +106,4 @@ def forward(self, x):
     
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 10, 10)])
 def test_add_constant_batch():
-    return AddConstantBatch()
+    return AddConstantBatch()
diff --git a/torch2trt/converters/cat.py b/torch2trt/converters/cat.py
@@ -9,7 +9,7 @@ def convert_cat(ctx):
 
     output = ctx.method_return
     trt_inputs = add_missing_trt_tensors(ctx.network, inputs)
-    trt_inputs = broadcast_trt_tensors(ctx.network, trt_inputs, output.ndim - 1)
+    trt_inputs = broadcast_trt_tensors(ctx.network, trt_inputs, len(output.shape) - 1)
 
     layer = ctx.network.add_concatenation(inputs=trt_inputs)
     layer.axis = dim - 1
@@ -25,4 +25,4 @@ def forward(self, *x):
 
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 4, 4), (1, 3, 4), (1, 17, 4)])
 def test_Cat_basic():
-    return Cat(1)
+    return Cat(1)
diff --git a/torch2trt/converters/compare.py b/torch2trt/converters/compare.py
@@ -6,7 +6,7 @@ def convert_elementwise(ctx, op):
     input_b = ctx.method_args[1]
     output = ctx.method_return
     input_a_trt, input_b_trt = add_missing_trt_tensors(ctx.network, [input_a, input_b])
-    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], output.ndim - 1)
+    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], len(output.shape) - 1)
     layer = ctx.network.add_elementwise(input_a_trt, input_b_trt, op)
     output._trt = layer.get_output(0)
 
diff --git a/torch2trt/converters/div.py b/torch2trt/converters/div.py
@@ -12,7 +12,7 @@ def convert_div(ctx):
     input_b = ctx.method_args[1]
     output = ctx.method_return
     input_a_trt, input_b_trt = add_missing_trt_tensors(ctx.network, [input_a, input_b])
-    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], output.ndim - 1)
+    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], len(output.shape) - 1)
     layer = ctx.network.add_elementwise(input_a_trt, input_b_trt, trt.ElementWiseOperation.DIV)
     output._trt = layer.get_output(0)
 
@@ -24,7 +24,7 @@ def convert_rdiv(ctx):
     input_b = ctx.method_args[0]
     output = ctx.method_return
     input_a_trt, input_b_trt = add_missing_trt_tensors(ctx.network, [input_a, input_b])
-    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], output.ndim - 1)
+    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], len(output.shape) - 1)
     layer = ctx.network.add_elementwise(input_a_trt, input_b_trt, trt.ElementWiseOperation.DIV)
     output._trt = layer.get_output(0)
     
@@ -120,4 +120,4 @@ def forward(self, x):
     
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 10, 10)])
 def test_div_constant_batch():
-    return DivConstantBatch()
+    return DivConstantBatch()
diff --git a/torch2trt/converters/getitem.py b/torch2trt/converters/getitem.py
@@ -31,7 +31,7 @@ def convert_tensor_getitem(ctx):
     
     # Step 1 - Replace ellipsis with expanded slices
     
-    num_ellipsis = input.ndim - num_slice_types(slices)
+    num_ellipsis = len(input.shape) - num_slice_types(slices)
     
     new_slices = []
     for s in slices:
@@ -152,4 +152,4 @@ def test_tensor_getitem_2d_append_2dim():
 
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 5, 4, 3)])
 def test_tensor_getitem_2d_weird_combo():
-    return LambdaModule(lambda x: x[:, 0:3:4, None, None, 1, ...])
+    return LambdaModule(lambda x: x[:, 0:3:4, None, None, 1, ...])
diff --git a/torch2trt/converters/instance_norm.py b/torch2trt/converters/instance_norm.py
@@ -58,7 +58,7 @@ def convert_instance_norm(ctx):
         
         eps_np = np.array([eps], dtype=np.float32)
         keep_dims = True
-        reduce_axes = torch_dim_to_trt_axes(tuple(range(2, input.ndim)))
+        reduce_axes = torch_dim_to_trt_axes(tuple(range(2, len(input.shape))))
         
         # compute mean over spatial
         mean_trt = ctx.network.add_reduce(input._trt, trt.ReduceOperation.AVG, reduce_axes, keep_dims).get_output(0)
@@ -147,4 +147,4 @@ def test_instance_norm_2d_dynamic_affine():
 
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 10, 3, 3, 3)])
 def test_instance_norm_3d_dynamic_affine():
-    return torch.nn.InstanceNorm3d(10, affine=True, track_running_stats=False)
+    return torch.nn.InstanceNorm3d(10, affine=True, track_running_stats=False)
diff --git a/torch2trt/converters/max.py b/torch2trt/converters/max.py
@@ -8,14 +8,14 @@ def __convert_max_elementwise(ctx):
     input_b = ctx.method_args[1]
     output = ctx.method_return
     input_a_trt, input_b_trt = add_missing_trt_tensors(ctx.network, [input_a, input_b])
-    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], output.ndim - 1)
+    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], len(output.shape) - 1)
     layer = ctx.network.add_elementwise(input_a_trt, input_b_trt, trt.ElementWiseOperation.MAX)
     output._trt = layer.get_output(0)
     
 
 def __convert_max_reduce(ctx):
     input = ctx.method_args[0]
-    dim = get_arg(ctx, 'dim', pos=1, default=tuple(range(1, input.ndim)))
+    dim = get_arg(ctx, 'dim', pos=1, default=tuple(range(1, len(input.shape))))
     keepdim = get_arg(ctx, 'keepdim', pos=2, default=False)
     input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
     output_val = ctx.method_return[0]
@@ -59,4 +59,4 @@ def forward(self, x, y):
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 3), (1,)]) # broadcast
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 3, 3), (1, 3, 3)]) # broadcast
 def test_max_elementwise():
-    return MaxElementwise()
+    return MaxElementwise()
diff --git a/torch2trt/converters/min.py b/torch2trt/converters/min.py
@@ -8,14 +8,14 @@ def __convert_min_elementwise(ctx):
     input_b = ctx.method_args[1]
     output = ctx.method_return
     input_a_trt, input_b_trt = add_missing_trt_tensors(ctx.network, [input_a, input_b])
-    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], output.ndim - 1)
+    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], len(output.shape) - 1)
     layer = ctx.network.add_elementwise(input_a_trt, input_b_trt, trt.ElementWiseOperation.MIN)
     output._trt = layer.get_output(0)
     
 
 def __convert_min_reduce(ctx):
     input = ctx.method_args[0]
-    dim = get_arg(ctx, 'dim', pos=1, default=tuple(range(1,input.ndim)))
+    dim = get_arg(ctx, 'dim', pos=1, default=tuple(range(1,len(input.shape))))
     keepdim = get_arg(ctx, 'keepdim', pos=2, default=False)
     input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
     output_val = ctx.method_return[0]
@@ -59,4 +59,4 @@ def forward(self, x, y):
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 3), (1,)]) # broadcast
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 3, 3), (1, 3, 3)]) # broadcast
 def test_min_elementwise():
-    return MinElementwise()
+    return MinElementwise()
diff --git a/torch2trt/converters/mul.py b/torch2trt/converters/mul.py
@@ -11,7 +11,7 @@ def convert_mul(ctx):
     input_b = ctx.method_args[1]
     output = ctx.method_return
     input_a_trt, input_b_trt = add_missing_trt_tensors(ctx.network, [input_a, input_b])
-    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], output.ndim - 1)
+    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], len(output.shape) - 1)
     layer = ctx.network.add_elementwise(input_a_trt, input_b_trt, trt.ElementWiseOperation.PROD)
     output._trt = layer.get_output(0)
 
@@ -105,4 +105,4 @@ def forward(self, x):
     
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 10, 10)])
 def test_mul_constant_batch():
-    return MulConstantBatch()
+    return MulConstantBatch()
diff --git a/torch2trt/converters/pow.py b/torch2trt/converters/pow.py
@@ -10,7 +10,7 @@ def convert_pow(ctx):
     input_b = ctx.method_args[1]
     output = ctx.method_return
     input_a_trt, input_b_trt = add_missing_trt_tensors(ctx.network, [input_a, input_b])
-    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], output.ndim - 1)
+    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], len(output.shape) - 1)
     layer = ctx.network.add_elementwise(input_a_trt, input_b_trt, trt.ElementWiseOperation.POW)
     output._trt = layer.get_output(0)
 
@@ -21,7 +21,7 @@ def convert_pow(ctx):
     input_b = ctx.method_args[0]  # flipped for rpow
     output = ctx.method_return
     input_a_trt, input_b_trt = add_missing_trt_tensors(ctx.network, [input_a, input_b])
-    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], output.ndim - 1)
+    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], len(output.shape) - 1)
     layer = ctx.network.add_elementwise(input_a_trt, input_b_trt, trt.ElementWiseOperation.POW)
     output._trt = layer.get_output(0)
     
@@ -89,4 +89,4 @@ def forward(self, x):
 
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 224, 224)])
 def test_rpow_float():
-    return RpowFloat()
+    return RpowFloat()
diff --git a/torch2trt/converters/prod.py b/torch2trt/converters/prod.py
@@ -7,7 +7,7 @@
 @tensorrt_converter('torch.Tensor.prod')
 def convert_prod(ctx):
     input = ctx.method_args[0]
-    dim = get_arg(ctx, 'dim', pos=1, default=tuple(range(1, input.ndim)))
+    dim = get_arg(ctx, 'dim', pos=1, default=tuple(range(1, len(input.shape))))
     keepdim = get_arg(ctx, 'keepdim', pos=2, default=False)
     input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
     output = ctx.method_return
@@ -35,4 +35,4 @@ def test_prod_reduce_dim22():
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3)])
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 3)])
 def test_prod_reduce_dim1_keepdim():
-    return UnaryModule(lambda x: torch.prod(x, 1, keepdim=True))
+    return UnaryModule(lambda x: torch.prod(x, 1, keepdim=True))
diff --git a/torch2trt/converters/softmax.py b/torch2trt/converters/softmax.py
@@ -18,7 +18,7 @@ def convert_softmax(ctx):
 #     import pdb
 #     pdb.set_trace()
     if dim < 0:
-        dim = input.ndim + dim
+        dim = len(input.shape) + dim
 
     axes = 1 << (dim - 1)
 
@@ -47,4 +47,4 @@ def test_softmax_module_neg1():
 
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 3, 3)])
 def test_softmax_module_dim_neg2():
-    return torch.nn.Softmax(-2)
+    return torch.nn.Softmax(-2)
diff --git a/torch2trt/converters/sub.py b/torch2trt/converters/sub.py
@@ -10,7 +10,7 @@ def convert_sub(ctx):
     input_b = ctx.method_args[1]
     output = ctx.method_return
     input_a_trt, input_b_trt = add_missing_trt_tensors(ctx.network, [input_a, input_b])
-    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], output.ndim - 1)
+    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], len(output.shape) - 1)
     layer = ctx.network.add_elementwise(input_a_trt, input_b_trt, trt.ElementWiseOperation.SUB)
     output._trt = layer.get_output(0)
 
@@ -21,7 +21,7 @@ def convert_sub(ctx):
     input_b = ctx.method_args[0]  # flipped for rsub
     output = ctx.method_return
     input_a_trt, input_b_trt = add_missing_trt_tensors(ctx.network, [input_a, input_b])
-    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], output.ndim - 1)
+    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], len(output.shape) - 1)
     layer = ctx.network.add_elementwise(input_a_trt, input_b_trt, trt.ElementWiseOperation.SUB)
     output._trt = layer.get_output(0)
     
@@ -115,4 +115,4 @@ def forward(self, x):
     
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 10, 10)])
 def test_sub_constant_batch():
-    return SubConstantBatch()
+    return SubConstantBatch()
diff --git a/torch2trt/converters/sum.py b/torch2trt/converters/sum.py
@@ -7,7 +7,7 @@
 @tensorrt_converter('torch.Tensor.sum')
 def convert_sum(ctx):
     input = ctx.method_args[0]
-    dim = get_arg(ctx, 'dim', pos=1, default=tuple(range(1, input.ndim)))
+    dim = get_arg(ctx, 'dim', pos=1, default=tuple(range(1, len(input.shape))))
     keepdim = get_arg(ctx, 'keepdim', pos=2, default=False)
     input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
     output = ctx.method_return
@@ -49,4 +49,4 @@ def forward(self, x):
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 10, 3, 3)])
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 10, 23, 23)])
 def test_disparity_reg():
-    return DisparityRegression(10)    
+    return DisparityRegression(10)    
diff --git a/torch2trt/torch2trt.py b/torch2trt/torch2trt.py
@@ -140,7 +140,7 @@ def add_missing_trt_tensors(network, tensors):
             
             # remove all preceding ones, these can be re-inserted later when broadcasting
             num_preceding_ones = 0
-            for j in range(t.ndim):
+            for j in range(len(t.shape)):
                 if int(t.shape[j]) == 1:
                     num_preceding_ones += 1
                 else:
