Skip to content

BatchCentering and BatchLipNorm classes compatible with multi-GPU - #55

Merged
franckma31 merged 19 commits into
masterfrom
develop
Mar 23, 2026
Merged

BatchCentering and BatchLipNorm classes compatible with multi-GPU #55
franckma31 merged 19 commits into
masterfrom
develop

Conversation

@franckma31

Copy link
Copy Markdown
Collaborator
  • BatchCentering support for multi-GPU training, based on a running_mean accumulator and a running_num_batches (implemented as a special case of batchLipNorm class)
  • BatchLipNorm compute the variance on batch and full set by using running means of square. This layer is not Lipschitz and can be used with a SharedLipFactory to maintain the global Lipschitz constant
  • Vanilla export for BatchCentering and BatchLipNorm with the layer ScaleBiasLayer
  • pytests on BatchCentering and BatchLipNorm
  • LipschitzModule class inherits from ScaledLipschitzModule to be able to register the coefficient_lip to a SharedLipFactory
    multigpu tests are excluded from pytest : use torchrun on a multigpu config to run tests/multigpu scripts

…ng_mean accumulator and a running_num_batches
…seful only if we want to register scaling values), + modify scaling factor to get directly the multiplied value 1.0 / sqrt(max(var))
…ovide the get_scaling_factor and support factory also in LipschitzModule class
BatchCentering and BatchLipNorm classes compatible with multi-GPU
@franckma31
franckma31 requested a review from cofri March 16, 2026 22:21
@franckma31
franckma31 merged commit a477697 into master Mar 23, 2026
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants