checkmate
is a package to compute schedules for rematerializing tensors in DFGraphs (tensor dataflow graphs).
pip install https://github.com/parasj/checkmate/archive/master.zip#egg=checkmate
If you use Checkmate in your work, please cite us with:
@article{jain2019checkmate,
title={Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization},
author={Jain, Paras and Jain, Ajay and Nrusimha, Aniruddha and Gholami, Amir and Abbeel, Pieter and Keutzer, Kurt and Stoica, Ion and Gonzalez, Joseph E},
journal={arXiv preprint arXiv:1910.02653},
year={2019}
}