# Setting Flux custom gradients

**URL:** <https://discourse.julialang.org/t/setting-flux-custom-gradients/18728>\
**Category:** Machine Learning\
**Created:** [December 17, 2018, 12:59am UTC](https://discourse.julialang.org/t/setting-flux-custom-gradients/18728 "2018-12-17T00:59:12Z")\
**Posts on this page:** 7\
**Page:** 1

<div class="post-metadata">

**Author:** ![platawiec](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/platawiec/32/31914_2.png) [@platawiec](https://discourse.julialang.org/u/platawiec)\
**Post date:** [December 17, 2018, 12:59am UTC](https://discourse.julialang.org/t/setting-flux-custom-gradients/18728/1 "2018-12-17T00:59:12Z")

</div>

I have been trying to make a Spatial Transformer Network in Flux ([https://arxiv.org/pdf/1506.02025.pdf](https://arxiv.org/pdf/1506.02025.pdf)) mostly for my own education and because I think it would be slick to make one in Julia via Flux + Interpolations.

I am having trouble defining the custom gradient, however. I’ve posted an example below, based off of MNIST example in the model zoo. The forward pass goes fine, but when I do `train!` it gives me a `MethodError: objects of type Float64 are not callable`.

After inspecting, it seems to me that the custom gradient I defined for `FilledExtrapolation` objects is not being called, and it is instead trying to do the gradient itself, rather than by the Interpolations version.

What can I do to figure out more about what is being called in the backward pass? If I define a custom gradient, will it automatically work with broadcasting, as I’ve done here? What else am I missing?

Minimal excerpt:

```julia
using Flux
using Flux.Tracker: TrackedArray, TrackedReal, track, @grad, data
import Interpolations; const ITP = Interpolations;
# Fix overly-restrictive call signature in Interpolations (will submit PR)
@inline ITP.lispyprod(p, v::AbstractVector{T}, rest...) where T = ITP.lispyprod(p*zero(T), rest...)

(itp::ITP.FilledExtrapolation{<:Any, 4, <:ITP.AbstractInterpolation, <:Any, <:Any})(x::TrackedReal, y::TrackedReal, c, n) = track(itp, x, y, c, n)
@grad function (itp::ITP.FilledExtrapolation{<:Any, 4, <:ITP.AbstractInterpolation, <:Any, <:Any})(x::TrackedReal, y::TrackedReal, c::Int, n::Int)
  g = ITP.gradient(itp, x, y, c, n)
  return (itp(data(x), data(y), data(c), data(n)), D -> (D*g[1], D*g[2], nothing, nothing))
end

z = rand(5, 5, 1, 5)
itp = ITP.extrapolate(ITP.interpolate(z,
           (QuadInterp, QuadInterp, ITP.NoInterp())), zero(eltype(z)))

y, back = Flux.Tracker.forward(itp, 1, 2, 1, 1) # this gives an error within Interpolations, suggesting custom gradient wasn't called

```

Full example:

```julia
using Flux, Flux.Data.MNIST, Statistics
using Flux: onehotbatch, onecold, crossentropy, throttle
using Flux.Tracker: TrackedArray, TrackedReal, track, @grad, data
using Base.Iterators: repeated, partition
import Interpolations; const ITP = Interpolations;
# using CuArrays

# Classify MNIST digits with a Spatial Transformer Network
# "Spatial Transformer Networks" M Jaderberg, K Simonyan, A Zisserman, K Kavukcuoglu
# ArXiv https://arxiv.org/abs/1506.02025

imgs = MNIST.images()

labels = onehotbatch(MNIST.labels(), 0:9)

# Partition into batches of size 1,000
train = [(cat(float.(imgs[i])..., dims = 4), labels[:,i])
         for i in partition(1:60_000, 1000)]

train = gpu.(train)

# Prepare test set (first 1,000 images)
tX = cat(float.(MNIST.images(:test)[1:1000])..., dims = 4) |> gpu
tY = onehotbatch(MNIST.labels(:test)[1:1000], 0:9) |> gpu

# Fix overly-restrictive call signature in Interpolations (will submit PR)
@inline ITP.lispyprod(p, v::AbstractVector{T}, rest...) where T = ITP.lispyprod(p*zero(T), rest...)

# Define Interpolations gradients for Flux
const QuadInterp = ITP.BSpline(ITP.Quadratic(ITP.Line(ITP.OnCell())))
(itp::ITP.FilledExtrapolation{<:Any, 4, <:ITP.AbstractInterpolation, <:Any, <:Any})(x::TrackedReal, y::TrackedReal, c::TrackedReal, n::TrackedReal) = track(itp, x, y, c, n)
@grad function (itp::ITP.FilledExtrapolation{<:Any, 4, <:ITP.AbstractInterpolation, <:Any, <:Any})(x::TrackedReal, y::TrackedReal, c::Int, n::Int)
  g = ITP.gradient(itp, x, y, c, n)
  return (itp(data(x), data(y), data(c), data(n)), D -> (D*g[1], D*g[2], nothing, nothing))
end

mutable struct STN{NNT, GT}
  localizer::NNT
  grid::GT
end
function (m::STN)(x)
  img_transform = m.localizer(x)
  interp_grid_x = reshape(m.grid * img_transform[:, 1, :], size(x))
  interp_grid_y = reshape(m.grid * img_transform[:, 2, :], size(x))

  itp = ITP.extrapolate(ITP.interpolate(x,
    (QuadInterp, QuadInterp, ITP.NoInterp(), ITP.NoInterp())), zero(eltype(x)))

  itp.(interp_grid_x, interp_grid_y, 1, reshape(1:size(x, 4), 1, 1, 1, size(x, 4)))
end
Flux.@treelike STN

# this is the "localizer", which learns the affine transform to perform
localizer = Chain(x -> reshape(x, :, size(x, 4)),
                  Dense(28^2, 32, relu),
                  Dense(32, 6),
                  x -> reshape(x, 3, 2, size(x, 2))) |> gpu
# this makes a meshgrid, which is multiplied by the learned affine transform
grid = [i == 3 ? 1 : i == 1 ? (x-1)%28+1 : x÷28+1 for x=1:(28^2), i=1:3]
m = Chain(
  STN(localizer, grid),
  Conv((2,2), 1=>16, relu),
  x -> maxpool(x, (2,2)),
  Conv((2,2), 16=>8, relu),
  x -> maxpool(x, (2,2)),
  x -> reshape(x, :, size(x, 4)),
  Dense(288, 10), softmax) |> gpu

m(train[1][1][:, :, :, 1:10])

loss(x, y) = crossentropy(m(x), y)

accuracy(x, y) = mean(onecold(m(x)) .== onecold(y))

evalcb = throttle(() -> @show(accuracy(tX, tY)), 10)
opt = ADAM(params(m))

Flux.train!(loss, train, opt, cb = evalcb)

```

---

<div class="post-metadata">

**Author:** ![improbable22](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/improbable22/32/5464_2.png) [@improbable22](https://discourse.julialang.org/u/improbable22)\
**Post date:** [December 17, 2018, 12:02pm UTC](https://discourse.julialang.org/t/setting-flux-custom-gradients/18728/2 "2018-12-17T12:02:28Z")

</div>

Maybe you will get a more detailed answer, but for a start: I’m not certain that `@grad` understands callable types which aren’t `Function`s. Can you make this work in the simplest possible example? Perhaps you need to define `forward(f, args...)` yourself?

---

<div class="post-metadata">

**Author:** ![platawiec](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/platawiec/32/31914_2.png) [@platawiec](https://discourse.julialang.org/u/platawiec)\
**Post date:** [December 17, 2018, 12:58pm UTC](https://discourse.julialang.org/t/setting-flux-custom-gradients/18728/3 "2018-12-17T12:58:42Z")

</div>

Yes, if I do the following, using `itp` defined above:

```julia
foo = (x, y) -> itp(x, y, 1, 1)
y, back = Flux.Tracker.forward(foo, 1, 2)

```

then it works, but I’m not sure how to “broadcast” over the closure (by which I mean not broadcasting `Flux.Tracker.forward`, but broadcasting over a `TrackedArray` `x` and `y`, as well as the last integer in the `itp` argument, which in this case represents the mini-batch number.

---

<div class="post-metadata">

**Author:** ![improbable22](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/improbable22/32/5464_2.png) [@improbable22](https://discourse.julialang.org/u/improbable22)\
**Post date:** [December 17, 2018, 2:09pm UTC](https://discourse.julialang.org/t/setting-flux-custom-gradients/18728/4 "2018-12-17T14:09:08Z")

</div>

OK, what I said isn’t the problem, nice to know that that works.

What exactly is the error in your minimal example, and can you reproduce this from some simple type you make up, not from another package?

In the full example, note that Flux’s broadcasting machinery uses ForwardDiff dual numbers, not its own TrackedReal, for broadcasting operations:

```julia
g(x) = (println(typeof(x)); x^2)
Flux.gradient(x -> sum(g.(x)), rand(2))

```

Thus the gradient you define for `itp` won’t be used. I wonder if it wouldn’t be simpler to define one custom gradient for `m::STN`.

---

<div class="post-metadata">

**Author:** ![platawiec](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/platawiec/32/31914_2.png) [@platawiec](https://discourse.julialang.org/u/platawiec)\
**Post date:** [December 17, 2018, 2:29pm UTC](https://discourse.julialang.org/t/setting-flux-custom-gradients/18728/5 "2018-12-17T14:29:19Z")

</div>

Thanks, I will try making a pared-down example. In retrospect the minimal example is a bit too minimal and misses the issue that comes up in the full example.

> Thus the gradient you define for `itp` won’t be used. I wonder if it wouldn’t be simpler to define one custom gradient for `m::STN` .

This is a good insight, I tried to define a custom gradient on a non-broadcasted `itp` (so the function call would be `itp(x::TrackedArray, y::TrackedArray, c, n)` but couldn’t get it to work at the time. I will experiment some more.

---

<div class="post-metadata">

**Author:** ![platawiec](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/platawiec/32/31914_2.png) [@platawiec](https://discourse.julialang.org/u/platawiec)\
**Post date:** [December 17, 2018, 7:49pm UTC](https://discourse.julialang.org/t/setting-flux-custom-gradients/18728/6 "2018-12-17T19:49:58Z")

</div>

I tried the following:

```julia
function (m::STN)(x)
  itp = ITP.extrapolate(ITP.interpolate(x,
    (QuadInterp, QuadInterp, ITP.NoInterp(), ITP.NoInterp())), zero(eltype(x)))

  interp_grid_x = x[:, :, :, :]
  interp_grid_y = x[:, :, :, :]

  itp.(interp_grid_x, interp_grid_y, 1, reshape(1:size(x, 4), 1, 1, 1, size(x, 4)))
end

```

And this actually worked, which suggests to me that it had something to do with the `localizer` network. I also tried the other variants, which forced `itp` to `extrapolate`, and I tested on subsets of `x` (like `x[1:3, 1:2, :, :]`) with the corresponding modifications in the `Chain` definition. Those all seemed to work.

I started running into trouble here:

```julia
function (m::STN)(x)
  itp = ITP.extrapolate(ITP.interpolate(x,
    (QuadInterp, QuadInterp, ITP.NoInterp(), ITP.NoInterp())), zero(eltype(x)))

  y = m.localizer(x)
  interp_grid_x = reshape(y, 3, 2, 1, size(x, 4))
  interp_grid_y = reshape(y, 3, 2, 1, size(x, 4))
    
  itp.(interp_grid_x, interp_grid_y, 1, reshape(1:size(x, 4), 1, 1, 1, size(x, 4)))
end

```

Which gave me an error `DimensionMismatch("arrays could not be broadcast to a common size")` during the backward pass, even though `size(interp_grid_x)` is the same in both cases and the forward pass was fine.

The main difference is that I call `y = m.localizer(x)`, where `localizer` is a neural network itself. Maybe I shouldn’t be calling it in this function, and should structure it differently?

**EDIT:** If I use

```julia
function (m::STN)(x)
  y = m.localizer(x)
  reshape(y, 3, 2, 1, size(x, 4))
end

```

then there are no issues, so it is somehow the combination of broadcasting + the call to the model?

---

<div class="post-metadata">

**Author:** ![improbable22](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/improbable22/32/5464_2.png) [@improbable22](https://discourse.julialang.org/u/improbable22)\
**Post date:** [December 18, 2018, 8:24am UTC](https://discourse.julialang.org/t/setting-flux-custom-gradients/18728/7 "2018-12-18T08:24:24Z")

</div>

No great ideas… but just to be clear, `m.localizer(x)` is a Flux network, and the plan now is to let ForwardDiff act on `itp` and not provide any derivatives by hand?

Note that you have `reshape(y` twice. Also that `zero(eltype(x)))` may not match what the broadcasting backward pass will feed to `itp`. I don’t know whether these matter.