Differentiable Image Rotatation with CUDA

[roflmaostc]

Motivation and description

Hi!
I was working on DiffImageRotation.jl to rotate images differentiably with KernelAbstractions.jl.

So far only bilinear interpolation is supported but nearest neighbors would be straightforward too.

Should we try efforts to integrate this into this NNlib.jl?

Best,

Felix

CC: @maxfreu (he suggested it on discourse)

[maxfreu]

Disclaimer: I haven't worked on the NNlib codebase for quite a while, so the following assumes that there is no rotation code present as of now. Searching for "rotat" in the repo didn't yield any results, so...

Your implementation is nice and small, just one file. So the source could go directly into src, KA is already a dependency of NNlib. It only needs some tweaks to include a channel dimension, which should be easy to add. In NNlib, the convention is to have reverse pytorch dim ordering, so WHCN. It makes sense to treat the channel dimension differently, depending which device the kernel is run on. On the CPU one would treat the batch and channel dimension as one and parallelize over that. On the GPU, it's faster to parallelize spatially, so across width and height, looping over the channel dimension in the kernel. You can look at the code for bilinear upsampling, @pxl-th has done magnificent work on that here. Oh and we should think forward a bit to possibly include nearest neighbor rotation in the future. So either name the function imrotate_bilinear or include a switch method=:bilinear or so.

Regarding the tests: That should also be pretty straight-forward to include here, as you already have them, which is great! NNlib offers some tools to test your reverse adjoints against the forward diff results. The upsampling tests are an example how you can write a test that either runs on CPU or GPU and uses automatic tests against forwarddiff via gradtest_fn (which is just gradtest or gputest from test_utils.jl, depending on the device it's being run on).

Just fork NNlib, copy your files over, open a draft PR and we guide you through! :)

[roflmaostc]

So restricting to 4D array should be good?
I also implemented method=:nearest right now.

[maxfreu]

4D usually corresponds to a batch of 2D images with a color channel. That is the most common use case for rotations. But if you want to implement 5D arrays, so colored voxels and possibly rotations around arbitrary axes nobody will stop you :D So yes, I think 4D arrays are enough for now.

[maxfreu]

I think this can be closed, right?