How can I compute regular deep learning 2D convolution layer?

[BugQualia]

EDIT: What I was trying to implement was locally connected layer. The convolution works fine.

Yes, its a stupid question but I cannot make it work despite the fact that I put unreasonable amount of time reading docs.

It seems like that Convolutions in JAX document only gives example for applying single kernel to every part of an image, which is not the case in deep learning convolution layer.

jax.lax.conv_general_dilated (doc) requires both lhs and rhs to have dimension of n+2. for reular image, n is 2, and lhs and rhs have demension of N H(image) W(image) C and H(kernel) W(kernel) I O respectively.

In Convolutions in JAX, it describes the meaning of each alphabet as:
N - batch dimension
H - spatial height
W - spatial height
C - channel dimension
I - kernel input channel dimension
O - kernel output channel dimension

Below is a working example:

import jax
import jax.numpy as jnp

x = jnp.ones((1, 3, 100, 100), dtype='float32')  # NCHW 
w = jnp.ones((7, 3, 5, 5), dtype='float32')  # OIWH

out1 = jax.lax.conv(x, w, window_strides=(1, 1), padding=((2, 2), (2, 2))) # works same with padding='SAME'
print(out1.shape)  
# result: (1, 7, 100, 100)

However, my problem is that convolution in deep learning requires different kernels for different part of a image.
So, 4 dimensional weight (w = jnp.ones((7, 3, 5, 5), dtype='float32') # OIWH) is not enough. Instead, it should have dimension of 6, like W(output) H(output) I(image) W(channel) H(channel) Oor its variant. (in this case).

jax.lax.conv, jax.lax.conv_general_dilated all throws error message when x and w have different dimension: TypeError: convolution requires lhs and rhs ndim to be equal, got {} and {}.

What really boggles me is that I think I solved this problem before. I might am completely missing the point but I have spent way too much time in this. It will be great to get an answer.

[hawkinsp]

It's not normally the case in deep learning that different parts of an image get different convolution kernels. Instead, the same kernel is advanced across the input image, and that is the operation that lax.conv performs.

Can you point to an example of what you're trying to implement in another deep learning system? I don't think the operation you have described exists in any of the standard systems.

[BugQualia]

Maybe I didn't described it very well.

For simple 1-d example:

'Mathematical' convolution:
data = [1, 2, 3, 4, 5]
kernel = [100, 10, 1]
stride = 1
than the output would be:
output = [123, 234, 345]
It works like definition of convolution

'Convolution' that I want:
data = [1, 2, 3, 4, 5]
w = [[100, 10, 1], [1, 10, 100], [1, 100, 10]]
stride = 1
than the output would be:
output = [123, 432, 453]

I always thought that second case is what deep learning does.. Is it not?

[hawkinsp]

No, the convolution used in deep learning is usually the first case.

If you did want to implement something like your second case, you could extract image patches using lax.conv_general_dilated_patches and then perform a matrix multiplication. The former operation is sometimes called im2col in other systems. But caution: this operation is quite memory intensive.

[BugQualia]

Thanks for the answer. It really helped me..
I found out that what I was trying to implement is locally connected layer.

To add some context about why I made such a mistake: I have a neuroscience background and I'm making a brain simulation(similar to SNN) with JAX. And in the context of neuroscience, weight sharing is impossible. So I always assumed that convolutions would worked like locally connected layer, without thinking much about it. Most DL libraries are so successful in hiding the internal computation, that I've never had a problem while using tensorflow/pytorch.
Again, thanks for the answer.

[BugQualia]

It turns out that what I was trying to implement was locally connected layer.
Below is my implementation. Note that its not that generalized.
You can copy and paste below code and run it in colab.

import numpy as np
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

@jax.jit
def locally_connected_2d(x, w, window_strides=(1, 1)):
    patches = jax.lax.conv_general_dilated_patches(
        lhs=x.reshape(1, *x.shape), 
        filter_shape=(w.shape[3], w.shape[4]), 
        window_strides=window_strides, 
        padding='SAME', 
        dimension_numbers  = ('NHWC', 'OIHW', 'NHWC')
    )
    w = w.reshape(w.shape[0], w.shape[1], w.shape[2], -1)
    patches = patches.reshape(patches.shape[1], patches.shape[2], -1)
    return jax.lax.dot_general(w, patches, dimension_numbers=(([3], [2]), ([0, 1], [0, 1])))

x = jnp.arange(100*100*5, dtype='float32').reshape(100, 100, 5)
w = jnp.arange(100*100*3*7*7*5, dtype='float32').reshape(100, 100, 3, 7, 7, 5)

print("x.shape=", x.shape)
print("w.shape=", w.shape)

r = locally_connected_2d(x, w)
print("r.shape=", r.shape)
plt.imshow(np.array(r/np.max(r)*255, dtype='uint8'))
plt.show()

%timeit locally_connected_2d(x, w).block_until_ready()

# x.shape= (100, 100, 5)
# w.shape= (100, 100, 3, 7, 7, 5)
# r.shape= (100, 100, 3)
# 1000 loops, best of 5: 607 µs per loop