AutoGrad.jl is an automatic differentiation package for Julia. It started as a port of the popular Python autograd package and forms the foundation of the Knet Julia deep learning framework. AutoGrad can differentiate regular Julia code that includes loops, conditionals, helper functions, closures etc. by keeping track of the primitive operations and using this execution trace to compute gradients. It uses reverse mode differentiation (a.k.a. backpropagation) so it can efficiently handle functions with large array inputs and scalar outputs. It can compute gradients of gradients to handle higher order derivatives.
You can install AutoGrad in Julia using:
julia> using Pkg; Pkg.add("AutoGrad")In order to use it in your code start with:
using AutoGradx = Param([1,2,3]) # user declares parameters
x => P([1,2,3]) # they are wrapped in a struct
value(x) => [1,2,3] # we can get the original value
sum(abs2,x) => 14 # they act like regular values outside of differentiation
y = @diff sum(abs2,x) # if you want the gradients
y => T(14) # you get another struct
value(y) => 14 # which represents the same value
grad(y,x) => [2,4,6] # but also contains gradients for all ParamsPre v1.1 AutoGrad only supported the following grad interface. This is still supported.
x = [1,2,3]
f(x) = sum(abs2,x)
g = grad(f)
f(x) => 14
g(x) => [2,4,6]Here is a linear regression example using callable objects:
struct Linear; w; b; end # user defines a model
(f::Linear)(x) = (f.w * x .+ f.b)
# Initialize a model as a callable object with parameters:
f = Linear(Param(randn(10,100)), Param(randn(10)))
# SGD training loop:
for (x,y) in data
loss = @diff sum(abs2,f(x)-y)
for w in params(f)
g = grad(loss,w)
axpy!(-0.01, g, w)
end
endSee the examples directory for more examples.
AutoGrad can only handle a function if the primitives it uses have known gradients. You can
add your own primitives with gradients using the @primitive and @zerograd macros in
macros.jl Here is an
example:
@primitive log(x),dy,y (dy .* (1 ./ x))The @primitive macro marks the log(::Any) method as a new primitive and the next
expression defines a gradient function wrt the first argument. The gradient expressions can
refer to the parameter(s) x, the return variable y and its gradient dy (optionally
indicated after the argument list) in the method declaration. For functions with multiple
inputs multiple gradient expressions may be given. Non-existent or zero gradients can be
specified by omitting a gradient expression or using nothing in place of one. By default
the broadcasting version log.(x) is also defined as a primitive, use the @primitive1
macro if you don't want this.
Note that Julia supports multiple-dispatch, i.e. a function may have multiple methods each
supporting different argument types. For example log(::Float32) and log(::BigFloat) are
two different log methods. In AutoGrad.jl each method can be defined independently as a
primitive and can have its own specific gradient. Generally AutoGrad defines gradients
without using argument types to keep the rules generic.
To view the contents of the computational graph after differentiating a function you can use the following:
julia> AutoGrad.gcnode(::AutoGrad.Node)=nothing # without this some values may be lost
julia> w = Param(rand(2,3)); b = Param(rand(2,1)); x = rand(3,4); y = rand(2,4);
julia> J = @diff sum(abs2, w*x .+ b - y)
T(14.695603907991153)
julia> [J] # displaying J in an Array causes pretty printing
1. P(Array{Float64,2}(2,3)) ∇=Array{Float64,2}(2,3)
2. Array{Float64,2}(2,4) = *(Array{Float64,2}(2,3), Array{Float64,2}(3,4))) ∇=Array{Float64,2}(2,4)
3. P(Array{Float64,2}(2,1)) ∇=Array{Float64,2}(2,1)
4. Array{Float64,2}(2,4) = broadcast(+, Array{Float64,2}(2,4), Array{Float64,2}(2,1))) ∇=Array{Float64,2}(2,4)
5. Array{Float64,2}(2,4) = -(Array{Float64,2}(2,4), Array{Float64,2}(2,4))) ∇=Array{Float64,2}(2,4)
6. 14.695603907991153 = sum(abs2, Array{Float64,2}(2,4))) ∇=1.0
julia> z = collect(J.list) # collect creates a Node array with reverse order
julia> dump(z[5], maxdepth=1) # allowing you to look at individual Nodes and Values
AutoGrad.Node
Value: AutoGrad.Result{Array{Float64,2}}
parents: Array{AutoGrad.Node}((2,))
children: Array{AutoGrad.Node}((1,))
outgrad: Array{Float64}((2, 4)) [3.82753 2.19124 3.26769 3.0075; 2.81565 2.3903 1.84373 1.60228]
cdr: AutoGrad.Node
julia> dump(z[5].Value, maxdepth=2)
AutoGrad.Result{Array{Float64,2}}
value: Array{Float64}((2, 4)) [1.16724 1.07224 0.935047 0.895262; 0.687182 0.589704 0.517114 0.495718]
func: * (function of type typeof(*))
args: Tuple{Param{Array{Float64,2}},Array{Float64,2}}
1: Param{Array{Float64,2}}
2: Array{Float64}((3, 4)) [0.515282 0.257471 0.140791 0.127632; 0.705288 0.783289 0.361965 0.311965; 0.780549 0.691645 0.853317 0.843374]
kwargs: Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}
data: NamedTuple{(),Tuple{}} NamedTuple()
itr: Tuple{} ()To profile AutoGrad using TimerOutputs.jl, set the environment variable
ENV["AUTOGRAD_TIMER"]="true" and rebuild AutoGrad with Pkg.build("AutoGrad"), before
evaluating using AutoGrad. The environment variable AUTOGRAD_TIMER is only checked at
compile time, not at run time for performance reasons. This will collect detailed timing
information but slows the code down, when you are done don't forget to
delete!(ENV,"AUTOGRAD_TIMER") and rebuild AutoGrad. In the example below, the symbol sum
indicates the time spent on the forward pass of the sum function and sum[2] indicates
the time spent on the backward pass for the second argument. record and sum_outgrads are
functions internal to AutoGrad.
julia> ENV["AUTOGRAD_TIMER"]="true"
julia> using Pkg; Pkg.build("AutoGrad")
julia> using AutoGrad, TimerOutputs
julia> reset_timer!(AutoGrad.to)
julia> w = Param(rand(2,3)); b = Param(rand(2,1)); x = rand(3,4); y = rand(2,4);
julia> J = @diff sum(abs2, w*x .+ b - y)
julia> AutoGrad.to
───────────────────────────────────────────────────────────────────────
Time Allocations
────────────────────── ───────────────────────
Tot / % measured: 4.62s / 30.4% 546MiB / 25.0%
Section ncalls time %tot avg alloc %tot avg
───────────────────────────────────────────────────────────────────────
+.[2] 1 328ms 23.3% 328ms 46.4MiB 34.1% 46.4MiB
sum[2] 1 288ms 20.5% 288ms 40.0MiB 29.4% 40.0MiB
* 1 38.8ms 2.76% 38.8ms 595KiB 0.43% 595KiB
* 1 269ms 19.2% 269ms 955KiB 0.68% 955KiB
+. 1 139ms 9.92% 139ms 20.4MiB 15.0% 20.4MiB
*[1] 1 117ms 8.33% 117ms 9.41MiB 6.90% 9.41MiB
record 4 88.7ms 6.31% 22.2ms 3.49MiB 2.56% 894KiB
-[1] 1 65.9ms 4.69% 65.9ms 10.0MiB 7.32% 10.0MiB
- 1 55.8ms 3.97% 55.8ms 929KiB 0.67% 929KiB
sum 1 50.0ms 3.56% 50.0ms 4.68MiB 3.44% 4.68MiB
+.[1] 1 1.78ms 0.13% 1.78ms 37.7KiB 0.03% 37.7KiB
sum_outgrads 5 1.41ms 0.10% 282μs 28.2KiB 0.02% 5.64KiB
───────────────────────────────────────────────────────────────────────core.jl implements the
main functionality and acts as the main documentation source.
macros.jl has some
support functions to define and test new primitives.
getindex.jl,
iterate.jl and
cat.jl set up support
for common data structures including Arrays, Tuples, and Dictionaries. The numerical
gradients are defined in files such as base.jl and math.jl.
The gradient coverage and unit testing are spotty, I am still adding more gradients and
tests to cover the Julia base. Documentation needs to be improved. Overwriting functions
(e.g. setindex!) are not supported. Efficiency could be improved by reducing runtime
compilation, memoization, and support for static computation.
AutoGrad.jl was written by Deniz Yuret. Parts of the code were initially ported from the Python autograd package. I'd like to thank autograd author Dougal Maclaurin for his support. See (Baydin et al. 2015) for a general review of automatic differentiation, autograd tutorial for some Python examples, and Dougal's PhD thesis for design principles. JuliaDiff and FluxML have alternative differentiation tools for Julia. I would like to thank the current contributors:
- Carlo Lucibello
- Ekin Akyürek
- Emre Yolcu
- Jarrett Revels
- Mike Innes
- Ozan Arkan Can
- Rene Donner
The suggested citation for AutoGrad is:
@inproceedings{knet2016mlsys,
author={Yuret, Deniz},
title={Knet: beginning deep learning with 100 lines of Julia},
year={2016},
booktitle={Machine Learning Systems Workshop at NIPS 2016}
}