TorchSyncBatchNorm

Bases: LayerSync

A plugin that wraps all batch normalization layers of a model with synchronization logic for multiprocessing.

This plugin has no effect in single-device operation.

apply(model)

Add global batchnorm for a model spread across multiple GPUs and nodes.

Override this method to synchronize batchnorm layers between specific process groups instead of the whole world.

Parameters:

model (Module) – Reference to the current LightningModule

Return type:

Module

Returns:

LightningModule with batchnorm layers synchronized within the process groups.

revert(model)

Convert the wrapped batchnorm layers back to regular batchnorm layers.

Parameters:

model (Module) – Reference to the current LightningModule

Return type:

Module

Returns:

LightningModule with regular batchnorm layers that will no longer sync across processes.