Skip to content

Commit

Permalink
syntax
Browse files Browse the repository at this point in the history
  • Loading branch information
@DhairyaLGandhi
DhairyaLGandhi committed Aug 25, 2021
1 parent d0217ba commit 66f2a50
Showing 1 changed file with 44 additions and 15 deletions.
59 changes: 44 additions & 15 deletions src/packages/diffeqflux.jl
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
group = addgroup!(SUITE, "DiffEqFlux")
# group = addgroup!(SUITE, "DiffEqFlux")

function diffeqflux_add_neuralode(abstol = 1f-3, reltol = 1f-3, solver = Tsit5(), batchsize = 256)
group = addgroup!(SUITE, "DiffEqFlux")
down = Chain(flatten, Dense(784, 512, tanh))
nn = Chain(Dense(512, 256, tanh),
Dense(256, 256, tanh),
Expand All @@ -18,18 +19,29 @@ function diffeqflux_add_neuralode(abstol = 1f-3, reltol = 1f-3, solver = Tsit5()

ip = rand(Float32, 784, batchsize)

group["DiffEqFlux - Forward Pass - NeuralODE with abstol $abstol, reltol $reltol, batchsize $batchsize, and solver $solver"] = b = @benchmarkable(
group["DiffEqFlux_Forward_Pass_NeuralODE_with_abstol_$(abstol)_reltol_$(reltol)_batchsize_$(batchsize)_and_solver_$(solver)"] = b = @benchmarkable(
fw(model, gip),
setup = (nn_gpu = $nn |> gpu; model = Chain($down, $nn_ode(nn_gpu), $diffeqarray_to_array, $fc); gip = $ip |> gpu),
setup = (nn_gpu = gpu($nn);
model = Chain($down,
$nn_ode(nn_gpu),
$diffeqarray_to_array,
$fc) |> gpu;
gip = gpu($ip)),
teardown = (GC.gc(); CUDA.reclaim()))

group["DiffEqFlux - Backward Pass - NeuralODE with abstol $abstol, reltol $reltol, batchsize $batchsize, and solver $solver"] = b = @benchmarkable(
group["DiffEqFlux_Backward_Pass_NeuralODE_with_abstol_$(abstol)_reltol_$(reltol)_batchsize_$(batchsize)_and_solver_$(solver)"] = b = @benchmarkable(
bw(model, gip),
setup = (nn_gpu = $nn |> gpu; model = Chain($down, $nn_ode(nn_gpu), $diffeqarray_to_array, $fc); gip = $ip |> gpu),
setup = (nn_gpu = $nn |> gpu;
model = Chain($down,
$nn_ode(nn_gpu),
$diffeqarray_to_array,
$fc) |> gpu;
gip = gpu($ip)),
teardown = (GC.gc(); CUDA.reclaim()))
end

function diffeqflux_add_neuralsde(batchsize = 16, ntrajectories = 100)
group = addgroup!(SUITE, "DiffEqFlux")
diffusion = Chain(Dense(2, 8, tanh), Dense(8, 2))
drift = Chain(Dense(2, 32, tanh), Dense(32, 32, tanh), Dense(32, 2))
nn_sde = (f, g) -> NeuralDSDE(f, g, (0.0f0, 1.0f0), SOSRI(), abstol = 1f-1, reltol = 1f-1)
Expand All @@ -41,18 +53,27 @@ function diffeqflux_add_neuralsde(batchsize = 16, ntrajectories = 100)

ip = repeat(rand(Float32, 2, batchsize), inner = (1, ntrajectories))

group["DiffEqFlux - Forward Pass - NeuralSDE with batchsize $batchsize, and ntrajectories $ntrajectories"] = b = @benchmarkable(
group["DiffEqFlux_Forward_Pass_NeuralSDE_with_batchsize_$(batchsize)_and_ntrajectories_$(ntrajectories)"] = b = @benchmarkable(
fw(model, gip),
setup = (drift_gpu = $drift |> gpu; diffusion_gpu = $diffusion; model = Chain($nn_sde(drift_gpu, diffusion_gpu), $sdesol_to_array); gip = $ip |> gpu),
setup = (drift_gpu = gpu($drift);
diffusion_gpu = gpu($diffusion);
model = Chain($nn_sde(drift_gpu, diffusion_gpu),
$sdesol_to_array);
gip = gpu($ip)),
teardown = (GC.gc(); CUDA.reclaim()))

group["DiffEqFlux - Backward Pass - NeuralSDE with batchsize $batchsize, and ntrajectories $ntrajectories"] = b = @benchmarkable(
group["DiffEqFlux_Backward_Pass_NeuralSDE_with_batchsize_$(batchsize)_and_ntrajectories_$(ntrajectories)"] = b = @benchmarkable(
bw(model, gip),
setup = (drift_gpu = $drift |> gpu; diffusion_gpu = $diffusion; model = Chain($nn_sde(drift_gpu, diffusion_gpu), $sdesol_to_array); gip = $ip |> gpu),
setup = (drift_gpu = gpu($drift);
diffusion_gpu = $diffusion;
model = Chain($nn_sde(drift_gpu, diffusion_gpu),
$sdesol_to_array) |> gpu;
gip = gpu($ip)),
teardown = (GC.gc(); CUDA.reclaim()))
end

function diffeqflux_add_ffjord(ndims = 2, batchsize = 256)
group = addgroup!(SUITE, "DiffEqFlux")
nn = Chain(Dense(ndims, ndims * 8, tanh), Dense(ndims * 8, ndims * 8, tanh), Dense(ndims * 8, ndims * 8, tanh), Dense(ndims * 8, ndims))
cnf_ffjord = f -> FFJORD(f, (0.0f0, 1.0f0), Tsit5(), monte_carlo = true)
ffjordsol_to_logpx(x) = -mean(x[1])[1]
Expand All @@ -71,18 +92,26 @@ function diffeqflux_add_ffjord(ndims = 2, batchsize = 256)
x_gen = solve(prob, cnf_ffjord.args...; sensealg = InterpolatingAdjoint(), cnf_ffjord.kwargs...)[1:end-1, :, end]
end

group["DiffEqFlux - Forward Pass - FFJORD with batchsize $batchsize, and ndims $ndims"] = b = @benchmarkable(
group["DiffEqFlux_Forward_Pass_FFJORD_with_batchsize_$(batchsize)_and_ndims_$(ndims)"] = b = @benchmarkable(
fw(model, gip),
setup = (nn_gpu = $nn |> gpu; model = Chain($cnf_ffjord(nn_gpu), $ffjordsol_to_logpx); gip = $ip |> gpu),
setup = (nn_gpu = gpu($nn);
model = Chain($cnf_ffjord(nn_gpu),
$ffjordsol_to_logpx) |> gpu;
gip = gpu($ip)),
teardown = (GC.gc(); CUDA.reclaim()))

group["DiffEqFlux - Backward Pass - FFJORD with batchsize $batchsize, and ndims $ndims"] = b = @benchmarkable(
group["DiffEqFlux_Backward_Pass_FFJORD_with_batchsize_$(batchsize)_and_ndims_$(ndims)"] = b = @benchmarkable(
bw(model, gip),
setup = (nn_gpu = $nn |> gpu; model = Chain($cnf_ffjord(nn_gpu), $ffjordsol_to_logpx); gip = $ip |> gpu),
setup = (nn_gpu = gpu($nn);
model = Chain($cnf_ffjord(nn_gpu),
$ffjordsol_to_logpx) |> gpu;
gip = gpu($ip)),
teardown = (GC.gc(); CUDA.reclaim()))

group["DiffEqFlux - Sampling - FFJORD with nsamples $nsamples, and ndims $ndims"] = b = @benchmarkable(
group["DiffEqFlux_Sampling_FFJORD_with_nsamples_$(nsamples)_and_ndims_$(ndims)"] = b = @benchmarkable(
fw(sampler, model),
setup = (nn_gpu = $nn |> gpu; model = $cnf_ffjord(nn_gpu); sampler = $sample_from_learned_model),
setup = (nn_gpu = gpu($nn);
model = gpu($cnf_ffjord(nn_gpu));
sampler = gpu($sample_from_learned_model)),
teardown = (GC.gc(); CUDA.reclaim()))
end

0 comments on commit 66f2a50

Please sign in to comment.