chiragnagpal / deep_cox_mixtures

Code for the paper "Deep Cox Mixtures for Survival Regression", Machine Learning for Healthcare Conference 2021

Deep Cox Mixtures

Chirag Nagpal^1,2 Steve Yadlowsky¹, Negar Rostamzadeh¹ and Katherine Heller¹

¹Google Brain Team, ²Carnegie Mellon University

---

❗⚠️❗⚠️❗⚠️❗⚠️❗ IMPORTANT NOTE ❗⚠️❗⚠️❗⚠️❗⚠️❗⚠️❗

Deep Cox Mixtures now has a more stable pytorch implementation here:

https://autonlab.github.io/auton-survival/models/dcm/

tensorflow version is no longer supported. Please use the version above. The repository is kept for legacy purposes.

---

This repository contains code for the MLHC 2021 paper:

Deep Cox Mixtures for Survival Regression

Installation

To download and run Deep Cox Mixtures:

foo@bar:~$ git clone https://github.com/chiragnagpal/deep_cox_mixtures.git
foo@bar:~$ cd deep_cox_mixtures
foo@bar:~$ pip install -r requirements.txt

Usage

To run DCM on a standard survival analysis dataset like SUPPORT, please see the following example notebook:

1. Deep Cox Mixtures on the SUPPORT Dataset

To run the original experiments from the paper, please use:

from dcm import deep_cox_mixture
results = deep_cox_mixture.experiment(dataset='SUPPORT', prot_att='race', groups=('white', 'other'))
deep_cox_mixture.display_results(results)

Requirements

dcm depends on tensorflow2 and scikit-survival,

Running baseline models for comparison requires lifelines, pycox and dsm

Citing

Please cite using the following bib-entry:

@article{nagpal2021dcm,
  title={Deep Cox mixtures for survival regression},
  author={Nagpal, Chirag and Yadlowsky, Steve and Rostamzadeh, Negar and Heller, Katherine},
  journal={Machine Learning for Healthcare Conference},
  year={2021}
  organization={PMLR}
}