How do I pass Optim.Options(...) to DiffEqFlux.sciml_train(...)

Matthew_Overlin · May 8, 2020, 6:46pm

I am using DiffEqFlux to train neural ODEs. Following the example at this link, I am using the function DiffEqFlux.scil_train to train the neural network. I’d like to pass/customize some things with the optimizer. For example, I might like to specify f_tol, f_abstol, show_trace. So, I might send “Optim.Options(f_tol=1.0e-3, f_abstol=1.0e-6, show_trace=true)” as an argument to the “sciml_train” function. What exactly would the syntax like if I want to do this? I’m trying to get a better fit with my data, so I want to pass in a low tolerance for f, so that the training completes with a lower loss than the loss it returns now.
Thanks,
Matt

ChrisRackauckas · May 10, 2020, 6:22am

Any extra keyword arguments get thrown down to Optim:

github.com

SciML/DiffEqFlux.jl/blob/master/src/train.jl#L271


      
          function sciml_train(loss, θ, opt=OptimizationPolyalgorithms.PolyOpt(), adtype=nothing, args...;

                               lower_bounds=nothing, upper_bounds=nothing, cb = nothing,

                               callback = (args...) -> (false),

                               maxiters=nothing, kwargs...)

          

              @warn "sciml_train is being deprecated in favor of direct usage of Optimization.jl. Please consult the Optimization.jl documentation for more details. Optimization.jl's PolyOpt solver is the polyalgorithm of sciml_train"

          

              if adtype === nothing

                  if length(θ) < 50

                      fdtime = try

                          ForwardDiff.gradient(x -> first(loss(x)), θ)

                          @elapsed ForwardDiff.gradient(x -> first(loss(x)), θ)

                      catch

                          Inf

                      end

                      zytime = try

                          Zygote.gradient(x -> first(loss(x)), θ)

                          @elapsed Zygote.gradient(x -> first(loss(x)), θ)

                      catch

                          Inf

                      end

          

                      if fdtime == zytime == Inf

                          @warn "AD methods failed, using numerical differentiation. To debug, try ForwardDiff.gradient(loss, θ) or Zygote.gradient(loss, θ)"

                          adtype = Optimization.AutoFiniteDiff()

                      elseif fdtime < zytime

                          adtype = Optimization.AutoForwardDiff()

                      else

                          adtype = Optimization.AutoZygote()

                      end

          

                  else

                      adtype = Optimization.AutoZygote()

                  end

              end

              if !isnothing(cb)

                callback = cb

              end

          

              optf = Optimization.OptimizationFunction((x, p) -> loss(x), adtype)

              optprob = Optimization.OptimizationProblem(optf, θ; lb=lower_bounds, ub=upper_bounds, kwargs...)

              if maxiters !== nothing

                  Optimization.solve(optprob, opt, args...; maxiters, callback = callback, kwargs...)

              else

                  Optimization.solve(optprob, opt, args...; callback = callback, kwargs...)

              end

          end

so you can just add show_trace = true, f_tol = 1e-3, etc.

Note that sciml_train ended up becoming a pretty big setup due to wrapping all of the Julia nonlinear optimization libraries we could find and throwing every reverse mode AD on it… so that’s all getting refactored out to https://github.com/SciML/GalacticOptim.jl fairly soon. After it gets moved out and the API gets cleaned up for its role as a differentiable wrapper over optimization, we’ll make sure all of this gets proper documentation.