RevBiFPN: The Fully Reversible Bidirectional Feature Pyramid Network
This work introduces the RevSilo, the first reversible module for
bidirectional multi-scale feature fusion. Like other reversible methods,
RevSilo eliminates the need to store hidden activations by recomputing them.
Existing reversible methods, however, do not apply to multi-scale feature
fusion and are therefore not applicable to a large class of networks.
Bidirectional multi-scale feature fusion promotes local and global coherence
and has become a de facto design principle for networks targeting spatially
sensitive tasks e.g. HRNet and EfficientDet. When paired with high-resolution
inputs, these networks achieve state-of-the-art results across various computer
vision tasks, but training them requires substantial accelerator memory for
saving large, multi-resolution activations. These memory requirements cap
network size and limit progress. Using reversible recomputation, the RevSilo
alleviates memory issues while still operating across resolution scales.
Stacking RevSilos, we create RevBiFPN, a fully reversible bidirectional feature
pyramid network. For classification, RevBiFPN is competitive with networks such
as EfficientNet while using up to 19.8x lesser training memory. When fine-tuned
on COCO, RevBiFPN provides up to a 2.5% boost in AP over HRNet using fewer MACs
and a 2.4x reduction in training-time memory.