Source-linked AI summary
In-Place Activated BatchNorm for Memory-Optimized Training of DNNs
Samuel Rota Bulò, Lorenzo Porzi, Peter Kontschieder
TL;DR
Training modern deep networks is constrained by GPU memory, particularly for semantic segmentation. The paper introduces INPLACE-ABN, a fused BN-and-activation layer that reconstructs backward-pass quantities through inversion. It reports approximately on-par ImageNet classification and significantly improved semantic-segmentation results, while targeting training-time rather than test-time memory efficiency.
Problem
Deep, wide networks create high GPU-memory requirements, forcing compromises in feature extraction, output resolution, and training data, especially for semantic segmentation.
Method
INPLACE-ABN merges batch normalization and activation into a self-contained in-place layer, storing one buffer and recovering backward-pass quantities by inverting forward computations.
Results
ImageNet classification is approximately on par with state-of-the-art models, while semantic segmentation shows significantly improved results across the reported tasks.
Takeaways & Limitations
The approach provides a plug-in strategy for reducing training memory in modern networks and supports stronger configurations for memory-critical semantic segmentation.
Takeaways & Limitations
The evaluation targets training-time memory efficiency rather than test-time optimization, and reported results may differ from original papers because training protocols are not identical.
Abstract
from arXiv · showhide
In this work we present In-Place Activated Batch Normalization (InPlace-ABN) - a novel approach to drastically reduce the training memory footprint of modern deep neural networks in a computationally efficient way. Our solution substitutes the conventionally used succession of BatchNorm + Activation layers with a single plugin layer, hence avoiding invasive framework surgery while providing straightforward applicability for existing deep learning frameworks. We obtain memory savings of up to 50% by dropping intermediate results and by recovering required information during the backward pass through the inversion of stored forward results, with only minor increase (0.8-2%) in computation time. Also, we demonstrate how frequently used checkpointing approaches can be made computationally as efficient as InPlace-ABN. In our experiments on image classification, we demonstrate on-par results on ImageNet-1k with state-of-the-art approaches. On the memory-demanding task of semantic segmentation, we report results for COCO-Stuff, Cityscapes and Mapillary Vistas, obtaining new state-of-the-art results on the latter without additional training data but in a single-scale and -model scenario. Code can be found at https://github.com/mapillary/inplace_abn .
1. Introduction
Modern deep networks create substantial training-memory demands, especially for semantic segmentation. The paper introduces INPLACE-ABN, which merges batch normalization and activation to reduce memory while maintaining computational efficiency and strong task performance.
- Deep, wide backbones increase GPU memory requirements and force trade-offs involving feature-extractor capacity, output resolution, and training data size.
- The method targets memory-efficient training for modern architectures and can provide additional memory for demanding applications such as semantic segmentation.
- INPLACE-ABN merges batch normalization and activation into one self-contained layer, enabling in-place computation with a single buffer and backward recovery by inversion.
- Up to 50% theoretical training-memory reduction is achieved by avoiding stored BN+ACT intermediates and reconstructing required quantities during backpropagation.
- The optimized checkpointing variant reduces recomputation and can match INPLACE-ABN computational efficiency, while conventional checkpointing requires invasive graph manipulation.
- ImageNet-1k classification is approximately on par with state-of-the-art models, while semantic segmentation results improve substantially across the evaluated datasets.
2. Related Work
Prior memory-reduction approaches trade memory against computation, precision, or implementation complexity. The paper positions INPLACE-ABN as a self-contained alternative focused on efficient training-memory management.
- Checkpointing saves memory by storing activation checkpoints and recomputing required quantities during the backward pass, at the cost of runtime.
- Reduced-precision methods decrease memory-footprint datatypes but typically degrade overall performance.
- Reversible residual blocks reconstruct activations without storing intermediates but recompute residual functions during backward propagation, incurring checkpointing-like overhead.
- The paper scopes its optimization target to training-time memory efficiency and excludes test-time optimization such as TensorRT.
3. In-Place Activated Batch Normalization
The paper develops memory-efficient alternatives to standard BN and checkpointing by recovering backward-pass information from stored forward results. The proposed methods are self-contained, reduce memory substantially, and minimize computational overhead.
- INPLACE-ABN: INPLACE-ABN merges BN and activation into one self-contained layer, stores one buffer, and recovers required quantities by inverting forward computations.Its main components are invertible activation and scale-and-shift operations, plus a BN backward pass expressed using the forward output y.
- Batch Normalization Review: Batch Normalization whitens minibatch activations, then applies learnable scaling and shifting while retaining representation capacity.The minibatch statistics are the empirical mean and variance; the scale-and-shift parameters are γ and β.
- Standard and Checkpointed Computation: Standard BN stores both its input and activation output because backward computation requires them for BN, activation, and subsequent layers.In the illustrated BN+ACT+CONV block, x supports BN gradients, while z supports activation and convolution operations.
- Standard and Checkpointed Computation: Checkpointing stores x and recomputes the activation output during backward propagation, reducing the illustrated block’s storage from two buffers to one.The recomputation reiterates forward operations from x, trading computation for memory.
- Proposed Checkpointing: The proposed optimized checkpointing stores the normalized value x̂ instead of x, allowing backward recovery through scale-and-shift followed by activation.This avoids recomputing x̂ during BN backward computation and is therefore more computationally efficient than direct checkpointing for the block.
- INPLACE-ABN: INPLACE-ABN II rewrites gradients directly as functions of y, absorbing scale-and-shift recovery into the gradient computation at O(1) cost per feature channel.The paper summarizes both optimized main contributions as more memory-efficient than state-of-the-art checkpointing, with INPLACE-ABN II computationally more efficient than optimized checkpointing.
4. Experiments
The experiments evaluate InPlace-ABN across image classification, semantic segmentation, and runtime, examining activation choices, architectures, memory utilization, and BN strategies. Results include comparable ImageNet performance, improved segmentation outcomes under expanded memory settings, and low computational overhead.
- Experimental setup: Experiments use four NVIDIA Titan Xp GPUs with 12GB RAM each, typically applying LEAKY RELU with slope a = 0.01.
- Image classification: Replacing RELU with LEAKY RELU, or vice versa, changes validation results by no more than a single point per training except for one 3202 center-crop top-1 evaluation.
- Image classification: InPlace-ABN memory savings can support larger batches or larger architectures, and both options produce a noticeable ImageNet performance increase.ResNeXt-101 with batch size 512 achieves similar accuracy to the deeper ResNeXt-152 model, while neither larger architecture fits with standard BN at batch size 256.
- Semantic segmentation: Segmentation experiments use Cityscapes, COCO-Stuff, and Mapillary Vistas with DeepLabV3 heads and single-scale evaluation settings.The evaluated bodies include ResNeXt-101, ResNeXt-152, and WideResNet-38, with INPLACE-ABN variants also used in the head.
- Semantic segmentation: 79.16% validation score is obtained on Cityscapes with WideResNet-38, INPLACE-ABNsync, and 12 crops at 872 × 872.This improves by approximately 0.9% over the best-performing setting in Table 3.
- Semantic segmentation: 53.12% validation score on Mapillary Vistas significantly outperforms the LSUN 2017 winner’s 51.59% single-scale approach.The setting uses 12 crops at 776×776 with INPLACE-ABNsync, without additional training data described in the passage.
5. Conclusions
INPLACE-ABN fuses batch normalization and activation layers to reduce training memory while reconstructing backward-pass quantities by inverting the forward computation. The method is computationally efficient, easy to deploy, and validated on image classification and semantic segmentation.
- INPLACE-ABN fuses batch normalization and activation layers for memory-optimized training of modern deep neural networks.
- Backward-pass quantities are reconstructed by inverting the forward computation from the storage buffer.
Appendix – Derivation of Gradient ∂L
The approach reconstructs discarded buffers during the backward pass, allowing BN+ACT to be encapsulated as a self-contained layer. This supports implementation and deployment across modern deep learning frameworks, with validation on classification and segmentation tasks.
- Discarded buffers are reconstructed backwards during the backward pass rather than recomputed conventionally.
- Encapsulating BN+ACT as a self-contained layer makes the method easy to implement and deploy in virtually all modern deep learning frameworks.
- The approach was validated through experiments on ImageNet-1k image classification and semantic segmentation.