High memory footprint for big dataframes in CausalForest model

py-why / EconML

ALICE (Automated Learning and Intelligence for Causation and Economics) is a Microsoft Research project aimed at applying Artificial Intelligence concepts to economic decision making. One of its goals is to build a toolkit that combines state-of-the-art machine learning techniques with econometrics in order to bring automation to complex causal inference problems. To date, the ALICE Python SDK (econml) implements orthogonal machine learning algorithms such as the double machine learning work of Chernozhukov et al. This toolkit is designed to measure the causal effect of some treatment variable(s) t on an outcome variable y, controlling for a set of features x.

https://www.microsoft.com/en-us/research/project/alice/

Other

3.64k stars 687 forks source link

High memory footprint for big dataframes in CausalForest model #883

Open gabrieldaiha opened 1 month ago

gabrieldaiha commented 1 month ago

Hi!

I'm facing multiple memory issues when using CausalForestDML class with big dataframes.

In both the fit method and the effect / effect inference methods, larger dataframes (even not so large, ex: 2M rows / 20 columns) tend to have an exponential memory occupation (~180GB ram).

We could start an investigation inside the code to try to minimize the memory footprint / optimize both methods below:

predict_point_and_var, from https://github.com/py-why/EconML/blob/db1e254f1e5cd476b917cfe1c2984db4c264bb66/econml/grf/_base_grf.py#L650, used in both effect / effect_inference methods
fit, from https://github.com/py-why/EconML/blob/db1e254f1e5cd476b917cfe1c2984db4c264bb66/econml/grf/_base_grf.py#L192

My first investigations are tending to indicate that problem is originating from these methods above.

Any clues? @kbattocchi

kbattocchi commented 1 month ago

One simple comparison that would be useful is how the memory consumption of standard sklearn RandomForests compares on dataframes of the same size, since much of the EconML tree code was forked from sklearn (version 0.24, I believe).

Although 180GB does seem excessive, I don't think it is really exponential - if your input has 40M floating point values, the raw data for that alone is 320MB, so this is ~560 times the size of your dataset. Certainly if we can easily optimize things to bring this down we should, but it's not even quadratic in the number of elements.

You mention that memory is high for both fit and effect: do you mean that while running those methods memory usage spikes but then comes back down to a more reasonable amount when the method calls complete?

gdaiha commented 1 month ago

You mention that memory is high for both fit and effect: do you mean that while running those methods memory usage spikes but then comes back down to a more reasonable amount when the method calls complete?

Yes, memory usage spikes, but then comes back down.

I'm trying to investigate better inside fit, but in predict_point_and_var, I identified that the spike of memory comes after the second Parallel call inside var condition, so I think memory spike is probably origined on these rows:

https://github.com/py-why/EconML/blob/db1e254f1e5cd476b917cfe1c2984db4c264bb66/econml/grf/_base_grf.py#L703-L763

gdaiha commented 1 month ago

Another important detail. I was using a treatment dataframe with featurizer, making me have 6 columns in T. I was inspecting code, and, in many steps, they use a cross product of T over T. I think this is contributing for this memory spike too.