lungd
commented
4 years ago I think you can define an in-place problem function to make this work (rather a workaround to not hit this issue)
using Flux
using OrdinaryDiffEq
using DiffEqFlux
model1 = FastDense(2,2)
model2 = FastChain(model1)
#initial conditions
u0 = Float32[1.0,1.0]
u1 = Float32[1.5,1.5]
u2 = Float32[2.0,2.0]
#minibatch = [(u0[:,:],)] # note this is 2x1. this matters to error
minibatch = [([u0 u1 u2],)]
# test both models
function do_it()
for model in (model1,model2)
println("\nmodel ",model)
p = initial_params(model)
#set up ODE Problem
#f(x,p,t) = model(x,p) # error tracebakc juno highlits this in red
function f(dx,x,p,t)
dx .= model(x,p)
nothing
end
prob = ODEProblem(f,minibatch[1][1],(0.0f0,1.0f0),p) # u0 is a placeholder.
pred(u,p) = concrete_solve(prob,Tsit5(),u,p, saveat=0.01)
#validate prediction working
pred(u0,p)
pred(minibatch[1][1],p)
# set up loss function
# setup loss function
loss_vec_batch(p,mb) = sum(abs, pred(mb,p).- mb) # vectorized over batch
loss_vec(p) = loss_vec_batch(p,minibatch[1]...)
loss = loss_vec # rename
println(loss(p))
#validate Loss Function working
loss_vec_batch(p,minibatch[1]...) == loss_vec(p)
# set up the gradient parameters
ps = Flux.params(p)
#use zygote on vector case
gsv = Flux.Zygote.gradient(ps) do
x = loss(p)
first(x)
end
println("\nMODEL RAN CLEANLY. ",gsv.grads)
end
end
do_it()julia> do_it()
model FastDense{typeof(identity),DiffEqFlux.var"#initial_params#48"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}(2, 2, identity, DiffEqFlux.var"#initial_params#48"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}(Flux.glorot_uniform, Flux.zeros, 2, 2))
76.2751
MODEL RAN CLEANLY. IdDict{Any,Any}(Float32[0.7815552, -0.026053874, -0.60071194, -0.05056247, 0.0, 0.0] => Float32[325.44238, -300.28796, 296.0859, -272.6577, 202.24295, -186.32872])
model FastChain{Tuple{FastDense{typeof(identity),DiffEqFlux.var"#initial_params#48"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}((FastDense{typeof(identity),DiffEqFlux.var"#initial_params#48"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}(2, 2, identity, DiffEqFlux.var"#initial_params#48"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}(Flux.glorot_uniform, Flux.zeros, 2, 2)),))
251.10022
MODEL RAN CLEANLY. IdDict{Any,Any}(Float32[0.41792107, 0.47331968, 0.24224015, -0.25262898, 0.0, 0.0] => Float32[374.25388, 288.6827, 330.43677, 253.09106, 202.13618, 153.87003])
cems2
ChrisRackauckas
gabrevaya
In a NeuralODE problem I am testing two trivial models
I believe these are nominally identical models. Yet when called in a batch mode, the Zygote Gradient of my loss function fails for the second model but not the first. The failure ONLY occurs when broadcasting over a batch.
code to reproduce:
Analysis The results of running this code are pasted in below. What happens is model1 runs cleanly and model2 has an error in the Zygote.gradient.
This is the MWE because there are two critical things required to trigger this error.
By batch orientation I mean: If you only wanted to run one intial condition at a time in the concrete_solver then the initial condition u0 can be a simple vector. But if you wanted to run multiple initial conditions at the same time then the initial condition is a 2D matrix (a list of initial condition vectors).
In this example, I am using the matrix format-- as though it were a batch-- but only doing it for one initial condition. The error is the same if I tack on additional initial conditions to this matrix. I left it at one to simplify this. The conversion step from a vector to a matrix is u0[:,:] and if you want add more initial conditions in batch orientation:
minibatch = [([u0 u1 u2],)]This error does not happen if the loss function loops over the initial conditions solving them one at a time as simple vectors. Thus this is a broadcasting problem.
Regression in the code you can see each component of the process is validated. that is the model, the ode solver, and the loss functions all work fine in this batch oriented input. The error happens in the gradient step. It only happens on the second model, not the first.
Output
Versions