class_docstr: Fast re-implementation of ``torch.nn.SyncBatchNorm`` that can achieve a speedup
  of 5x or more over the default implementation depending on size of the input
  and number of distributed workers.
class_docstr_tok: ["Fast","re-implementation","of","``","torch.nn.SyncBatchNorm","``","that","can","achieve","a","speedup","of","5x","or","more","over","the","default","implementation","depending","on","size","of","the","input","and","number","of","distributed","workers","."]
docstr: Helper function to convert all :attr:`BatchNorm*D` layers in the model to
  :class:`fairscale.experimental.nn.SyncBatchNorm` layers.
  
  Args:
      module (nn.Module): module containing one or more attr:`BatchNorm*D` layers
      process_group (optional): process group to scope synchronization,
          default is the whole world
  
  Returns:
      The original :attr:`module` with the converted :class:`torch.nn.SyncBatchNorm`
      layers. If the original :attr:`module` is a :attr:`BatchNorm*D` layer,
      a new :class:`torch.nn.SyncBatchNorm` layer object will be returned
      instead.
  
  Example::
  
      >>> # Network with nn.BatchNorm layer
      >>> module = torch.nn.Sequential(
      >>>            torch.nn.Linear(20, 100),
      >>>            torch.nn.BatchNorm1d(100),
      >>>          ).cuda()
      >>> # creating process group (optional)
      >>> # ranks is a list of int identifying rank ids.
      >>> ranks = list(range(8))
      >>> r1, r2 = ranks[:4], ranks[4:]
      >>> # Note: every rank calls into new_group for every
      >>> # process group created, even if that rank is not
      >>> # part of the group.
      >>> process_groups = [torch.distributed.new_group(pids) for pids in [r1, r2]]
      >>> process_group = process_groups[0 if dist.get_rank() <= 3 else 1]
      >>> sync_bn_module = fairscale.experimental.nn.SyncBatchNorm.convert_sync_batchnorm(module, process_group)
docstr_tok: ["Helper","function","to","convert","all",":","attr",":","`","BatchNorm","*","D","`","layers","in","the","model","to",":","class",":","`","fairscale.experimental.nn.SyncBatchNorm","`","layers",".","Args",":","module","(","nn.Module",")",":","module","containing","one","or","more","attr",":","`","BatchNorm","*","D","`","layers","process_group","(","optional",")",":","process","group","to","scope","synchronization",",","default","is","the","whole","world","Returns",":","The","original",":","attr",":","`","module","`","with","the","converted",":","class",":","`","torch.nn.SyncBatchNorm","`","layers",".","If","the","original",":","attr",":","`","module","`","is","a",":","attr",":","`","BatchNorm","*","D","`","layer",",","a","new",":","class",":","`","torch.nn.SyncBatchNorm","`","layer","object","will","be","returned","instead",".","Example",":",":",">",">",">","#","Network","with","nn.BatchNorm","layer",">",">",">","module","=","torch.nn.Sequential","(",">",">",">","torch.nn.Linear","(","20",",","100",")",",",">",">",">","torch.nn.BatchNorm1d","(","100",")",",",">",">",">",")",".cuda","(",")",">",">",">","#","creating","process","group","(","optional",")",">",">",">","#","ranks","is","a","list","of","int","identifying","rank","ids",".",">",">",">","ranks","=","list","(","range","(","8",")",")",">",">",">","r1",",","r2","=","ranks","[",":4","]",",","ranks","[","4",":","]",">",">",">","#","Note",":","every","rank","calls","into","new_group","for","every",">",">",">","#","process","group","created",",","even","if","that","rank","is","not",">",">",">","#","part","of","the","group",".",">",">",">","process_groups","=","[","torch.distributed.new_group","(","pids",")","for","pids","in","[","r1",",","r2","]","]",">",">",">","process_group","=","process_groups","[","0","if","dist.get_rank","(",")","<","=","3","else","1","]",">",">",">","sync_bn_module","=","fairscale.experimental.nn.SyncBatchNorm.convert_sync_batchnorm","(","module",",","process_group",")"]
code: def convert_sync_batchnorm(
      cls,
      module: torch.nn.Module,
      process_group: Optional[ProcessGroup] = None,
  ) -> torch.nn.Module:
      r"""Helper function to convert all :attr:`BatchNorm*D` layers in the model to
      :class:`fairscale.experimental.nn.SyncBatchNorm` layers.
  
      Args:
          module (nn.Module): module containing one or more attr:`BatchNorm*D` layers
          process_group (optional): process group to scope synchronization,
              default is the whole world
  
      Returns:
          The original :attr:`module` with the converted :class:`torch.nn.SyncBatchNorm`
          layers. If the original :attr:`module` is a :attr:`BatchNorm*D` layer,
          a new :class:`torch.nn.SyncBatchNorm` layer object will be returned
          instead.
  
      Example::
  
          >>> # Network with nn.BatchNorm layer
          >>> module = torch.nn.Sequential(
          >>>            torch.nn.Linear(20, 100),
          >>>            torch.nn.BatchNorm1d(100),
          >>>          ).cuda()
          >>> # creating process group (optional)
          >>> # ranks is a list of int identifying rank ids.
          >>> ranks = list(range(8))
          >>> r1, r2 = ranks[:4], ranks[4:]
          >>> # Note: every rank calls into new_group for every
          >>> # process group created, even if that rank is not
          >>> # part of the group.
          >>> process_groups = [torch.distributed.new_group(pids) for pids in [r1, r2]]
          >>> process_group = process_groups[0 if dist.get_rank() <= 3 else 1]
          >>> sync_bn_module = fairscale.experimental.nn.SyncBatchNorm.convert_sync_batchnorm(module, process_group)
  
      """
      module_output = module
      if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
          module_output = SyncBatchNorm(
              module.num_features,  # type: ignore
              module.eps,  # type: ignore
              module.momentum,  # type: ignore
              module.affine,  # type: ignore
              module.track_running_stats,  # type: ignore
              process_group=process_group,
          )
          if module.affine:
              with torch.no_grad():
                  module_output.weight = module.weight
                  module_output.bias = module.bias
          module_output.running_mean = module.running_mean
          module_output.running_var = module.running_var
          module_output.num_batches_tracked = module.num_batches_tracked
          if hasattr(module, "qconfig"):
              module_output.qconfig = module.qconfig
      for name, child in module.named_children():
          module_output.add_module(
              name, cls.convert_sync_batchnorm(child, process_group)
          )
      del module
      return module_output
  
file_path: fairscale/fairscale/experimental/nn/sync_batchnorm.py
imports: from typing import Any, Dict, List, Optional, Tuple
  import torch
  from torch import Tensor
  import torch.distributed
  from torch.distributed import ProcessGroup
  from fairscale.internal import torch_version
  from fairscale.nn.checkpoint import is_checkpointing, is_recomputing
label_desc: Use image node_id 3 for calling the SyncBatchNorm obj's underlying member method code with example usage: obj.convert_sync_batchnorm(cls, module, process_group) and returns: module_output