Im2col export (pytorch#30972)

Summary: Added im2col to opset 11. This symbolic is used to export torch.nn.Unfold Pull Request resolved: pytorch#30972 Reviewed By: hl475 Differential Revision: D18946921 Pulled By: houseroad fbshipit-source-id: 13dd0cbae899700df32fd74d6dff1f29033a2b4c

Author: neginraoof

test/onnx/test_pytorch_onnx_onnxruntime.py

@@ -2296,7 +2296,6 @@ def forward(self, x, pad):
  y = pad = (torch.tensor(2, dtype=torch.int64), torch.tensor(4, dtype=torch.int64))
  self.run_test(Pad(), (x, y))
 
-
  def test_reflection_pad(self):
  model = torch.nn.ReflectionPad1d(2)
  x = torch.randn(2, 4, 4)
@@ -2315,6 +2314,17 @@ def test_replication_pad(self):
  x = torch.randn(2, 2, 4, 4)
  self.run_test(model, x)
 
+ @skipIfUnsupportedMinOpsetVersion(11)
+ def test_im2col(self):
+ class Unfold(torch.nn.Module):
+ def forward(self, input):
+ return torch.nn.functional.unfold(input, kernel_size=(10, 15), dilation=2, padding=5, stride=3), \
+ torch.nn.functional.unfold(input, kernel_size=(2, 2), dilation=1, padding=0, stride=3), \
+ torch.nn.functional.unfold(input, kernel_size=(1, 1), dilation=5, padding=2, stride=3)
+
+ x = torch.rand(1, 1, 200, 100)
+ self.run_test(Unfold(), x)
+
  @skipIfNoLapack
  @skipIfUnsupportedMinOpsetVersion(11)
  def test_det(self):
@@ -2582,23 +2592,23 @@ def setup_rnn_tests():
  dict(TestONNXRuntime.__dict__, opset_version=11))
 
 
-# opset 9 tests, with keep_initializers_as_inputs=False for 
+# opset 9 tests, with keep_initializers_as_inputs=False for
 # IR version 4 style export.
 TestONNXRuntime_opset9_IRv4 = type(str("TestONNXRuntime_opset9_IRv4"),
  (unittest.TestCase,),
  dict(TestONNXRuntime.__dict__,
  keep_initializers_as_inputs=False))
 
 
-# opset 10 tests, with keep_initializers_as_inputs=False for 
+# opset 10 tests, with keep_initializers_as_inputs=False for
 # IR version 4 style export.
 TestONNXRuntime_opset10_IRv4 = type(str("TestONNXRuntime_opset10_IRv4"),
  (unittest.TestCase,),
  dict(TestONNXRuntime.__dict__, opset_version=10,
  keep_initializers_as_inputs=False))
 
 
-# opset 11 tests, with keep_initializers_as_inputs=False for 
+# opset 11 tests, with keep_initializers_as_inputs=False for
 # IR version 4 style export.
 TestONNXRuntime_opset11_IRv4 = type(str("TestONNXRuntime_opset11_IRv4"),
  (unittest.TestCase,),

torch/onnx/symbolic_opset11.py

@@ -5,6 +5,7 @@
 import torch
 import torch.onnx.symbolic_helper as sym_help
 import warnings
+import numpy
 
 from torch.onnx.symbolic_helper import parse_args, _unimplemented
 from torch.onnx.symbolic_opset9 import expand
@@ -514,6 +515,101 @@ def __lshift_(g, self, other):
  return lshift
 
 
+def _get_im2col_indices_along_dim(g, input_d, kernel_size_d, dilation_d, padding_d, stride_d):
+ # Input is always 4-D (N, C, H, W)
+ # Calculate indices of sliding blocks along spatial dimension
+ # Slide kernel over input each dim d:
+ # each dimension d ranges from 0 to input[d]+2×padding[d]−dilation[d]×(kernel_size[d]−1)
+ # with steps = stride
+
+ blocks_d = g.op("Add", input_d, g.op("Constant", value_t=torch.tensor(padding_d * 2)))
+ blocks_d = g.op("Sub", blocks_d, g.op("Constant", value_t=torch.tensor(dilation_d * (kernel_size_d - 1))))
+
+ # Stride kernel over input and find starting indices along dim d
+ blocks_d_indices = g.op("Range", g.op("Constant", value_t=torch.tensor(0)),
+ blocks_d, g.op("Constant", value_t=torch.tensor(stride_d)))
+
+ # Apply dilation on kernel and find its indices along dim d
+ kernel_grid = numpy.arange(0, kernel_size_d * dilation_d, dilation_d)
+ kernel_grid = g.op("Constant", value_t=torch.tensor([kernel_grid]))
+
+ # Broadcast and add kernel staring positions (indices) with
+ # kernel_grid along dim d, to get block indices along dim d
+ blocks_d_indices = g.op('Unsqueeze', blocks_d_indices, axes_i=[0]) # Reshape to [1, -1]
+ kernel_mask = g.op('Reshape', kernel_grid, g.op('Constant', value_t=torch.tensor([-1, 1])))
+ block_mask = g.op("Add", blocks_d_indices, kernel_mask)
+
+ return block_mask
+
+
+def _get_im2col_padded_input(g, input, padding_h, padding_w):
+ # Input is always 4-D tensor (N, C, H, W)
+ # Padding tensor has the following format: (padding_h, padding_w)
+ # Reshape the padding to follow ONNX format: (dim1_begin, dim2_begin,...,dim1_end, dim2_end,...)
+ pad = g.op("Constant", value_t=torch.LongTensor([0, 0, padding_h, padding_w] * 2))
+ return g.op("Pad", input, pad)
+
+
+def _get_im2col_output_shape(g, input, kernel_h, kernel_w):
+ batch_dim = size(g, input, g.op("Constant", value_t=torch.tensor(0)))
+ channel_dim = size(g, input, g.op("Constant", value_t=torch.tensor(1)))
+ channel_unfolded = g.op("Mul", channel_dim,
+ g.op("Constant", value_t=torch.tensor(kernel_h * kernel_w)))
+
+ return g.op("Concat",
+ g.op("Unsqueeze", batch_dim, axes_i=[0]),
+ g.op("Unsqueeze", channel_unfolded, axes_i=[0]),
+ g.op("Constant", value_t=torch.tensor([-1])), axis_i=0)
+
+
+@parse_args('v', 'is', 'is', 'is', 'is')
+def im2col(g, input, kernel_size, dilation, padding, stride):
+ # Input is always 4-D tensor (N, C, H, W)
+ # All other args are int[2]
+
+ input_h = size(g, input, g.op("Constant", value_t=torch.tensor(2)))
+ input_w = size(g, input, g.op("Constant", value_t=torch.tensor(3)))
+
+ stride_h, stride_w = stride[0], stride[1]
+ padding_h, padding_w = padding[0], padding[1]
+ dilation_h, dilation_w = dilation[0], dilation[1]
+ kernel_h, kernel_w = kernel_size[0], kernel_size[1]
+
+ blocks_row_indices = _get_im2col_indices_along_dim(g, input_h, kernel_h, dilation_h, padding_h, stride_h)
+ blocks_col_indices = _get_im2col_indices_along_dim(g, input_w, kernel_w, dilation_w, padding_w, stride_w)
+
+ output_shape = _get_im2col_output_shape(g, input, kernel_h, kernel_w)
+ padded_input = _get_im2col_padded_input(g, input, padding_h, padding_w)
+
+ # For a 4D matrix of size (1, 1, 3, 3) as below with kernel_size=2, stride=1, and dilation=1
+ # [[[[1., 2., 3.,],
+ # [4., 5., 6.,],
+ # [7., 8., 9.,]]]]
+ # First gather indices along rows (dim=2) with blocks_row_indices = [[0,1], [1,2]] to get:
+ # [[[[[1., 2., 3.],
+ # [4., 5., 6.]],
+ # [[4., 5., 6.],
+ # [7., 8., 9.]]]]]
+ # And then gather along cols (dim=4) with blocks_row_indices = [[0,1], [1,2]] to get:
+ # [[[[[[1., 2.],
+ # [4., 5.]],
+ # [[2., 3.],
+ # [5., 6]]],
+ # [[[4., 5.],
+ # [7., 8.]],
+ # [[5., 6.],
+ # [8., 9.]]]]]]
+ # Transpose dims 3 (depth) and 4 (rows), and then reshape to output shape (1, 1, 4, 4) to get:
+ # [[[1., 2., 4., 5.],
+ # [2., 3., 5., 6.],
+ # [4., 5., 7., 8.],
+ # [5., 6., 8., 9.]]]
+ output = g.op("Gather", padded_input, blocks_row_indices, axis_i=2)
+ output = g.op("Gather", output, blocks_col_indices, axis_i=4)
+ output = g.op("Transpose", output, perm_i=[0, 1, 2, 4, 3, 5])
+ return g.op("Reshape", output, output_shape)
+
+
 @parse_args('v', 'i', 'i')
 def flatten(g, input, start_dim, end_dim):
  dim = input.type().dim()