compatibility with Flux

[CarloLucibello]

Hi,
this is to enquire about the possibility of using Interpolations.jl to build image upsampling or downsampling layers in Flux.jl.
We recently added a bilinear upsampling function FluxML/NNlib.jl#262, but I was wondering if we could leverage instead some of the code here. The requisites would be

• handle batch dimension
• gradient computation
• compatible with CuArrays

I'm not familiar with the codebase here, maybe this is a long shot, but worth making an attempt and creating awareness about the possibility of this kind of interaction.
Moreover gpu and automatic diff friendly interpolations would generally benefit the ML ecosystem.

Best,
CL

cc @johnnychen94

[mkitti]

Would compatible with CuArrays create a dependency on CUDA?

[CarloLucibello]

Typically no, CuArrays just follow abstract array interface, although scalar indexing leads to horrible performance.
To address such cases a cuda kernel is needed (and the CUDA.jl dependence)

[maxfreu]

Just having completed the upsampling code: I would now rather use KernelAbstractions.jl for such work, as it saves you from writing almost the same code twice. Plus I think it's a much lighter dependency than CUDA. In the end I imagine having super portable (CPU, GPU (Nvidia & AMD!)) code with gradients, which can be recycled in julia math, images and nnlib. It could even be used to write specialized methods for different dimension orderings and dispatch via NamedDims.

[moesphere]

Related to the gradient computation, it would be helpful to have a rrule and frule function defined, package ChainRulesCore, to be able to use Interpolations with AD, e.g. Zygote. A gradient function is already defined. Thus the only part missing that needs to be implemented is a gradient with respect to the fields of Interpolations objects, if I understand it correctly, see here

[DhairyaLGandhi]

Could you point to the gradient function?

[rick2047]

I don't know too much about the code, but simple debugging tells me there are 6 gradient functions

These have special definitions

• Indexing.jl x2
• monotonic.jl

These call the general gradient functions

• Interpolations.jl
• extrapolation.jl
• scaling.jl

I suspect the indexing.jl one is the one which needs to be used here.

[DhairyaLGandhi]

Found it!

Interpolations.jl/src/Interpolations.jl

Line 413 in ec35013

function gradient(itp::AbstractInterpolation, x::Vararg{UnexpandedIndexTypes})

[rick2047]

Now that #414 is finished, how do we rewrite the upsampling functions? Going through the NNLib code, it seems like now that we have chainrules integration we can just rewrite stuff like upsample_nearest and its gradient. If its just that, I will open a PR their.

[mkitti]

Have you taken a look at http://juliamath.github.io/Interpolations.jl/latest/devdocs/ yet?

[rick2047]

I have read that page and I think I get it, a bit. But I was wondering if we should rather write them as simple functions using interpolation objects. Like the upsample_nearest is something like

itp = interpolate(x,(NoInterp(),NoInterp()))

We already have parent axes, which can be used to determine size.

[DhairyaLGandhi]

Would it work similarly to the kernels we have now in NNlib? Do we get the correct gradients out etc?

Worth doing a comparison imo

[rick2047]

I dont know what the kernel system is in NNLib (not familiar with that codebase at all). But I was thinking I can replace the code for upsample_nearest and its gradient function to the Zygote call. They have tests for these functions so rewriting should be easier.

[maxfreu]

Hi! @DhairyaLGandhi this serves as a reply to FluxML/NNlibCUDA.jl#5 (comment), but I think it might be better to keep the discussion in one place :)

Yes, I would find it cool if it works and we can get rid of duplicate code (of which I really wrote a lot, sorry for that). But I don't know if the performance can match the handwritten kernels I ported from pytorch (at least for now). For example I benchmarked the CPU implementation of bilinear upsampling against imresize! in ImageTransformations.jl (which uses interpolations.jl under the hood) and NNlib was 2.2x faster, even single threaded and with sub-optimal memory layout. Furthermore the ImageTransformations.jl code doesn't work on gpus when scalar getindex is disallowed. So I don't really know what the best way forward is. Maybe to enhance interpolations.jl and then use it? But how should all the specialized GPU code be handled, like the setup for the kernel calls (setting thread & block size)? Create CUDAinterpolations.jl so that interpolations.jl doesn't have to depend on CUDA? Or simply proceed with collecting specialized code for neural networks / dl in NNlib and just pick the stuff which really makes things faster and simpler from other packages? Note that at least for upsampling / interpolations things can be simplified a lot by smarter dispatch, which I can maybe hand in two PRs down the line or so.

@rick2047

I dont know what the kernel system is in NNLib

I don't know either 😄 I just ported pytorch code brute force and it turned out to be quite fast, but also a bit unwieldy, which can probably be improved by properly leveraging the dispatch. The nearest neighbour part was written by @mcabott.

Lastly, maybe someone else can also benchmark the interpolations to check my results?