# Spatial transformer (i.e. GPU-friendly interpolations with gradients) with Flux.jl?

**URL:** <https://discourse.julialang.org/t/spatial-transformer-i-e-gpu-friendly-interpolations-with-gradients-with-flux-jl/70836>\
**Category:** Machine Learning\
**Tags:** flux, zygote, pytorch, interpolations\
**Created:** [November 3, 2021, 3:24am UTC](https://discourse.julialang.org/t/spatial-transformer-i-e-gpu-friendly-interpolations-with-gradients-with-flux-jl/70836 "2021-11-03T03:24:19Z")\
**Posts on this page:** 8\
**Page:** 1

<div class="post-metadata">

**Author:** ![afishy](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/afishy/32/18870_2.png) [@afishy](https://discourse.julialang.org/u/afishy)\
**Post date:** [November 3, 2021, 3:24am UTC](https://discourse.julialang.org/t/spatial-transformer-i-e-gpu-friendly-interpolations-with-gradients-with-flux-jl/70836/1 "2021-11-03T03:24:19Z")

</div>

Originally posted this on the flux-bridged Slack channel, but thought this might be a good place to post. I need a [spatial transformer](https://proceedings.neurips.cc//paper/2015/file/33ceb07bf4eeb3da587e268d663aba1a-Paper.pdf) or something equivalent for a project, written so far entirely in Julia. I thought Interpolations.jl might be a good solution, but it isn’t GPU-friendly so far. There’s a [sample implementation](https://github.com/thebhatman/Spatial-Transformer-Network/blob/master/src/stn.jl) but it’s old and not gpu-friendly and i’m not sure how to adapt it to be.

What I really need is a way to differentiably do this (from the paper)

 ![image](https://global.discourse-cdn.com/julialang/original/3X/2/d/2d8c884b3a55686bf31273f7d5c149a19762bbb4.jpeg)

The other options are a) moving everything to Pytorch (which has the [parts I’m missing baked in](https://pytorch.org/tutorials/intermediate/spatial_transformer_tutorial.html) – would rather not ☹ or b) getting a Pytorch model which somehow works with Zygote + the rest of my model upstream? Is there any hope?

Also I think it would be nice to provide an example for people (maybe in the Model Zoo or wherever) who might want a spatial transformer in Julia (or some of the functionality anyway) without having to do all this soul searching 🙂

Any help greatly appreciated!

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [November 3, 2021, 4:23am UTC](https://discourse.julialang.org/t/spatial-transformer-i-e-gpu-friendly-interpolations-with-gradients-with-flux-jl/70836/2 "2021-11-03T04:23:59Z")

</div>

Coincidentally, a [PR for grid sampling](https://github.com/FluxML/NNlib.jl/pull/361) was just opened against NNlib (CUDA version [here](https://github.com/FluxML/NNlibCUDA.jl/pull/31)). That leaves the affine grid generation, for which the code at [https://github.com/thebhatman/Spatial-Transformer-Network/blob/master/src/stn.jl#L43-L58](https://github.com/thebhatman/Spatial-Transformer-Network/blob/master/src/stn.jl#L43-L58) might be workable with some minor changes since it’s 100% vectorized already. Also have a look at [NNlib · Flux](https://fluxml.ai/Flux.jl/stable/models/nnlib/#Upsampling) to see if you can make use of anything there.

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [November 3, 2021, 4:27am UTC](https://discourse.julialang.org/t/spatial-transformer-i-e-gpu-friendly-interpolations-with-gradients-with-flux-jl/70836/3 "2021-11-03T04:27:16Z")

</div>

> [@afishy](#):
>
> Also I think it would be nice to provide an example for people (maybe in the Model Zoo or wherever) who might want a spatial transformer in Julia (or some of the functionality anyway) without having to do all this soul searching

As for this, the Julia ML community is orders of magnitude smaller than those around TF or PyTorch, so the likelihood of newer/more specialized models like spatial transformers being added to the zoo is pretty low (case in point, this is the first time I’ve heard of this architecture!) The best way to remedy that is for folks with the know-how to contribute PRs to the zoo, so those are very much welcome 🙂

---

<div class="post-metadata">

**Author:** ![afishy](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/afishy/32/18870_2.png) [@afishy](https://discourse.julialang.org/u/afishy)\
**Post date:** [November 3, 2021, 4:31am UTC](https://discourse.julialang.org/t/spatial-transformer-i-e-gpu-friendly-interpolations-with-gradients-with-flux-jl/70836/4 "2021-11-03T04:31:13Z")

</div>

Yes absolutely, happy to put this up if it works - I just need a few steps to get there!

---

<div class="post-metadata">

**Author:** ![afishy](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/afishy/32/18870_2.png) [@afishy](https://discourse.julialang.org/u/afishy)\
**Post date:** [November 3, 2021, 7:08am UTC](https://discourse.julialang.org/t/spatial-transformer-i-e-gpu-friendly-interpolations-with-gradients-with-flux-jl/70836/6 "2021-11-03T07:08:13Z")

</div>

The [PR](https://github.com/FluxML/NNlib.jl/pull/361) works! Note for anyone who wants to try: you need to merge the PR into the branches of NNlib.jl and NNlibCUDA.jl and add/precompile the packages that way (i think?) for the functionality to work, otherwise you may get a “this intrinsic must be compiled to be called” error when trying to take a gradient.

Now the only issue is the grid generator. This is my minor adaptation from [thebhatman/Spatial-Transformer-Network](https://github.com/thebhatman/Spatial-Transformer-Network/blob/master/src/stn.jl#L43-L58). It works so far, I’ll put the full implementation up soon!

```julia

x = LinRange(-1f0, 1f0, width) |> gpu
y = LinRange(-1f0, 1f0, height) |> gpu

function affine_grid_generator(x, y, theta)
    batch_size = size(theta)[3]
    one1 = one(eltype(theta))
    x = LinRange(-one1, one1, width)
    y = LinRange(-one1, one1, height)
    x_t_flat = reshape(repeat(x, height), 1, height * width)
    y_t_flat = reshape(repeat(transpose(y), width), 1, height * width)
    all_ones = ones(eltype(x_t_flat), 1, size(x_t_flat)[2])

    sampling_grid = vcat(x_t_flat, y_t_flat, all_ones)
    sampling_grid = reshape(
        transpose(repeat(transpose(sampling_grid), batch_size)),
        3,
        size(x_t_flat)[2],
        batch_size,
    )

    batch_grids = batched_mul(theta, sampling_grid)
    y_s = reshape(batch_grids[2, :, :], width, height, batch_size)
    x_s = reshape(batch_grids[1, :, :], width, height, batch_size)
    return x_s, y_s

end

```

---

<div class="post-metadata">

**Author:** ![dhairyagandhi96](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dhairyagandhi96/32/7589_2.png) [@dhairyagandhi96](https://discourse.julialang.org/u/dhairyagandhi96)\
**Post date:** [December 20, 2021, 12:06pm UTC](https://discourse.julialang.org/t/spatial-transformer-i-e-gpu-friendly-interpolations-with-gradients-with-flux-jl/70836/7 "2021-12-20T12:06:17Z")

</div>

Would love to see an updated spatial transformer in the zoo too!

---

<div class="post-metadata">

**Author:** ![afishy](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/afishy/32/18870_2.png) [@afishy](https://discourse.julialang.org/u/afishy)\
**Post date:** [December 21, 2021, 8:08am UTC](https://discourse.julialang.org/t/spatial-transformer-i-e-gpu-friendly-interpolations-with-gradients-with-flux-jl/70836/8 "2021-12-21T08:08:15Z")

</div>

Ooh forgot about this, I’ll try and do it over the break!

---

<div class="post-metadata">

**Author:** ![afishy](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/afishy/32/18870_2.png) [@afishy](https://discourse.julialang.org/u/afishy)\
**Post date:** [July 1, 2022, 4:56am UTC](https://discourse.julialang.org/t/spatial-transformer-i-e-gpu-friendly-interpolations-with-gradients-with-flux-jl/70836/9 "2022-07-01T04:56:17Z")

</div>

Pardon the delay, just made a PR for a spatial transformer example up on the model zoo, [https://github.com/FluxML/model-zoo/pull/361](https://github.com/FluxML/model-zoo/pull/361)