Flux.jl: Compatibility with Tracker

[MikeInnes]

We could use an option for going back to Tracker for AD. This would be useful for debugging in cases like #876. It would also serve as a model for people who want to use more specialised AD implementations for whatever purpose.

It should more-or-less just be a case of (a) wrapping gradients in Tracker.param and (b) making Zygote.Params work with Tracker's gradient function (the same thing we originally did for Zygote in reverse).

[MikeInnes]

This works well enough with https://github.com/FluxML/Zygote.jl/pull/362.

using Flux, Tracker

track(m) = fmap(x -> x isa AbstractArray ? Tracker.param(x) : x, m)

model = Chain(Dense(10, 5, relu), Dense(5, 2), softmax) |> track

gs = Tracker.gradient(params(model)) do
  Flux.mse(model(rand(10)), rand(2))
end

IdDict(p=>gs[p] for p in params(model))

This archive contains a project and manifest to get it running quickly. cc @jessebett

[jessebett]

In the recent update in #876 I show that Tracker gradients are slightly different than Zygote's. However I am unable to test if this causes Zygote NaN's during training where Tracker wouldn't because I am unable to optimize both Tracker and Zygote models with the same package statuses. Even with the above compatibility changes, Optimise is incompatible with Tracker.

# Set up Models

using Random
using Flux
using Tracker
using MLDataUtils
using Zygote
using Statistics: mean
using MLDataUtils

# dummy data
x = Float32.(rand(113,100000))
y = sum(x.^2,dims=1)
dataset = batchview((x,y),size=256)
x1,y1 = first(dataset)

# Flux Model
model = Chain(
              Dense(113,1000),
              Dense(1000,1)
             )
θ = params(model)

# Compatibility with Tracker
track(m) = fmap(x -> x isa AbstractArray ? Tracker.param(x) : x, m)
t_model= track(model)
t_θ = params(t_model)

# Objective
criterion(logits,y) = mean(Flux.logitbinarycrossentropy.(logits,y))

# Optimizer
optimizer = Flux.ADAM(1e-4,(0.9,0.99))

# Tracker Optimization, will NaN after a ~200 iterations
for (x,y) in dataset
  θ = params(model)
  loss,back = Flux.Zygote.pullback(()->criterion(model(x),y),θ)
  println(loss)
  grads = back(1.)
  Flux.Optimise.update!(optimizer,θ,grads)
end

# Tracker Optimization
for (x,y) in dataset
  t_loss,t_back = Tracker.forward(()->criterion(t_model(x),y),t_θ)
  println(t_loss)
  t_grads = t_back(1.)
  Flux.Optimise.update!(optimizer,t_θ,t_grads)
end

While the Zygote Optimization block eventually NaNs, I can't run the Tracker block as is because I get the following error:

ERROR: MethodError: no method matching Float32(::Tracker.TrackedReal{Float64})
Closest candidates are:
  Float32(::Real, ::RoundingMode) where T<:AbstractFloat at rounding.jl:200
  Float32(::T) where T<:Number at boot.jl:718
  Float32(::Int8) at float.jl:60
  ...
Stacktrace:
 [1] convert(::Type{Float32}, ::Tracker.TrackedReal{Float64}) at ./number.jl:7
 [2] setindex! at ./array.jl:782 [inlined]
 [3] setindex! at ./multidimensional.jl:488 [inlined]
 [4] macro expansion at ./broadcast.jl:909 [inlined]
 [5] macro expansion at ./simdloop.jl:77 [inlined]
 [6] copyto! at ./broadcast.jl:908 [inlined]
 [7] copyto! at ./broadcast.jl:863 [inlined]
 [8] materialize!(::Array{Float32,2}, ::Base.Broadcast.Broadcasted{Tracker.TrackedStyle,Nothing,typeof(+),Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{2},Nothing,typeof(*),Tuple{Float64,Array{Float32,2}}},Base.Broadcast.Broadcasted{Tracker.TrackedStyle,Nothing,typeof(*),Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0},Nothing,typeof(-),Tuple{Int64,Float64}},TrackedArray{…,Array{Float64,2}}}}}}) at ./broadcast.jl:822
 [9] apply!(::ADAM, ::TrackedArray{…,Array{Float32,2}}, ::TrackedArray{…,Array{Float64,2}}) at /Users/jesse/.julia/packages/Flux/WxpHb/src/optimise/optimisers.jl:104
 [10] update!(::ADAM, ::TrackedArray{…,Array{Float32,2}}, ::TrackedArray{…,Array{Float64,2}}) at /Users/jesse/.julia/packages/Flux/WxpHb/src/optimise/train.jl:10
 [11] update!(::ADAM, ::Zygote.Params, ::Tracker.Grads) at /Users/jesse/.julia/packages/Flux/WxpHb/src/optimise/train.jl:16
 [12] top-level scope at REPL[37]:5

[DhairyaLGandhi]

Defining this should be sufficient.

function Flux.Optimise.update!(opt, xs::Tracker.Params, gs)
  for x in xs
   Flux.Optimise.update!(opt, x, gs[x])
  end
end

function Flux.Optimise.update!(opt, x, x̄)
  Tracker.update!(x, -Flux.Optimise.apply!(opt, x, Tracker.data(x̄)))
end

Maybe could add this in a tracker compat section?

[SaremS]

@dhairyagandhi96 Tried your solution as I was facing a similar problem but now I am getting

`UndefRefError: access to undefined reference

Stacktrace:
[1] getproperty at ./Base.jl:20 [inlined]
[2] update!(::TrackedArray{…,Array{Float32,2}}, ::Array{Float64,2}) at /Users/saremseitz/.julia/packages/Tracker/JhqMQ/src/lib/array.jl:72
[4] update!(::ADAM, ::Zygote.Params, ::Tracker.Grads) at /Users/saremseitz/.julia/packages/Flux/oX9Pi/src/optimise/train.jl:16

[freddycct]

Hi, I'm still hoping we can use Tracker instead as a backwards compatibility option, it's more stable and does not crash, is it possible?

[AStupidBear]

@freddycct @SaremS I collected the above snippets and extended to create a package TrackerFlux.jl.