search

MilesCranmer / SymbolicRegression.jl

Distributed High-Performance Symbolic Regression in Julia
https://ai.damtp.cam.ac.uk/symbolicregression/
Apache License 2.0
644 stars 86 forks source link

Large memory consumption of EquationSearch #9

Closed baggepinnen closed 3 years ago

baggepinnen commented 3 years ago

I ran with the following options

using Distributed
nprocs() == 1 && (procs = addprocs(4))
@everywhere using SymbolicRegression
options = SymbolicRegression.Options(
    binary_operators=(+, *, /, -),
    unary_operators=(cos, sin, sqrt),
    npopulations=40
)
niterations = 15

on a problem where X had size 3x18, and before it completed the julia session had allocated some 12+GB of memory and crashed both itself and my browser. I am not familiar with the internals of this package, but are there perhaps excessive amounts of data stored in an internal data structure that are not really needed that can be discarded as the optimization progresses?

MilesCranmer commented 3 years ago

Sorry that it used up so much memory! I haven’t experienced this before, not sure what is causing it.

The internal representation of equations is super lightweight: https://github.com/MilesCranmer/SymbolicRegression.jl/blob/cec1d077655f2903283de0253918882daa8f9dc3/src/Equation.jl#L4-L11, containing enums which index feature, operator, and degree. So I don’t think the memory issue is coming from storing those.

I note that npopulations=40 means you are running 40 equation searches asynchronously, with IIRC 300 equations each. (The processes share their best equations at the end of each iteration, which “migrate” between processes). I wonder if setting npopulations=4 (equal to your number of procs), and npop to 3000, would give better memory usage?

baggepinnen commented 3 years ago

I was thinking, is the compiler involved for each equation, or how are equations evaluated? If the compiler compiles each equation, I can imagine that the compiled code would slowly fill up the memory. I' tried your suggested configuration and it indeed appears to be a much better choice.

MilesCranmer commented 3 years ago

There’s no compilation when evaluating equations; this is essentially why I need my own data structure. If I used an Expr it would take about 10,000x longer to allocate and evaluate equations—-I tried this initially but the improved flexibility wasn’t worth the compute price. I also found a similar performance hit when I used SymbolicUtils.jl’s internal type by default.

Since the struct holding equations is an enum, and I fix operators for each new search, I can basically just recursively traverse the tree with a branching statement over the operators: https://github.com/MilesCranmer/SymbolicRegression.jl/blob/174af5b7e42fcf9dbe1c9c95be0689093a03aa8c/src/EvaluateEquation.jl#L22-L57. By having the operators all fixed inside this evaluator, the compiler can really speed things up it seems! I also try to minimize array allocation in this loop as you can seen with the cumulator variable.

At one point i tried to use a stack to traverse and evaluate equations instead of recursion, which I think would be even better for memory usage, but I couldn’t seem to get this working yet.

One other idea: I wonder if it could even improve performance more to compile specialized functions for each (binary, unary, unary) subtree...? There aren’t too many possibilities so it might make sense to try to fuse them!

MilesCranmer commented 3 years ago

Want to try reproducing this on v0.4.1? The serial computation is much easier for debugging these sorts of things I think (and for benchmarking!).

Cheers, Miles

MilesCranmer commented 3 years ago

Closing for now as I think it should be fixed; let me know if this issue comes up again!

Cheers, Miles