Merge pull request #592 from SciML/dg/zeros

use Flux.zeros32

Author: ChrisRackauckas

src/fast_layers.jl

@@ -24,7 +24,7 @@ initial_params(c::FastChain) = vcat(initial_params.(c.layers)...)
 
 """
 FastDense(in,out,activation=identity;
-          bias = true ,initW = Flux.glorot_uniform, initb = Flux.zeros)
+          bias = true ,initW = Flux.glorot_uniform, initb = Flux.zeros32)
 
 A Dense layer `activation.(W*x + b)` with input size `in` and output size `out`.
 The `activation` function defaults to `identity`, meaning the layer is an affine
@@ -41,7 +41,7 @@ struct FastDense{F,F2} <: FastLayer
   initial_params::F2
   bias::Bool
   function FastDense(in::Integer, out::Integer, σ = identity;
-                 bias = true, initW = Flux.glorot_uniform, initb = Flux.zeros)
+                 bias = true, initW = Flux.glorot_uniform, initb = Flux.zeros32)
     temp = ((bias == false) ? vcat(vec(initW(out, in))) : vcat(vec(initW(out, in)),initb(out)))
     initial_params() = temp
     new{typeof(σ),typeof(initial_params)}(out,in,σ,initial_params,bias)
@@ -108,7 +108,7 @@ initial_params(f::FastDense) = f.initial_params()
 
 """
 StaticDense(in,out,activation=identity;
-          initW = Flux.glorot_uniform, initb = Flux.zeros)
+          initW = Flux.glorot_uniform, initb = Flux.zeros32)
 
 A Dense layer `activation.(W*x + b)` with input size `in` and output size `out`.
 The `activation` function defaults to `identity`, meaning the layer is an affine
@@ -124,7 +124,7 @@ struct StaticDense{out,in,bias,F,F2} <: FastLayer
   σ::F
   initial_params::F2
   function StaticDense(in::Integer, out::Integer, σ = identity;
-                 bias::Bool = true, initW = Flux.glorot_uniform, initb = Flux.zeros)
+                 bias::Bool = true, initW = Flux.glorot_uniform, initb = Flux.zeros32)
     temp = ((bias == true ) ? vcat(vec(initW(out, in)),initb(out)) : vcat(vec(initW(out, in))) )
     initial_params() = temp
     new{out,in,bias,typeof(σ),typeof(initial_params)}(σ,initial_params)

test/newton_neural_ode.jl

@@ -21,7 +21,7 @@ NN = Chain(Dense(n, 5n, tanh),
 @info "ROCK4"
 nODE = NeuralODE(NN, tspan, ROCK4(), reltol=1e-4, saveat=[tspan[end]])
 
-loss_function(θ) = Flux.mse(y, nODE(x, θ))
+loss_function(θ) = Flux.mse(y, nODE(x, θ)[end])
 l1 = loss_function(nODE.p)
 
 res = DiffEqFlux.sciml_train(loss_function, nODE.p, NewtonTrustRegion(), GalacticOptim.AutoZygote(), maxiters = 100, cb=cb)
@@ -35,7 +35,7 @@ NN = FastChain(FastDense(n, 5n, tanh),
 @info "ROCK2"
 nODE = NeuralODE(NN, tspan, ROCK2(), reltol=1e-4, saveat=[tspan[end]])
 
-loss_function(θ) = Flux.mse(y, nODE(x, θ))
+loss_function(θ) = Flux.mse(y, nODE(x, θ)[end])
 l1 = loss_function(nODE.p)
 optfunc = GalacticOptim.OptimizationFunction((x, p) -> loss_function(x), GalacticOptim.AutoZygote())
 optprob = GalacticOptim.OptimizationProblem(optfunc, nODE.p,)