Repository files navigation

This repository is the official TensorFlow implementation ofUncertainty Estimation in Deep Bayesian Survival Models, BHI 2023.

The proposed method is implemented based onTensorFlow Probability.

Full paper is available on IEEE Xplore:https://ieeexplore.ieee.org/document/10313466

Preprint:https://bhiconference.github.io/BHI2023/2023/pdfs/1570918354.pdf

In this work, we introduce the use of Bayesian inference techniques for survival analysis in neural networks that rely on the Cox’s proportional hazard assumption, for which we discuss a new flexible and effective architecture. We implement three architectures: a fully-deterministic neural network that acts as a baseline, a Bayesian model using variational inference and one using Monte-Carlo Dropout.

Experiments show that the Bayesian models improve predictive performance over SOTA neural networks in a test dataset with few samples (WHAS500, 500 samples) and provide comparable performance in two larger ones (SEER and SUPPORT, 4024 and 8873 samples, respectively)

To view the license for this work, visithttps://github.com/thecml/UE-BNNSurv/blob/main/LICENSE

To run the models, please refer toRequirements.txt.

Install auton-survival manually from Git:

pip install git+https://github.com/autonlab/auton-survival.git

Code was tested in virtual environment withPython 3.8,TensorFlow 2.11 andTensorFlow Probability 0.19

• Make directoriesmkdir results andmkdir models.

• Please refer topaths.py to set appropriate paths. By default, results are inresults and models inmodels

• Network configuration using best hyperparameters are found inconfigs/*

• Run the training code:

# SOTA modelspython train_sota_models.py# BNN Modelspython train_bnn_models.py

• After model training, view the results in theresults folder.

• Run the visualization notebooks:

jupyter notebook plot_survival_curves.ipynbjupyter notebook plot_survival_time.ipynb

@inproceedings{lillelund_uncertainty_2023,  author={Lillelund, Christian Marius and Magris, Martin and Pedersen, Christian Fischer},  booktitle={2023 IEEE EMBS International Conference on Biomedical and Health Informatics (BHI)},   title={Uncertainty Estimation in Deep Bayesian Survival Models},   year={2023},  pages={1-4},  doi={10.1109/BHI58575.2023.10313466}}