This repository provides a sublinear memory optimization for PyTorch deep learning models, targeting researchers and practitioners who need to train larger models or use larger batch sizes within limited GPU memory. It achieves this by reducing the memory footprint required for backpropagation from linear to square-root complexity.

How It Works

The core idea is to trade computation for memory by re-computing intermediate activations during the backward pass, rather than storing them all from the forward pass. This is implemented by replacing PyTorch's nn.Sequential with memonger.SublinearSequential. For models with non-deterministic layers like BatchNorm and Dropout, specific handling is included to maintain correctness, such as re-scaling BatchNorm momentum and memorizing RNG states for Dropout.

Quick Start & Requirements

• Install via pip: pip install pytorch-memonger
• Requires PyTorch.
• Supports nn.Sequential models.

Highlighted Details

• Reduces memory from O(N) to O(sqrt(N)) for backward pass.
• Achieves significant memory savings: ResNet152 with batch size 16 uses 2455MiB vs. 5459MiB.
• Supports BatchNorm and Dropout layers with specific handling for non-deterministic behavior.
• Re-implementation of the "Training Deep Nets with Sublinear Memory Cost" paper.

Maintenance & Community

• The project appears to be a personal implementation by Lyken17. No explicit community channels or roadmap are mentioned.

Licensing & Compatibility

• The README does not specify a license.

Limitations & Caveats

• Only supports nn.Sequential models due to PyTorch's define-by-run nature.
• Requires careful consideration for models with non-deterministic layers, though common ones are addressed.