Optimal checkpointing for heterogeneous chains: how to train deep neural
  networks with limited memory
  This paper introduces a new activation checkpointing method which allows to
significantly decrease memory usage when training Deep Neural Networks with the
back-propagation algorithm. Similarly to checkpoint-ing techniques coming from
the literature on Automatic Differentiation, it consists in dynamically
selecting the forward activations that are saved during the training phase, and
then automatically recomputing missing activations from those previously
recorded. We propose an original computation model that combines two types of
activation savings: either only storing the layer inputs, or recording the
complete history of operations that produced the outputs (this uses more
memory, but requires fewer recomputations in the backward phase), and we
provide an algorithm to compute the optimal computation sequence for this
model. This paper also describes a PyTorch implementation that processes the
entire chain, dealing with any sequential DNN whose internal layers may be
arbitrarily complex and automatically executing it according to the optimal
checkpointing strategy computed given a memory limit. Through extensive
experiments, we show that our implementation consistently outperforms existing
checkpoint-ing approaches for a large class of networks, image sizes and batch
sizes.
