How is regular 2-d convolution layer implemented in JAX(STAX)?

[jingtianweidizzf]

Dear all,

Hi, as I mentioned here, I'm searching for a function in JAX that does the same thing as torch.nn.functional.unfold does. (I was told to try conv_general_dilated_patches but I don't have an example of how to use it and there's always an error).

I implemented functions with for loops or 'jax.scan' to do this but its grad function runs really slow. (a forward operation only costs ~40ms but the backward gradient function costs 1.5s!)

So I'm now interested in how is regular 2-d convolutional layer implemented in stax? As I know, in pytorch they do im2col first and use GEMM to do this layerwise convolution. So I wonder if you also do similar things?

Thanks in advance!:)

Best,
Zifan

P.S. In case anyone is interested in why I need that: What I'm actually doing is to add some additional operations to the regular convolution (which can be seen as a sum of element-wise multiplication), and build a 'new convolution'. So my idea is to transform the image into large arrays, flatten the weights and then do my new operation based on the vectorized operations.

[jingtianweidizzf]

Problem already solved. Sorry for my stupid question again. And many thanks to @mattjj !!! :)

[BugQualia]

Hi! I'm looking for an answer for this exact same question. Do you mind if I ask how did you solved it?

[jingtianweidizzf]

Hi! I'm looking for an answer for this exact same question. Do you mind if I ask how did you solved it?

Hi, of course not, I was just using the conv_general_dilated_patches function to solve my problem:

What I was looking for is a function similar to the torch.nn.function.unfold which does the following thing:

import torch
x = torch.arange(32.0).reshape([1,2,4,4])
print(x)
x_unf = torch.nn.functional.unfold(x,kernel_size=3,dilation = 1, padding = 1).transpose(1,2)
print(x_unf.shape)
print(x_unf)

The output of the above block of code should be:

tensor([[[[ 0.,  1.,  2.],
          [ 3.,  4.,  5.],
          [ 6.,  7.,  8.]],

         [[ 9., 10., 11.],
          [12., 13., 14.],
          [15., 16., 17.]]]])
torch.Size([1, 9, 18])
tensor([[[ 0.,  0.,  0.,  0.,  0.,  1.,  0.,  3.,  4.,  0.,  0.,  0.,  0.,  9.,
          10.,  0., 12., 13.],
         [ 0.,  0.,  0.,  0.,  1.,  2.,  3.,  4.,  5.,  0.,  0.,  0.,  9., 10.,
          11., 12., 13., 14.],
         [ 0.,  0.,  0.,  1.,  2.,  0.,  4.,  5.,  0.,  0.,  0.,  0., 10., 11.,
           0., 13., 14.,  0.],
         [ 0.,  0.,  1.,  0.,  3.,  4.,  0.,  6.,  7.,  0.,  9., 10.,  0., 12.,
          13.,  0., 15., 16.],
         [ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12., 13.,
          14., 15., 16., 17.],
         [ 1.,  2.,  0.,  4.,  5.,  0.,  7.,  8.,  0., 10., 11.,  0., 13., 14.,
           0., 16., 17.,  0.],
         [ 0.,  3.,  4.,  0.,  6.,  7.,  0.,  0.,  0.,  0., 12., 13.,  0., 15.,
          16.,  0.,  0.,  0.],
         [ 3.,  4.,  5.,  6.,  7.,  8.,  0.,  0.,  0., 12., 13., 14., 15., 16.,
          17.,  0.,  0.,  0.],
         [ 4.,  5.,  0.,  7.,  8.,  0.,  0.,  0.,  0., 13., 14.,  0., 16., 17.,
           0.,  0.,  0.,  0.]]])

And this tensor x_unf containing the extracted patches of pixels in the input image is then used in my algorithm for further computation.

And the conv_general_dilated_patches works similarly in Google JAX, for example:

import jax
import jax.numpy as jnp

x2 = jnp.arange(18.0).reshape([1,2,3,3])
print(x2)

x2 = jnp.moveaxis(x2, 1,-1)
x2_unf = jax.lax.conv_general_dilated_patches(lhs = x2, filter_shape = (3,3), 
                        window_strides = (1,1), padding = 'SAME', 
                        dimension_numbers  = ('NHWC', 'HWIO', 'NHWC')) 

print(x2_unf.shape)
print(x2_unf)

And the output should be:

[[[[ 0.  1.  2.]
   [ 3.  4.  5.]
   [ 6.  7.  8.]]

  [[ 9. 10. 11.]
   [12. 13. 14.]
   [15. 16. 17.]]]]
(1, 3, 3, 18)
[[[[ 0.  0.  0.  0.  0.  1.  0.  3.  4.  0.  0.  0.  0.  9. 10.  0. 12.
    13.]
   [ 0.  0.  0.  0.  1.  2.  3.  4.  5.  0.  0.  0.  9. 10. 11. 12. 13.
    14.]
   [ 0.  0.  0.  1.  2.  0.  4.  5.  0.  0.  0.  0. 10. 11.  0. 13. 14.
     0.]]

  [[ 0.  0.  1.  0.  3.  4.  0.  6.  7.  0.  9. 10.  0. 12. 13.  0. 15.
    16.]
   [ 0.  1.  2.  3.  4.  5.  6.  7.  8.  9. 10. 11. 12. 13. 14. 15. 16.
    17.]
   [ 1.  2.  0.  4.  5.  0.  7.  8.  0. 10. 11.  0. 13. 14.  0. 16. 17.
     0.]]

  [[ 0.  3.  4.  0.  6.  7.  0.  0.  0.  0. 12. 13.  0. 15. 16.  0.  0.
     0.]
   [ 3.  4.  5.  6.  7.  8.  0.  0.  0. 12. 13. 14. 15. 16. 17.  0.  0.
     0.]
   [ 4.  5.  0.  7.  8.  0.  0.  0.  0. 13. 14.  0. 16. 17.  0.  0.  0.
     0.]]]]

So the output of the two functions are almost the same, only that the conv_general_dilated_patches function's output has two dimensions for different locations in the image and the torch.nn.functional.unfold outputs with only one dimension.

If you wanna recover the regular convolutional operation, just make GEMM of the x_unf tensor with the flattened convolutional kernel and it's done! You could also apply your custom tricks to the regular convolutional operation by using this function.

Hope my experience could help you :-)

Best,
Zifan