Julia port of the Python autograd package.
Author denizyuret
Popularity
148 Stars
Updated Last
6 Months Ago
Started In
August 2016

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.

## Installation

You can install AutoGrad in Julia using:

`julia> using Pkg; Pkg.add("AutoGrad")`

`using AutoGrad`

## Interface

```x = 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

## Old Interface

Pre v1.1 AutoGrad only supported the following `grad` interface. This is still supported.

```x = [1,2,3]
f(x) = sum(abs2,x)
f(x) => 14
g(x) => [2,4,6]```

## Example

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)
axpy!(-0.01, g, w)
end
end```

See 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.

## Debugging and Profiling

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
outgrad: Array{Float64}((2, 4)) [3.82753 2.19124 3.26769 3.0075; 2.81565 2.3903 1.84373 1.60228]
julia> dump(z[5].Value, maxdepth=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> 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)
───────────────────────────────────────────────────────────────────────
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
───────────────────────────────────────────────────────────────────────```

## Code structure

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`.

## Current status and future work

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.

## Acknowledgments and references

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}
}
``````