reduce SurvivalTree.predict's memory use

[cpoerschke]

Checklist

• pytest passes
• tests are included
• code is well formatted
• documentation renders correctly

What does this implement/fix? Explain your changes

issue encountered: Passing relatively lots of technical data samples to SurvivalTree.predict results in "kernel died" in a notebook environment.

analysis

The pred = self.tree_.predict(X) call in SurvivalTree.predict_cumulative_hazard_function returns an array of (n_samples, self.event_times_, 2) shape which is then returned as a (n_samples, self.event_times_) shape and then SurvivalTree.predict does a .sum(1) on it, resulting in an overall (n_samples) shape return value. The n_samples X self.event_times_ shapes use a lot of memory for larger samples especially if the number of event times is a proportion of samples.

proposed solution

• Do the self.tree_.predict call on a one-row-at-a-time basis, requiring one (1, self.event_times_, 2) shape for each iteration.
• Have SurvivalTree.predict_cumulative_hazard_function do the sum-ing for each row, removing the need for a (n_samples, self.event_times_) shape array.

testing

peak memory usage stats inspected before and after the change using memray run ./pr6-demo.py profiling.

code snippet:

#!/usr/bin/env python

import numpy as np
import sksurv.tree
import sksurv.util

n = 12345

verbose = (n <= 10)

feature1 = np.arange(0, n, 1)
feature2 = n - feature1

times = np.arange(0,n) + 1
events = (times < times[len(times)//2])

X = np.vstack((feature1, feature2)).T
y = sksurv.util.Surv.from_arrays(time=times, event=events)

if verbose:
    print(f"X={X}\ny={y}")
else:
    print(f"X.shape={X.shape}\ny.shape={y.shape}")

st = sksurv.tree.SurvivalTree(max_leaf_nodes=100)
print(st.fit(X, y))

risk_scores = st.predict(X)

if verbose:
    print(f"risk_scores={risk_scores}")
else:
    print(f"risk_scores.shape={risk_scores.shape}")

[sebp]

Thanks for your PR.

I agree that predict is quite heavy on memory usage. Part of it has to do with sklearn's Tree-related code, where predictions are returned via the split criterion's node_value method. It is currently implemented to return the full survival function and cumulative hazard function for each sample, disregarding whether predict, predict_survival_function, or predict_cumulative_hazard_function has been called. I could imagine adding a low-memory option that disables computing survival and cumulative hazard function (CHF) and just returns the event counts (sum over CHF, i.e. what predict returns).

[cpoerschke]

... Part of it has to do with ... I could imagine adding a low-memory option that disables computing survival and cumulative hazard function (CHF) and just returns the event counts (sum over CHF, i.e. what predict returns).

Thanks for the context and quick feedback!

I've added a low_memory=False option in the latest commit, though I guess its current use does not disable the computations as such. And of course for any new option there should be test coverage too.

[codecov]

Codecov Report

Patch coverage: 100.00% and no project coverage change.

Comparison is base (bd2e240) 97.94% compared to head (9e1b344) 97.95%.

Additional details and impacted files

@@           Coverage Diff           @@
##           master     #369   +/-   ##
=======================================
  Coverage   97.94%   97.95%           
=======================================
  Files          37       37           
  Lines        3361     3376   +15     
  Branches      509      511    +2     
=======================================
+ Hits         3292     3307   +15     
  Misses         33       33           
  Partials       36       36           

Impacted Files	Coverage Δ
sksurv/ensemble/forest.py	100.00% <100.00%> (ø)
sksurv/tree/tree.py	95.71% <100.00%> (+0.43%)	⬆️

☔ View full report in Codecov by Sentry.
📢 Do you have feedback about the report comment? Let us know in this issue.

[cpoerschke]

... Part of it has to do with sklearn's Tree-related code, where predictions are returned via the split criterion's node_value method. It is currently implemented to return the full survival function and cumulative hazard function for each sample, disregarding whether predict, predict_survival_function, or predict_cumulative_hazard_function has been called. I could imagine adding a low-memory option that disables computing survival and cumulative hazard function (CHF) and just returns the event counts (sum over CHF, i.e. what predict returns).

So I haven't worked with .pxd and .pyx code before but from a little bit of code reading just now ... is the idea conceptually that in low-memory mode:

• https://github.com/sebp/scikit-survival/blob/v0.20.0/sksurv/tree/tree.py#L227-L229 would request fewer outputs and classes from LogRankCriterion and Tree
• https://github.com/sebp/scikit-survival/blob/v0.20.0/sksurv/tree/_criterion.pyx#L305-L327 would fill the reduced size dest with the event counts
• predict_cumulative_hazard_function and predict_survival_function cannot be used in low-memory mode
• some but not all of the above and/or something else?

[cpoerschke]

Added work-in-progress notes inline.

[sebp]

Actually, ties should not cause any problems.

[cpoerschke]

Thanks @sebp for the feedback! I think I've addressed all points and tests pass locally. Please let me know if there's any additional suggestions. Thank you.