Deformable convolutions in Flax/JAX

[bhchiang]

Is there support deformable convolutions (https://arxiv.org/abs/1703.06211)? If not, any pointers on how to implement it myself would be greatly appreciated. Thanks!

[marcvanzee]

They aren't implemented in Flax yet. They do seem to be implemented in Pytorch and Tensorflow, so if I were to implement them I would take that as a starting point and rewrite it to Flax.

Happy to help you if you have more specific questions!

[bhchiang]

Thanks! I ended up writing a naive 2D deformable convolution implementation in JAX / Flax (link) based on the TensorFlow implementation you linked. The high-level flow is to extract convolutions patches, perturb by the offset and do bilinear interpolation, and then do the final convolution with stride equal to the kernel size.

For the given sample data it ran in ~200 ms on my machine, probably ~an order of magnitude slower than a CUDA-backend implementation.

I'd be interested in continuing to work on this to speed it up / extend to N-d cases (hopefully good enough to merge into Flax) - any pointers would be greatly appreciated.

One thing I'm going to try is just to use conv_general_dilated_patches to extract patches instead of extracting patches with manual vmaps - not sure if this will help.