# GPU simulation time of neural ODE program

**URL:** <https://discourse.julialang.org/t/gpu-simulation-time-of-neural-ode-program/27416>\
**Category:** Machine Learning\
**Tags:** diffeq, flux\
**Created:** [August 11, 2019, 1:01pm UTC](https://discourse.julialang.org/t/gpu-simulation-time-of-neural-ode-program/27416 "2019-08-11T13:01:31Z")\
**Posts on this page:** 9\
**Page:** 1

<div class="post-metadata">

**Author:** ![JunichiFukui](https://avatars.discourse-cdn.com/v4/letter/j/ad7895/32.png) [@JunichiFukui](https://discourse.julialang.org/u/JunichiFukui)\
**Post date:** [August 11, 2019, 1:01pm UTC](https://discourse.julialang.org/t/gpu-simulation-time-of-neural-ode-program/27416/1 "2019-08-11T13:01:31Z")

</div>

Hello,

I’m trying to run the following neural ODE program using GPU. If I run the code using GPU (NVIDIA GeForce GTX 1050), the simulation time of `Flux.train!(loss_n_ode, ps, data, opt, cb = cb)` part was much longer than simulation using only CPU as shown below.

GPU: 362.607972 seconds (719.08 M allocations: 24.526 GiB, 3.86% gc time)  
CPU: 35.925501 seconds (150.99 M allocations: 13.473 GiB, 9.14% gc time)

Could you tell me how to fix the issue to make simulation speed of GPU code faster?

```julia
using Flux, DiffEqFlux, DifferentialEquations, Plots, CSV, CuArrays

# Read CSV file
flamedata = CSV.read("./results.csv")
flamedata2 = flamedata[2:2:6002,:]

u0 = Float32[0.; 0.11189834407236525; 0.8881016559276348; 0.]|>gpu

datasize = 3001
tspan = (0.0f0,0.0003f0)|>gpu

t = range(tspan[1],tspan[2],length=datasize)|>gpu

ode_data2 = Matrix(flamedata2[[:1,:4,:7,:9]])
ode_data2 = transpose(ode_data2)
ode_data2 = convert(Array{Float32}, ode_data2)
ode_data = ode_data2|>gpu
ode_data[1,:] = tanh.(ode_data[1,:]*100)|>gpu

dudt = Chain(
        Dense(4,32,swish),
	Dense(32,16,swish),
	Dense(16,8,swish),
        Dense(8,4)
	)|>gpu

ps = Flux.params(dudt)|>gpu
n_ode = x->neural_ode(dudt,gpu(x),gpu(tspan),AutoTsit5(Rodas5(autodiff=false)),saveat=t,dtmin=1.0E-14,maxiters=1e10,reltol=1e-7,abstol=1e-9)

function predict_n_ode()
  n_ode(u0)
end
loss_n_ode() = sum(abs2,ode_data .- predict_n_ode())

data = Iterators.repeated((), 5)
opt = ADAM(0.1, (0.9, 0.999))
cb = function () #callback function to observe training
  display(loss_n_ode())
  # plot current prediction against data
  cur_pred = Flux.data(predict_n_ode())
  pl1 = plot(t,ode_data[1,:],label="data1",lw=2)
  plot!(pl1,t,cur_pred[1,:],label="prediction1",lw=2)
  plot!(pl1,t,ode_data[2,:],label="data2",lw=2)
  plot!(pl1,t,cur_pred[2,:],label="prediction2",lw=2)
  plot!(pl1,t,ode_data[3,:],label="data3",lw=2)
  plot!(pl1,t,cur_pred[3,:],label="prediction3",lw=2)
  plot!(pl1,t,ode_data[4,:],label="data4",lw=2)
  plot!(pl1,t,cur_pred[4,:],label="prediction4",lw=2)
  gui(plot(pl1))
end

# Display the ODE with the initial parameter values.
cb()

@time Flux.train!(loss_n_ode, ps, data, opt, cb = cb)

```

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [August 11, 2019, 1:15pm UTC](https://discourse.julialang.org/t/gpu-simulation-time-of-neural-ode-program/27416/2 "2019-08-11T13:15:04Z")

</div>

> [@JunichiFukui](#):
>
> Could you tell me how to fix the issue to make simulation speed of GPU code faster?

The problem is `Rodas5` here. If you just do `Tsit5` you’ll be fine, for now (or use a ROCK method). Jacobian construction without autodiff is slow on the GPU. We are building a new forward-mode autodiff precisely for this reason, and I think we will have something to show for this by the end of next week.

---

<div class="post-metadata">

**Author:** ![JunichiFukui](https://avatars.discourse-cdn.com/v4/letter/j/ad7895/32.png) [@JunichiFukui](https://discourse.julialang.org/u/JunichiFukui)\
**Post date:** [August 11, 2019, 2:04pm UTC](https://discourse.julialang.org/t/gpu-simulation-time-of-neural-ode-program/27416/3 "2019-08-11T14:04:48Z")

</div>

Dear Chris,

Thank you for your prompt reply.

I tried two types of ODE algorithms as shown below. However, simulation time for the runs using GPU was still slower than CPU only.

`n_ode = x->neural_ode(dudt,gpu(x),gpu(tspan),Tsit5(),saveat=t,dtmin=1.0E-14,maxiters=1e10,reltol=1e-7,abstol=1e-9)`  
286.258497 seconds (699.98 M allocations: 23.458 GiB, 3.58% gc time)

`n_ode = x->neural_ode(dudt,gpu(x),gpu(tspan),AutoTsit5(Tsit5()),saveat=t,dtmin=1.0E-14,maxiters=1e10,reltol=1e-7,abstol=1e-9)`  
302.667069 seconds (700.02 M allocations: 23.523 GiB, 3.62% gc time)

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [August 11, 2019, 2:17pm UTC](https://discourse.julialang.org/t/gpu-simulation-time-of-neural-ode-program/27416/4 "2019-08-11T14:17:57Z")

</div>

For your case, the matrices are still too small. A 32x32 matmul isn’t large enough to overcome the kernel call cost. The GPU will just help scale larger layers.

---

<div class="post-metadata">

**Author:** ![JunichiFukui](https://avatars.discourse-cdn.com/v4/letter/j/ad7895/32.png) [@JunichiFukui](https://discourse.julialang.org/u/JunichiFukui)\
**Post date:** [August 14, 2019, 9:40pm UTC](https://discourse.julialang.org/t/gpu-simulation-time-of-neural-ode-program/27416/5 "2019-08-14T21:40:10Z")

</div>

Dear Chris,

Difference of simulation time between CPU and GPU became smaller than before if number of parameters in neural network layers increased as shown below.

```julia
dudt = Chain(
        Dense(4,256,swish),
	Dense(256,128,swish),
	Dense(128,64,swish),
        Dense(64,4)
	)|>gpu

```

```
GPU: 279.341810 seconds (713.16 M allocations: 23.783 GiB, 4.39% gc time)
CPU: 172.805359 seconds (137.03 M allocations: 361.775 GiB, 21.18% gc time)

```

Simulation speed of GPU is still slower than that of CPU, since GPU in my laptop PC (NVIDIA GeForce GTX 1050) is not so powerful. Therefore, I will use CPU if I run neural ODE code on my laptop PC. I hope MPI or OpenMP support in future version of DiffEqFlux.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [August 14, 2019, 10:30pm UTC](https://discourse.julialang.org/t/gpu-simulation-time-of-neural-ode-program/27416/6 "2019-08-14T22:30:36Z")

</div>

What does your computer show at that size?

```julia
julia> using BenchmarkTools, CuArrays

julia> A = cu(rand(4,256)); B = cu(rand(256,128));

julia> @btime A*B
  3.406 μs (19 allocations: 592 bytes)
4×128 CuArray{Float32,2}:
 70.0315 66.4814 71.9346 71.5927 69.5655 65.5431 … 70.0417 63.0131 70.5124 67.8303 68.7845 65.3313
 59.1719 59.0919 63.0657 59.825 61.1214 57.4047 56.0934 52.0603 59.7451 61.2429 58.7901 57.6035
 61.2603 60.0343 65.5527 65.2218 60.3737 59.5472 59.8305 53.4624 64.5687 62.1284 63.2896 57.3726
 67.5751 64.2045 71.7379 71.1343 68.4805 63.8018 66.5502 60.405 68.8187 67.7905 65.8484 64.2003

julia> A = rand(4,256); B = rand(256,128);

julia> @btime A*B
  31.357 μs (1 allocation: 4.13 KiB)
4×128 Array{Float64,2}:
 62.8857 63.1846 60.4739 58.7771 59.831 62.9932 … 57.5166 59.5414 60.4711 59.8159 63.8812 68.3361
 65.5783 65.8799 62.7385 63.8594 63.1873 65.7564 62.3821 61.9301 64.7102 60.7978 62.7974 67.8307
 65.8959 68.8122 61.7253 65.4755 64.1282 66.0544 64.3825 63.6596 67.8998 62.6193 67.2302 71.0087
 65.2023 65.7254 63.5841 64.6367 62.0464 64.6357 64.951 63.2562 65.7277 62.2329 67.7674 66.351

```

> [@JunichiFukui](#):
>
> I hope MPI or OpenMP support in future version of DiffEqFlux.

You don’t need OpenMP support because all of the matmuls are already multithreaded. Check your CPU usage and you should see it using all cores automatically.

---

<div class="post-metadata">

**Author:** ![JunichiFukui](https://avatars.discourse-cdn.com/v4/letter/j/ad7895/32.png) [@JunichiFukui](https://discourse.julialang.org/u/JunichiFukui)\
**Post date:** [August 15, 2019, 12:01am UTC](https://discourse.julialang.org/t/gpu-simulation-time-of-neural-ode-program/27416/7 "2019-08-15T00:01:45Z")

</div>

Dear Chris,

Calculation time of the matmul on my computer is following.

```julia
julia> using BenchmarkTools, CuArrays

julia> A = cu(rand(4,256)); B = cu(rand(256,128));

julia> @btime A*B
  3.699 μs (19 allocations: 592 bytes)
4×128 CuArray{Float32,2}:
 63.3189 58.2083 67.3136 56.9396 63.9851 62.1672 … 62.191 62.884 65.5695 63.5641 62.2439 58.2934
 67.5529 59.881 63.778 60.0911 66.3862 61.5216 62.1403 59.062 67.0465 66.907 66.459 59.2792
 65.3215 59.3315 62.6717 54.5476 64.2242 63.1418 60.2677 61.5783 64.8999 64.8344 63.723 57.5102
 68.5587 61.1417 66.4793 59.7681 67.7273 64.8121 63.6989 61.4582 65.1183 67.98 65.8259 61.5581

julia> A = rand(4,256); B = rand(256,128);

julia> @btime A*B
  32.820 μs (1 allocation: 4.13 KiB)
4×128 Array{Float64,2}:
 66.8995 67.3647 62.8093 63.0263 64.6963 66.201 … 62.7533 63.4769 64.2497 63.7303 65.1254 64.2671
 60.0078 58.9639 60.5994 59.0173 58.324 61.5357 61.5452 63.7883 60.4779 60.7345 60.6196 61.6659
 66.8389 67.7833 66.9895 64.8955 66.6696 67.4492 65.1754 71.0193 66.63 66.6087 68.0931 67.0778
 63.7755 66.7935 66.5961 62.9272 66.9261 66.9876 65.8134 66.7046 65.4326 65.149 66.0973 71.374

```

The attached figure shows CPU usage during the neural ODE code run. My laptop PC has Intel Core i7-7700HQ (4 physical processors and 8 logical processors). It seems that all cores were not used in the simulation. My simulation environment is following.

- OS: Windows10 Pro 64 bit
- Package version:  
Julia v1.1.0  
DifferentialEquations v6.6.0  
DiffEqFlux version: 0.6.1  
Flux v0.8.3  
CuArrays v1.1.0  
 ![task_manager_window](https://global.discourse-cdn.com/julialang/original/3X/f/5/f53a9bccbfe4e13fd5365a5e24bd32326cf6b2d6.png)

---

<div class="post-metadata">

**Author:** ![maedoc](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/maedoc/32/8490_2.png) [@maedoc](https://discourse.julialang.org/u/maedoc)\
**Post date:** [August 15, 2019, 7:45am UTC](https://discourse.julialang.org/t/gpu-simulation-time-of-neural-ode-program/27416/8 "2019-08-15T07:45:17Z")

</div>

A 1050 isn’t a good contender against the 7700HQ, especially given their clock speeds, and a Dense(256, 128) layer is still fairly small for a GPU. At least you can validate it works and motivate trying it on a faster GPU.

---

<div class="post-metadata">

**Author:** ![JunichiFukui](https://avatars.discourse-cdn.com/v4/letter/j/ad7895/32.png) [@JunichiFukui](https://discourse.julialang.org/u/JunichiFukui)\
**Post date:** [August 16, 2019, 11:43pm UTC](https://discourse.julialang.org/t/gpu-simulation-time-of-neural-ode-program/27416/9 "2019-08-16T23:43:45Z")

</div>

Dear Marmaduke,

Thank you for your reply.

Actually, Dense(32,16) layers is sufficient for my case now. Therefore, I will use CPU. I confirmed that the GPU computation worked without any error for my case. If I use more large data set, I will try to use GPU.