How do I release Flux.Tracker allocated memory?

New to Julia
, ,
1

I’m diving into Flux and wanted to try out the reverse based autodiff via Flux.Tracker. I have written the following function to do a small test which seems to work and is reasonable fast. However, if I run this function multiple times it keeps increasing memory usage until I get an OutOfMemory error. How should I deallocated the memory between multiple runs of this function. From my naive impression this should happen automatically since nothing is holding on to the memory after the function is evaluated but obviously I’m wrong. :slight_smile:

I’m on Julia 1.0 on Ubuntu 18.04 on a Linux Surface Pro 4.

using Flux
using Flux.Tracker
using BenchmarkTools

function fluxderivtest(nhidden=100_000, ninput=1_000)
    W = param(rand(nhidden, ninput) .- 0.5)
    b = param(rand(nhidden) .- 0.5)
    x = rand(ninput)
    f(x) = tanh(sum(W*x .+ b))
    g = Tracker.gradient(() -> f(x), Params([W, b]))
    g[W], g[b]
end

@btime fluxderivtest()
@btime fluxderivtest()

ERROR: OutOfMemoryError()
Stacktrace:
 [1] Type at ./boot.jl:396 [inlined]
 [2] Type at ./boot.jl:404 [inlined]
 [3] Type at ./boot.jl:411 [inlined]
 [4] similar at ./abstractarray.jl:618 [inlined]
 [5] similar at ./abstractarray.jl:617 [inlined]
 [6] similar at ./broadcast.jl:193 [inlined]
 [7] copy at ./broadcast.jl:744 [inlined]
 [8] materialize at ./broadcast.jl:724 [inlined]
 [9] param(::Array{Float64,2}) at /home/michael/.julia/packages/Flux/UHjNa/src/tracker/Tracker.jl:105
 [10] fluxderivtest(::Int64, ::Int64) at ./REPL[55]:2 (repeats 2 times)
 [11] ##core#364() at /home/michael/.julia/packages/BenchmarkTools/dtwnm/src/execution.jl:293
 [12] ##sample#365(::BenchmarkTools.Parameters) at /home/michael/.julia/packages/BenchmarkTools/dtwnm/src/execution.jl:301
 [13] sample at /home/michael/.julia/packages/BenchmarkTools/dtwnm/src/execution.jl:316 [inlined]
 [14] #_lineartrial#20(::Int64, ::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}, ::Function, ::BenchmarkTools.Benchmark{Symbol("##benchmark#363")}, ::BenchmarkTools.Parameters) at /home/michael/.julia/packages/BenchmarkTools/dtwnm/src/execution.jl:71
 [15] _lineartrial(::BenchmarkTools.Benchmark{Symbol("##benchmark#363")}, ::BenchmarkTools.Parameters) at /home/michael/.julia/packages/BenchmarkTools/dtwnm/src/execution.jl:63
 [16] #invokelatest#1 at ./essentials.jl:686 [inlined]
 [17] invokelatest at ./essentials.jl:685 [inlined]
 [18] #lineartrial#17 at /home/michael/.julia/packages/BenchmarkTools/dtwnm/src/execution.jl:33 [inlined]
 [19] lineartrial at /home/michael/.julia/packages/BenchmarkTools/dtwnm/src/execution.jl:33 [inlined]
 [20] #tune!#23(::Bool, ::String, ::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}, ::Function, ::BenchmarkTools.Benchmark{Symbol("##benchmark#363")}, ::BenchmarkTools.Parameters) at /home/michael/.julia/packages/BenchmarkTools/dtwnm/src/execution.jl:135
 [21] tune! at /home/michael/.julia/packages/BenchmarkTools/dtwnm/src/execution.jl:134 [inlined] (repeats 2 times)
 [22] top-level scope at /home/michael/.julia/packages/BenchmarkTools/dtwnm/src/execution.jl:388

Most likely it’s a stupid error but I cannot see it.

2

Works, for me. Memory increases the first times it is executed but then the GC kicks in so OOM is never reached. Explicitly running GC.gc() a couple of times (3x) also reclaims all the memory. So it seems the GC correctly tracks the memory but is failing to run soon enough for you not to go OOM?

3

That could very well be since I assume that @btime calls the function twice. Once to compile and once to evaluate as standard. Maybe these two calls are done “too” quickly for the GC to keep up? Also if I run @time instead this does not seem to happen. Is it then best practice to run the GC manually when you have a hunch you might run out of memory?

4

The GC can fire in a lot of places so this shouldn’t be the case.

No, AFAIU this shouldn’t happen. If it is reproducible for you, perhaps you can open an issue. @btime works ok on my computer though.

5

Good to know.

Yeah it is indeed reproducible. Is Flux the right place to open an issue or would you think this belongs somewhere else?

6

Doesn’t seem like a Flux issue unless Flux does things in a way that it needs to handle memory itself.

7

Note that if any of the memory consumed is externally allocated and managed by finalizers, julia GC won’t be able to feel and react to it.

8

I’ve opened an issue here: Out of memory error caused by calling @btime on on a medium sized model · Issue #29085 · JuliaLang/julia · GitHub

9

I don’t know enough about this to say if this is the case, but to me there shouldn’t be any external allocations here since all I’m doing is calling the function twice( i.e. four times) with @btime. But I could be wrong. Is there anything you would do differently in this small case I have here?

10

@btime is not calling the function twice and that’s irrelevant.

I mean external libraries used by Flux.

11

That’s why I asked if this belongs in Flux or in Julia as an issue. Either way I’ve posted it in julia as an issue. :+1:

Regarding @btime I don’t understand that at all as in my mind that could definitely be part of the problem. But then again I don’t have to understand it. :slight_smile: I’m sure you know what you’re talking about. A few observations I did that may be of help:

fluxderivtest()
fluxderivtest()
fluxderivtest()
fluxderivtest()
fluxderivtest()

yields no problem at all. While

@btime fluxderivtest()
@btime fluxderivtest()

yields an Out of memory error. At the same time

@time fluxderivtest()
@time fluxderivtest()
@time fluxderivtest()
@time fluxderivtest()
@time fluxderivtest()

yields no issue at all.