nathanaelbosch / ProbNumDiffEq.jl

Probabilistic Numerical Differential Equation solvers via Bayesian filtering and smoothing
MIT License
119 stars 16 forks source link

Memory leaking when computing gradients #249

Closed jnsbck closed 9 months ago

jnsbck commented 1 year ago

Hey,

I encountered some weird issue, where memory seamingly leaks during ForwardDiff.gradient(ek1_loss). I was able to boil it down to the snippet below:

using ComponentArrays
using ForwardDiff
using ProbNumDiffEq
using BenchmarkTools

function get_lv_prob(tspan=(0.0, 10.0))
    function lotka_volterra(du, u, p, t)
        du[1] = p[1] * u[1] - p[2] * u[1] * u[2]
        du[2] = -p[3] * u[2] + p[4] * u[1] * u[2]
    end
    p = [1.5, 1.0, 3.0, 1.0]
    u0 = [1.0, 1.0]
    prob = ODEProblem(lotka_volterra, u0, tspan, p)
    return prob
end

prob = get_lv_prob((0.0, 17.0)) # no leak
prob = get_lv_prob((0.0, 20.0)) # slowly starts leaking
prob = get_lv_prob((0.0, 100.0)) # continues to leak

function ek1_loss(p, prob)
    # leaks for both EK0 and EK1
    sol = solve(prob, EK1(smooth=false), dense=false, p=p) # leak independet of order or dt
    l = sum(abs2, Array(sol)) / length(sol)
    return l
end

function test(; prob=prob)
    p = ComponentArray(prob.p)
    total_mem = Sys.total_memory() / 2^20
    i = 0
    while true
        i += 1
        ForwardDiff.gradient(p -> ek1_loss(p, prob), p)
        # GC.gc() # avoid memory leak (HOTFIX)
        free_mem = Sys.free_memory() / 2^20
        mem_usage = round((total_mem - free_mem) / total_mem * 100)
        print("\r Iteration $i, memory usage: $mem_usage%")
    end
end

test()

Lemme know in case you need any additional info. Thanks for looking into this.

nathanaelbosch commented 1 year ago

Thanks for opening this issue! I don't really have an idea where this is coming from, I am not familiar enough with the Julia internals and garbage collection. But I am surprised that manually calling GC.gc() fixes this. Either way, this should not be happening, and I will try to find out more to get this fixed.

nathanaelbosch commented 9 months ago

I just ran this code again and the memory usage des not seem to increase, so I will close this issue for now. But if I somehow missed something and this still persists, please do reopen the issue.

jnsbck commented 9 months ago

Just tested again and for me the issue still persists. I am on the lastest package versions

  [6e4b80f9] BenchmarkTools v1.4.0
  [b0b7db55] ComponentArrays v0.15.8
  [f6369f11] ForwardDiff v0.10.36
  [bf3e78b0] ProbNumDiffEq v0.15.0

EDIT: cannot re-open since you closed it and I am not a collaborator on this repo.

nathanaelbosch commented 9 months ago

Ok so this seems to be an issue with julia 1.9, but with 1.10 I can't reproduce this.