Autodiff support for im2col (unfold)

My understanding is that torch.nn.unfold does the equivalent of tf.extract_image_patches. The extract_image_patches in TF doesn’t support gradients and backprop to my knowledge.

I was wondering what the status of unfold in PyTorch is - does it support gradients / autodiff right now? Has anyone used this successfully as part of an architecture?

In terms of speedup, how does it compare to explicitly extracting local windows in a DWise Conv like fashion? Example, Scott Gray has these kernels in his blocksparse repo, where I believe the idea is to run the convolutions at a block level and simulate the channel wise index gathering through a convolution primitive (which I guess implicitly calls im2col or whatever’s faster under the hood)- blocksparse/ at master · openai/blocksparse · GitHub

Edit: More specifically, for the speedup part, if anyone has benchmarked implementing a regular conv2d with unfold + matmul + fold and compared the two on a GPU, that would be great to know!


PyTorch backprops through nn.unfold/fold.

The native aka THNN convolutions are implemented using im2col, so you can just disable CuDNN to benchmark whatever network you want. It uses a lot of memory and it won’t be fast. Part of that is the inefficiency in memory accesses when the im2col result is materialized in GPU memory (as opposed to living only in shared memory).

Best regards