Extending Col2Im to support 5-D input tensors

vinayakdsci · June 24, 2024, 5:54pm

Hi all,

I recently had to work with PyTorch’s Col2Im and Im2Col functions, and I found that Col2Im currently supports only 3D and 4D input.

I think it would be great if Col2Im could be extended to support 5D input, and in case the maintainers agree with me, I would be really happy to go ahead and open an issue to track progress, and implement 5D support myself.

Any thoughts on this? Would it be alright if I open an issue?

Thanks!

soulitzer · June 25, 2024, 12:50pm

Do you mean to support an additional batch dimension?

vinayakdsci · June 25, 2024, 1:43pm

No, I mean that along with of 2-d output_size and 2-d kernel_size, we support 3-d output_size, and 3-d kernel_size, too.
I am not sure if that would be adding an additional batch dimension?

I apologise if what I am saying is not clear, if you need more clarification, I would be happy to provide an example, too. Thanks!

soulitzer · June 25, 2024, 3:50pm

Gotcha thanks for the clarification. Opening an issue sounds good.

Though apparently there is a library that already implements it here. Manual Implementation of Unrolled 3D Convolutions - #5 by Youngseok0001. I guess the discussion would be whether want to upstream the implementation there?

vinayakdsci · June 25, 2024, 4:33pm

That is a great way to go forward, but my thought was that as there is already an iterative implementation of Col2Im(that uses for loops), we could extend that itself to support 5D images.

This code in the Onnx Col2Im reference implementation is similar to what I had in mind:

github.com

onnx/onnx/blob/4c3ed5e08be75bbe1eeb6818e490b1b6a370183e/onnx/reference/ops/op_col2im.py#L145


      
                  raise ValueError(
                      f"Given n_input_plane={n_input_plane}, X.shape={X.shape}, "
                      f"output_shape={output_shape}, kernel_shape={kernel_shape}, "
                      f"dilations={dilations}, pads={pads}, strides={strides}, "
                      f"expected size of input's dimension 2 to match the calculated number of "
                      f"sliding blocks {n_blocks} = {block_size}, "
                      f"but got input.size(2)={input_length}.",
                  )
          
          
          def col2im_naive_implementation(data, image_shape, kernel_shape, dilations, pads, strides):  # type: ignore
              """Naive implementation for `col2im`."""
              n_dims = len(pads) // 2
              new_pads = np.array([(pads[i], pads[i + n_dims]) for i in range(n_dims)])
              _col2im_shape_check(data, image_shape, kernel_shape, dilations, new_pads, strides)
          
              data_col = data
              data_im = np.zeros(image_shape, dtype=data.dtype)
          
              dim_col = []
              for i in range(n_dims):

The implementation in Felix Dangel’s library(the one that you mentioned) would in my opinion be difficult to integrate in the existing implementation of Col2Im, because of its use of torch functions on whole tensors, instead of the use of indexing.

Your thoughts on this?

vinayakdsci · June 26, 2024, 11:15am

Hi @soulitzer, I went ahead and opened an issue, and cc’d you there. Here’s the link: Support for 5-D output image tensors in `Col2Im`. · Issue #129564 · pytorch/pytorch · GitHub.

Thanks!

vinayakdsci · June 28, 2024, 3:59am

Hi @soulitzer! Just a gentle ping. Could you have a look at the issue when you find time? (Linked above).

I’ll then go ahead and start the implementation once we are able to decide if we would want to upstream the existing implementation or I should extend the current one.

Thanks a lot!