Huge News!Announcing our $40M Series B led by Abstract Ventures.Learn More
Socket
Sign inDemoInstall
Socket

gradient-accumulator

Package Overview
Dependencies
Maintainers
1
Alerts
File Explorer

Advanced tools

Socket logo

Install Socket

Detect and block malicious and high-risk dependencies

Install

gradient-accumulator

Package for gradient accumulation in TensorFlow

  • 0.5.2
  • PyPI
  • Socket score

Maintainers
1
gradient-accumulator

GradientAccumulator

Seemless gradient accumulation for TensorFlow 2

Pip Downloads PyPI version License DOI

GradientAccumulator was developed by SINTEF Health due to the lack of an easy-to-use method for gradient accumulation in TensorFlow 2.

The package is available on PyPI and is compatible with and have been tested against TensorFlow 2.2-2.12 and Python 3.6-3.11, and works cross-platform (Ubuntu, Windows, macOS).

Continuous integration

Build TypeStatus
Code coveragecodecov
DocumentationsDocumentation Status
Unit testsCI

Install

Stable release from PyPI:

pip install gradient-accumulator

Or from source:

pip install git+https://github.com/andreped/GradientAccumulator

Getting started

A simple example to add gradient accumulation to an existing model is by:

from gradient_accumulator import GradientAccumulateModel
from tensorflow.keras.models import Model

model = Model(...)
model = GradientAccumulateModel(accum_steps=4, inputs=model.input, outputs=model.output)

Then simply use the model as you normally would!

In practice, using gradient accumulation with a custom pipeline might require some extra overhead and tricks to get working.

For more information, see documentations which are hosted at gradientaccumulator.readthedocs.io

What?

Gradient accumulation (GA) enables reduced GPU memory consumption through dividing a batch into smaller reduced batches, and performing gradient computation either in a distributing setting across multiple GPUs or sequentially on the same GPU. When the full batch is processed, the gradients are then accumulated to produce the full batch gradient.

Note that how we implemented gradient accumulation is slightly different from this illustration, as our design does not require having the entire batch in CPU memory. More information on what goes under the hood can be seen in the documentations.

Why?

In TensorFlow 2, there did not exist a plug-and-play method to use gradient accumulation with any custom pipeline. Hence, we have implemented two generic TF2-compatible approaches:

MethodUsage
GradientAccumulateModelmodel = GradientAccumulateModel(accum_steps=4, inputs=model.input, outputs=model.output)
GradientAccumulateOptimizeropt = GradientAccumulateOptimizer(accum_steps=4, optimizer=tf.keras.optimizers.SGD(1e-2))

Both approaches control how frequently the weigths are updated but in their own way. Approach (1) overrides the train_step method of a given Model, whereas approach (2) wraps the optimizer. (1) is only compatible with single-GPU usage, whereas (2) also supports distributed training (multi-GPU).

Our implementations enable theoretically infinitely large batch size, with identical memory consumption as for a regular mini batch. If a single GPU is used, this comes at the cost of increased training runtime. Multiple GPUs could be used to improve runtime performance.

TechniqueUsage
Batch Normalizationlayer = AccumBatchNormalization(accum_steps=4)
Adaptive Gradient Clippingmodel = GradientAccumulateModel(accum_steps=4, agc=True, inputs=model.input, outputs=model.output)
Mixed precisionmodel = GradientAccumulateModel(accum_steps=4, mixed_precision=True, inputs=model.input, outputs=model.output)
  • As batch normalization (BN) is not natively compatible with GA, we have implemented a custom BN layer which can be used as a drop-in replacement.
  • Support for adaptive gradient clipping has been added as an alternative to BN.
  • Mixed precision can also be utilized on both GPUs and TPUs.
  • Multi-GPU distributed training using generic optimizer wrapper.

For more information on usage, supported techniques, and examples, refer to the documentations.

Applications

  • Helland et al., Segmentation of glioblastomas in early post-operative multi-modal MRI with deep neural networks, arXiv, 2023, https://arxiv.org/abs/2304.08881
  • Bouget et al., Raidionics: an open software for pre- and postoperative central nervous system tumor segmentation and standardized reporting, arXiv, 2023, https://arxiv.org/abs/2305.14351
  • Pedersen et al., H2G-Net: A multi-resolution refinement approach for segmentation of breast cancer region in gigapixel histopathological images, Frontiers in Medicine, 2022, https://doi.org/10.3389/fmed.2022.971873

Acknowledgements

The gradient accumulator model wrapper is based on the implementation presented in this thread on stack overflow. The adaptive gradient clipping method is based on the implementation by @sayakpaul. The optimizer wrapper is derived from the implementation by @fsx950223 and @stefan-falk.

The documentations hosted here was made possible by the incredible Read The Docs team which offer free documentation hosting!

How to cite?

If you used this package or found the project relevant in your research, please, include the following citation:

@software{andre_pedersen_2023_7905351,
  author       = {André Pedersen and Tor-Arne Schmidt Nordmo and Javier Pérez de Frutos and David Bouget},
  title        = {andreped/GradientAccumulator: v0.5.0},
  month        = may,
  year         = 2023,
  publisher    = {Zenodo},
  version      = {v0.5.0},
  doi          = {10.5281/zenodo.7905351},
  url          = {https://doi.org/10.5281/zenodo.7905351}
}

FAQs


Did you know?

Socket

Socket for GitHub automatically highlights issues in each pull request and monitors the health of all your open source dependencies. Discover the contents of your packages and block harmful activity before you install or update your dependencies.

Install

Related posts

SocketSocket SOC 2 Logo

Product

  • Package Alerts
  • Integrations
  • Docs
  • Pricing
  • FAQ
  • Roadmap
  • Changelog

Packages

npm

Stay in touch

Get open source security insights delivered straight into your inbox.


  • Terms
  • Privacy
  • Security

Made with ⚡️ by Socket Inc