# How to use hook to clip a gradient?

**URL:** <https://discourse.julialang.org/t/how-to-use-hook-to-clip-a-gradient/37896>\
**Category:** General Usage\
**Tags:** zygote\
**Created:** [April 20, 2020, 12:47pm UTC](https://discourse.julialang.org/t/how-to-use-hook-to-clip-a-gradient/37896 "2020-04-20T12:47:07Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![HenriDeh](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/henrideh/32/8316_2.png) [@HenriDeh](https://discourse.julialang.org/u/HenriDeh)\
**Post date:** [April 20, 2020, 12:47pm UTC](https://discourse.julialang.org/t/how-to-use-hook-to-clip-a-gradient/37896/1 "2020-04-20T12:47:07Z")

</div>

I’m trying to train a neural network using gradient clipping. I’m trying to understand how the hook function works.  
Let’s define a model and a clip function

```julia
using Zygote, LinearAlgebra
clip(g) = g./norm(g)
a = 2; b = 1; m(x,y) = a*x*(y-a)+b
gradient(m, 2, 3) # returns (2,4)
gradient(m,2,3) |> clip # returns (0.4472135954999579, 0.8944271909999159)

```

But I’m using Flux’s train function so I need to directly provide a loss function with a clipped gradient. But this does not work :

```julia
gradient(2,3) do (x,y)
       Zygote.hook(clip, m(x,y))
end
ERROR: MethodError: no method matching (::var"#31#32")(::Int64, ::Int64)
Stacktrace:
 [1] macro expansion at C:\Users\Henri\.julia\packages\Zygote\ApBXe\src\compiler\interface2.jl:0 [inlined]
 [2] _pullback(::Zygote.Context, ::var"#31#32", ::Int64, ::Int64) at C:\Users\Henri\.julia\packages\Zygote\ApBXe\src\compiler\interface2.jl:7
 [3] _pullback(::Function, ::Int64, ::Int64) at C:\Users\Henri\.julia\packages\Zygote\ApBXe\src\compiler\interface.jl:29
 [4] pullback(::Function, ::Int64, ::Int64) at C:\Users\Henri\.julia\packages\Zygote\ApBXe\src\compiler\interface.jl:35
 [5] gradient(::Function, ::Int64, ::Vararg{Int64,N} where N) at C:\Users\Henri\.julia\packages\Zygote\ApBXe\src\compiler\interface.jl:44
 [6] top-level scope at none:0

```

I cannot find any documentation on how to clip a gradient. Which is weird because it’s a frequently used tool in DL I think.

---

<div class="post-metadata">

**Author:** ![MikeInnes](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mikeinnes/32/3656_2.png) [@MikeInnes](https://discourse.julialang.org/u/MikeInnes)\
**Post date:** [April 20, 2020, 1:08pm UTC](https://discourse.julialang.org/t/how-to-use-hook-to-clip-a-gradient/37896/2 "2020-04-20T13:08:53Z")

</div>

The error here is actually in your `do` block; it should be `x, y`, not `(x, y)`, since as written it’s expecting a tuple of two things rather than two arguments.

However, that then returns `(2, 4)`, which isn’t what you want. The reason is that `hook` applies to the gradient of the variable passed to it. This is slightly clearer if you write

```julia
gradient(2,3) do x, y
  l = m(x, y)
  Zygote.hook(clip, l)
end

```

`clip` gets passed `l̄`, the gradient of the loss, which is always `1`, so this doesn’t do anything useful. To apply `clip` to `(x, y)` you have to pass that value to `hook`. Here’s one way to write that:

```julia
gradient(2,3) do x, y
  x, y = Zygote.hook(clip, (x, y))
  m(x, y)
end

```

> [@HenriDeh](#):
>
> I cannot find any documentation on how to clip a gradient. Which is weird because it’s a frequently used tool in DL I think.

If you can think of ways to make this stuff clearer in the docs a pr would be huge! This is definitely the kind of area we’d like to be a bit more comprehensive on.

---

<div class="post-metadata">

**Author:** ![HenriDeh](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/henrideh/32/8316_2.png) [@HenriDeh](https://discourse.julialang.org/u/HenriDeh)\
**Post date:** [April 20, 2020, 1:57pm UTC](https://discourse.julialang.org/t/how-to-use-hook-to-clip-a-gradient/37896/3 "2020-04-20T13:57:25Z")

</div>

Thank you, it’s getting clearer. And how would you proceed to clip the gradient of a loss function? I tried this

```julia
using Flux, LinearAlgebra
model = Chain(Dense(2,1))
clip(g) = g./norm(g)
loss(x,y) = mean((x .- y).^2)
batch = (rand(2,2), rand(1,2))

gs = gradient(Flux.params(model)) do
           loss(model(batch[1]), batch[2])
end

gs_clip = gradient(Flux.params(model)) do
           Zygote.hook(clip, loss(model(batch[1]), batch[2]))
end

```

But gs and gs\_clip are equal. I’d happily make a PR but I need to fully understand this first since the typical use-case in ML is to clip the gradient of the loss function.

Maybe it would also be interesting to provide clip as a built-in utility function.

---

<div class="post-metadata">

**Author:** ![MikeInnes](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mikeinnes/32/3656_2.png) [@MikeInnes](https://discourse.julialang.org/u/MikeInnes)\
**Post date:** [April 20, 2020, 2:02pm UTC](https://discourse.julialang.org/t/how-to-use-hook-to-clip-a-gradient/37896/4 "2020-04-20T14:02:27Z")

</div>

Right, that has the same issue as the original example: you’re applying `clip` to the gradient of the loss value, not to the parameters.

One way to approach this would be to define a function like `gradclip(x) = hook(clip, x)` and then do `model = mapleaves(gradclip, model)` within the forward pass. Definitely agree that we could use some examples and/or utilities for that kind of thing as it’s not obvious.

---

<div class="post-metadata">

**Author:** ![HenriDeh](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/henrideh/32/8316_2.png) [@HenriDeh](https://discourse.julialang.org/u/HenriDeh)\
**Post date:** [April 20, 2020, 2:17pm UTC](https://discourse.julialang.org/t/how-to-use-hook-to-clip-a-gradient/37896/6 "2020-04-20T14:17:03Z")

</div>

Ha yes indeed, this time it was on purpose, I thought the clipping heuristic was to clip the gradient of the error before backpropagating in the parameters. But it’s the gradient w.r.t. the parameters which is clipped. Thanks for your time, I’ll practice.