update docs
This commit is contained in:
Mike J Innes
parent
10a169bb77
commit
dda51a0140
|
|
@ -6,6 +6,52 @@ Backpropagation, or reverse-mode automatic differentiation, is handled by the `F
|
|||
julia> using Flux.Tracker
|
||||
```
|
||||
|
||||
Here we discuss some more advanced uses of this module, as well as covering its internals.
|
||||
|
||||
## Taking Gradients
|
||||
|
||||
In the [basics section](../models/basics.md) we covered basic usage of the `gradient` function.
|
||||
|
||||
```julia
|
||||
using Flux.Tracker
|
||||
|
||||
Tracker.gradient((a, b) -> a*b, 2, 3) # (3.0 (tracked), 2.0 (tracked))
|
||||
```
|
||||
|
||||
`gradient` is actually just a thin wrapper around the backpropagator-based interface, `forward`.
|
||||
|
||||
```julia
|
||||
using Flux.Tracker: forward
|
||||
|
||||
y, back = forward((a, b) -> a*b, 2, 3) # (6.0 (tracked), Flux.Tracker.#9)
|
||||
|
||||
back(1) # (3.0 (tracked), 2.0 (tracked))
|
||||
```
|
||||
|
||||
The `forward` function returns two results. The first, `y`, is the original value of the function (perhaps with tracking applied). The second, `back`, is a new function which, given a sensitivity, returns the sensitivity of the inputs to `forward` (we call this a "backpropagator"). One use of this interface is to provide custom sensitivities when outputs are not scalar.
|
||||
|
||||
```julia
|
||||
julia> y, back = forward((a, b) -> a.*b, [1,2,3],[4,5,6])
|
||||
(param([4.0, 10.0, 18.0]), Flux.Tracker.#9)
|
||||
|
||||
julia> back([1,1,1])
|
||||
(param([4.0, 5.0, 6.0]), param([1.0, 2.0, 3.0]))
|
||||
```
|
||||
|
||||
We can also take gradients in-place. This can be useful if you only care about first-order gradients.
|
||||
|
||||
```julia
|
||||
a, b = param(2), param(3)
|
||||
|
||||
c = a*b # 6.0 (tracked)
|
||||
|
||||
Tracker.back!(c)
|
||||
|
||||
Tracker.grad(a), Tracker.grad(b) # (3.0, 2.0)
|
||||
```
|
||||
|
||||
## Tracked Arrays
|
||||
|
||||
The `param` function converts a normal Julia array into a new object that, while behaving like an array, tracks extra information that allows us to calculate derivatives. For example, say we multiply two parameters:
|
||||
|
||||
```julia
|
||||
|
|
@ -41,7 +87,48 @@ julia> x.grad
|
|||
-2.0
|
||||
```
|
||||
|
||||
## Internals
|
||||
You may sometimes want to drop derivative information and just get the plain value back. You can do this by calling `Tracker.data(W)`.
|
||||
|
||||
## Custom Gradients
|
||||
|
||||
We can hook in to the processes above to implement custom gradients for a function or kernel. For a toy example, imagine a custom implementation of `minus`:
|
||||
|
||||
```julia
|
||||
minus(a, b) = a - b
|
||||
```
|
||||
|
||||
Firstly, we must tell the tracker system to stop when it sees a call to `minus`, and record it. We can do this using dispatch:
|
||||
|
||||
```julia
|
||||
using Flux.Tracker: TrackedReal, track, @grad
|
||||
|
||||
minus(a::TrackedArray, b::TrackedArray) = Tracker.track(minus, a, b)
|
||||
```
|
||||
|
||||
`track` takes care of building a new `Tracked` object and recording the operation on the tape. We just need to provide a gradient definition.
|
||||
|
||||
```julia
|
||||
@grad function minus(a, b)
|
||||
return minus(data(a),data(b)), Δ -> (Δ, -Δ)
|
||||
end
|
||||
```
|
||||
|
||||
This is essentially just a way of overloading the `forward` function we saw above. We strip tracking from `a` and `b` so that we are calling the original definition of `minus` (otherwise, we'd just try to track the call again and hit an infinite regress).
|
||||
|
||||
Note that in the backpropagator we don't call `data(a)`; we *do* in fact want to track this, since nest AD will take a derivative through the backpropagator itself. For example, the gradient of `*` might look like this.
|
||||
|
||||
```julia
|
||||
@grad a * b = data(a)*data(b), Δ -> (Δ*b, a*Δ)
|
||||
```
|
||||
|
||||
For multi-argument functions with custom gradients, you likely want to catch not just `minus(::TrackedArray, ::TrackedArray)` but also `minus(::Array, TrackedArray)` and so on. To do so, just define those extra signatures as needed:
|
||||
|
||||
```julia
|
||||
minus(a::AbstractArray, b::TrackedArray) = Tracker.track(minus, a, b)
|
||||
minus(a::TrackedArray, b::AbstractArray) = Tracker.track(minus, a, b)
|
||||
```
|
||||
|
||||
## Tracked Internals
|
||||
|
||||
All `Tracked*` objects (`TrackedArray`, `TrackedReal`) are light wrappers around the `Tracked` type, which you can access via the `.tracker` field.
|
||||
|
||||
|
|
@ -50,14 +137,9 @@ julia> x.tracker
|
|||
Flux.Tracker.Tracked{Array{Float64,1}}(0x00000000, Flux.Tracker.Call{Void,Tuple{}}(nothing, ()), true, [5.0, 6.0], [-2.0, -2.0])
|
||||
```
|
||||
|
||||
The `Tracker` stores the value and gradient of a given object, which we've seen before.
|
||||
The `Tracker` stores the gradient of a given object, which we've seen before.
|
||||
|
||||
```julia
|
||||
julia> x.tracker.data
|
||||
2-element Array{Float64,1}:
|
||||
5.0
|
||||
6.0
|
||||
|
||||
julia> x.tracker.grad
|
||||
2-element Array{Float64,1}:
|
||||
-2.0
|
||||
|
|
@ -86,71 +168,4 @@ When we call `back!(y, [1, -1])`, the sensitivities `[1, -1]` simply get forward
|
|||
Tracker.back(*, [1, -1], W, x)
|
||||
```
|
||||
|
||||
which in turn calculates the sensitivities of the arguments (`W` and `x`) and backpropagates through their calls. This is recursive, so it will walk the entire program graph and propagate gradients to the original model parameters.
|
||||
|
||||
## Custom Gradients
|
||||
|
||||
We can hook in to the processes above to implement custom gradients for a function or kernel. For a toy example, imagine a custom implementation of `minus`:
|
||||
|
||||
```julia
|
||||
julia> minus(a, b) = a - b
|
||||
```
|
||||
|
||||
Firstly, we must tell the tracker system to stop when it sees a call to `minus`, and record it. We can do this using dispatch:
|
||||
|
||||
```julia
|
||||
julia> minus(a::TrackedArray, b::TrackedArray) = Tracker.track(minus, a, b)
|
||||
minus (generic function with 2 methods)
|
||||
```
|
||||
|
||||
`Tracker.track` does two things: (1) it makes sure `minus` is called with *normal* array, not tracked ones (you can use `@show` inside `minus` to verify this), and (2) it uses the result to add a `minus` node to the tape. Look inside the result of calling `minus` to see what happened:
|
||||
|
||||
```julia
|
||||
julia> a, b = param([6,5,4]), param([1,2,3])
|
||||
(param([6.0, 5.0, 4.0]), param([1.0, 2.0, 3.0]))
|
||||
|
||||
julia> c = minus(a, b)
|
||||
Tracked 3-element Array{Float64,1}:
|
||||
5.0
|
||||
3.0
|
||||
1.0
|
||||
|
||||
julia> c.tracker.f
|
||||
Flux.Tracker.Call{...}(minus, (param([6.0, 5.0, 4.0]), param([1.0, 2.0, 3.0])))
|
||||
```
|
||||
|
||||
Finally, we have to specify the gradient of `minus`.
|
||||
|
||||
```julia
|
||||
julia> Tracker.back(::typeof(minus), Δ, a, b) =
|
||||
(Tracker.@back(a, Δ); Tracker.@back(b, -Δ))
|
||||
```
|
||||
|
||||
`@back(x, Δ)` tells the tracker to continue propagating the sensitivity `Δ` through `x`. Now, AD will work with any program that calls `minus`.
|
||||
|
||||
```julia
|
||||
julia> Flux.back!(c, 1)
|
||||
|
||||
julia> a.grad
|
||||
3-element Array{Float64,1}:
|
||||
1.0
|
||||
1.0
|
||||
1.0
|
||||
|
||||
julia> b.grad
|
||||
3-element Array{Float64,1}:
|
||||
-1.0
|
||||
-1.0
|
||||
-1.0
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
For multi-argument functions with custom gradients, you likely want to catch not just `minus(::TrackedArray, ::TrackedArray)` but also `minus(::Array, TrackedArray)` and so on. To do so, just define those extra signatures as needed:
|
||||
|
||||
```julia
|
||||
minus(a::AbstractArray, b::TrackedArray) = Tracker.track(minus, a, b)
|
||||
minus(a::TrackedArray, b::AbstractArray) = Tracker.track(minus, a, b)
|
||||
```
|
||||
|
||||
`@back` *must* be called exactly once on each tracked input argument. You do not need to do any special handling if one of the arguments is not tracked, as `@back` will just become a no-op.
|
||||
which in turn calculates the sensitivities of the arguments (`W` and `x`) and back-propagates through their calls. This is recursive, so it will walk the entire program graph and propagate gradients to the original model parameters.
|
||||
|
|
|
|||
|
|
@ -2,20 +2,74 @@
|
|||
|
||||
## Taking Gradients
|
||||
|
||||
Consider a simple linear regression, which tries to predict an output array `y` from an input `x`. (It's a good idea to follow this example in the Julia repl.)
|
||||
Flux's core feature is taking gradients of Julia code. The `gradient` function takes another Julia function `f` and a set of arguments, and returns the gradient with respect to each argument. (It's a good idea to try pasting these examples in the Julia terminal.)
|
||||
|
||||
```julia
|
||||
using Flux.Tracker
|
||||
|
||||
f(x) = 3x^2 + 2x + 1
|
||||
|
||||
# df/dx = 6x + 2
|
||||
f′(x) = Tracker.gradient(f, x)[1]
|
||||
|
||||
f′(2) # 14.0 (tracked)
|
||||
|
||||
# d²f/dx² = 6
|
||||
f′′(x) = Tracker.gradient(f′, x)[1]
|
||||
|
||||
f′′(2) # 6.0 (tracked)
|
||||
```
|
||||
|
||||
(We'll learn more about why these numbers show up as `(tracked)` below.)
|
||||
|
||||
When a function has many parameters, we can pass them all in explicitly:
|
||||
|
||||
```julia
|
||||
f(W, b, x) = W * x + b
|
||||
|
||||
Tracker.gradient(f, 2, 3, 4)
|
||||
(4.0 (tracked), 1.0, 2.0 (tracked))
|
||||
```
|
||||
|
||||
But machine learning models can have *hundreds* of parameters! Flux offers a nice way to handle this. We can tell Flux to treat something as a parameter via `param`. Then we can collect these together and tell `gradient` to collect the gradients of all of them at once.
|
||||
|
||||
```julia
|
||||
W = param(2) # 2.0 (tracked)
|
||||
b = param(3) # 3.0 (tracked)
|
||||
|
||||
f(x) = W * x + b
|
||||
|
||||
params = Params([W, b])
|
||||
grads = Tracker.gradient(() -> f(4), params)
|
||||
|
||||
grads[W] # 4.0
|
||||
grads[b] # 1.0
|
||||
```
|
||||
|
||||
There are a few things to notice here. Firstly, `W` and `b` now show up as *tracked*. Tracked things behave like normal numbers or arrays, but keep records of everything you do with them, allowing Flux to calculate their gradients. `gradient` takes a zero-argument function; no arguments are necessary because the `Params` tell it what to differentiate.
|
||||
|
||||
This will come in really handy when dealing with big, complicated models. For now, though, let's start with something simple.
|
||||
|
||||
## Simple Models
|
||||
|
||||
Consider a simple linear regression, which tries to predict an output array `y` from an input `x`.
|
||||
|
||||
```julia
|
||||
W = rand(2, 5)
|
||||
b = rand(2)
|
||||
|
||||
predict(x) = W*x .+ b
|
||||
loss(x, y) = sum((predict(x) .- y).^2)
|
||||
|
||||
function loss(x, y)
|
||||
ŷ = predict(x)
|
||||
sum((y .- ŷ).^2)
|
||||
end
|
||||
|
||||
x, y = rand(5), rand(2) # Dummy data
|
||||
loss(x, y) # ~ 3
|
||||
```
|
||||
|
||||
To improve the prediction we can take the gradients of `W` and `b` with respect to the loss function and perform gradient descent. We could calculate gradients by hand, but Flux will do it for us if we tell it that `W` and `b` are trainable *parameters*.
|
||||
To improve the prediction we can take the gradients of `W` and `b` with respect to the loss and perform gradient descent. Let's tell Flux that `W` and `b` are parameters, just like we did above.
|
||||
|
||||
```julia
|
||||
using Flux.Tracker
|
||||
|
|
@ -23,17 +77,15 @@ using Flux.Tracker
|
|||
W = param(W)
|
||||
b = param(b)
|
||||
|
||||
l = loss(x, y)
|
||||
|
||||
back!(l)
|
||||
gs = Tracker.gradient(() -> loss(x, y), Params([W, b]))
|
||||
```
|
||||
|
||||
`loss(x, y)` returns the same number, but it's now a *tracked* value that records gradients as it goes along. Calling `back!` then accumulates the gradient of `W` and `b`. We can see what this gradient is, and modify `W` to train the model.
|
||||
Now that we have gradients, we can pull them out and update `W` to train the model. The `update!(W, Δ)` function applies `W = W + Δ`, which we can use for gradient descent.
|
||||
|
||||
```julia
|
||||
using Flux.Tracker: grad, update!
|
||||
using Flux.Tracker: update!
|
||||
|
||||
Δ = grad(W)
|
||||
Δ = gs[W]
|
||||
|
||||
# Update the parameter and reset the gradient
|
||||
update!(W, -0.1Δ)
|
||||
|
|
@ -43,7 +95,7 @@ loss(x, y) # ~ 2.5
|
|||
|
||||
The loss has decreased a little, meaning that our prediction `x` is closer to the target `y`. If we have some data we can already try [training the model](../training/training.md).
|
||||
|
||||
All deep learning in Flux, however complex, is a simple generalisation of this example. Of course, models can *look* very different – they might have millions of parameters or complex control flow, and there are ways to manage this complexity. Let's see what that looks like.
|
||||
All deep learning in Flux, however complex, is a simple generalisation of this example. Of course, models can *look* very different – they might have millions of parameters or complex control flow. Let's see how Flux handles more complex models.
|
||||
|
||||
## Building Layers
|
||||
|
||||
|
|
|
|||
|
|
@ -103,7 +103,7 @@ m.(seq)
|
|||
|
||||
## Truncating Gradients
|
||||
|
||||
By default, calculating the gradients in a recurrent layer involves the entire history. For example, if we call the model on 100 inputs, calling `back!` will calculate the gradient for those 100 calls. If we then calculate another 10 inputs we have to calculate 110 gradients – this accumulates and quickly becomes expensive.
|
||||
By default, calculating the gradients in a recurrent layer involves its entire history. For example, if we call the model on 100 inputs, we'll have to calculate the gradient for those 100 calls. If we then calculate another 10 inputs we have to calculate 110 gradients – this accumulates and quickly becomes expensive.
|
||||
|
||||
To avoid this we can *truncate* the gradient calculation, forgetting the history.
|
||||
|
||||
|
|
|
|||
|
|
@ -3,6 +3,8 @@
|
|||
Consider a [simple linear regression](../models/basics.md). We create some dummy data, calculate a loss, and backpropagate to calculate gradients for the parameters `W` and `b`.
|
||||
|
||||
```julia
|
||||
using Flux.Tracker
|
||||
|
||||
W = param(rand(2, 5))
|
||||
b = param(rand(2))
|
||||
|
||||
|
|
@ -11,7 +13,9 @@ loss(x, y) = sum((predict(x) .- y).^2)
|
|||
|
||||
x, y = rand(5), rand(2) # Dummy data
|
||||
l = loss(x, y) # ~ 3
|
||||
back!(l)
|
||||
|
||||
params = Params([W, b])
|
||||
grads = Tracker.gradient(() -> loss(x, y), params)
|
||||
```
|
||||
|
||||
We want to update each parameter, using the gradient, in order to improve (reduce) the loss. Here's one way to do that:
|
||||
|
|
@ -22,7 +26,7 @@ using Flux.Tracker: grad, update!
|
|||
function sgd()
|
||||
η = 0.1 # Learning Rate
|
||||
for p in (W, b)
|
||||
update!(p, -η * grad(p))
|
||||
update!(p, -η * grads[p])
|
||||
end
|
||||
end
|
||||
```
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ using MacroTools: @q, @forward
|
|||
|
||||
import Base: ==
|
||||
|
||||
export TrackedArray, TrackedVector, TrackedMatrix, param, back!
|
||||
export TrackedArray, TrackedVector, TrackedMatrix, Params, param, back!
|
||||
|
||||
tracker(x) = nothing
|
||||
|
||||
|
|
@ -61,7 +61,7 @@ macro grad(ex)
|
|||
end
|
||||
|
||||
function update!(x, Δ)
|
||||
tracker(x).data += Δ
|
||||
x.data .+= data(Δ)
|
||||
tracker(x).grad .= 0
|
||||
return x
|
||||
end
|
||||
|
|
|
|||
|
|
@ -72,10 +72,18 @@ end
|
|||
|
||||
@forward Params.params Base.start, Base.next, Base.done
|
||||
|
||||
function Base.show(io::IO, ps::Params)
|
||||
print(io, "Params([")
|
||||
join(io, ps.params, ", ")
|
||||
print(io, "])")
|
||||
end
|
||||
|
||||
struct Grads
|
||||
grads::ObjectIdDict
|
||||
end
|
||||
|
||||
Base.show(io::IO, ps::Grads) = println(io, "Grads(...)")
|
||||
|
||||
Grads() = Grads(ObjectIdDict())
|
||||
|
||||
Grads(ps::Params) = Grads(ObjectIdDict(tracker(p) => init_grad(data(p)) for p in ps))
|
||||
|
|
|
|||
Loading…
Reference in New Issue