PhdddlmZmZddlZddlmZddlmZmZmZddl m Z ddl m Z dd l mZdd lmZdd lmZgd ZGd deZGddeZGddeeZGddeZGddeeZGddeZGddeeZGddeZGddeeZGdd eZy)!)OptionalAnyN)Tensor) ParameterUninitializedParameterUninitializedBuffer) functional)init) SyncBatchNorm)LazyModuleMixin)Module) BatchNorm1dLazyBatchNorm1d BatchNorm2dLazyBatchNorm2d BatchNorm3dLazyBatchNorm3dr c eZdZUdZdZgdZeed<eed<eed<e ed<e ed< ddededede de d d f fd Z dd Z dd Z dZ dZfdZxZS) _NormBasez+Common base of _InstanceNorm and _BatchNormr )track_running_statsmomentumeps num_featuresaffinerrrrrNreturnc t||d}t |||_||_||_||_||_|j rIttj|fi||_ ttj|fi||_ n$|jdd|jdd|j r|jdtj|fi||jdtj|fi||||jdtj ddtj"i|j%D cic]\} } | dk7s | | c} } |n6|jdd|jdd|jdd|j'ycc} } w) Ndevicedtypeweightbias running_mean running_varnum_batches_trackedr!r)super__init__rrrrrrtorchemptyr"r#register_parameterregister_bufferzerosonestensorlongitemsreset_parameters) selfrrrrrr r!factory_kwargskv __class__s eC:\Users\daisl\Desktop\realtime-object-detection\venv\Lib\site-packages\torch/nn/modules/batchnorm.pyr)z_NormBase.__init__s%+U; (   #6 ;;#EKK $O$OPDK!%++l"Mn"MNDI  # #Hd 3  # #FD 1  # #  \1\^1\ ]   <0Z>0Z [    !6!&"kuzz"kBPBVBVBX0iBX$!Q\]ah\hABX0i"k l    6   5  !6 =  1js F4F4c|jrP|jj|jj d|j jyyNr )rr$zero_r%fill_r&r4s r9reset_running_statsz_NormBase.reset_running_statsDsJ  # #    # # %    " "1 %  $ $ * * , $c|j|jr?tj|jtj |j yyN)r?rr ones_r"zeros_r#r>s r9r3z_NormBase.reset_parametersLs:   " ;; JJt{{ # KK " r@ctrB)NotImplementedErrorr4inputs r9_check_input_dimz_NormBase._check_input_dimRs!!r@c:djdi|jS)Nzj{num_features}, eps={eps}, momentum={momentum}, affine={affine}, track_running_stats={track_running_stats})format__dict__r>s r9 extra_reprz_NormBase.extra_reprUs) ? 88> PAE P r@c  |jdd}||dkrU|jrI|dz} | |vr@|j |jn$tjdtj || <t ||||||||y)Nversionr r&r)r!)getrr&r*r0r1r(_load_from_state_dict) r4 state_dictprefixlocal_metadatastrict missing_keysunexpected_keys error_msgsrPnum_batches_tracked_keyr8s r9rRz_NormBase._load_from_state_dict[s!$$Y5 Ow{0H0H'-/D&D #&j8//;,,auzz:23 %        r@h㈵>皙?TTNNrN)__name__ __module__ __qualname____doc___version __constants__int__annotations__floatboolr)r?r3rIrNrR __classcell__r8s@r9rrs5HXM JO L$($ $ $  $  $ " $  $ L-# "   r@rc VeZdZ d dedededededdf fd Zd edefd ZxZ S) _BatchNormNrrrrrrc8||d}t ||||||fi|yNr)r(r)) r4rrrrrr r!r5r8s r9r)z_BatchNorm.__init__~s0%+U;  #x1D HV r@rHc |j||jd}n |j}|jrd|jrX|jL|jj d|jdt |jz }n |j} |jrd}n|jduxr|jdu} tj||jr |jr |jnd|jr |jr |jnd|j|j|||jS)Nr ?T)rIrtrainingrr&add_rgr$r%F batch_normr"r#r)r4rHexponential_average_factor bn_trainings r9forwardz_BatchNorm.forwards' e$ == ), &)- & ==T55''3((--a0==(14uT=U=U7V1V.15.  ==K,,4T4;K;Kt;SK || ==D$<$<   $(MMT5M5MD  SW KK II  & HH  r@r[) r_r`rarergrhr)rrxrirjs@r9rlrl}si$(          "     . V. . r@rlcReZdZUeed<eed< d dfd Zdfd ZddZxZS) _LazyNormBaser"r#c ||d}t |d||ddfi|||_||_|jr t di||_t di||_|jrttdi||_tdi||_ tj ddtji|jD cic]\}} |dk7s || c} }|_yycc} }w)NrrFr!rKr')r(r)rrrr"r#rr$r%r*r0r1r2r&) r4rrrrr r!r5r6r7r8s r9r)z_LazyNormBase.__init__s$*U;         #6 ;;0B>BDK.@@DI  # # 3 En ED 2D^DD ',||(b(b9G9M9M9O'`9OASTX_S_19O'`(bD $ $(as 6 CCcd|js|jdk7rt| yyy)Nr)has_uninitialized_paramsrr(r3)r4r8s r9r3z_LazyNormBase.reset_parameterss0,,.43D3D3I G $ &4J.r@c@|jr |jd|_|jrt |j t sJt |jt sJ|j j|jf|jj|jf|jrL|jj|jf|jj|jf|jyyr;) r}shaperr isinstancer"rr# materializerr$r%r3rGs r9initialize_parametersz#_LazyNormBase.initialize_parameterss  ( ( * % AD {{!$++/EFFF!$))-CDDD ''):):(<= %%t'8'8&:;''!!--t/@/@.BC  ,,d.?.?-AB  ! ! # +r@r[r^) r_r`rarrfr)r3rrirjs@r9rzrzs. ""  PT$(b-1b0' $r@rzceZdZdZdZy)ra Applies Batch Normalization over a 2D or 3D input as described in the paper `Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift `__ . .. math:: y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta The mean and standard-deviation are calculated per-dimension over the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors of size `C` (where `C` is the number of features or channels of the input). By default, the elements of :math:`\gamma` are set to 1 and the elements of :math:`\beta` are set to 0. At train time in the forward pass, the standard-deviation is calculated via the biased estimator, equivalent to ``torch.var(input, unbiased=False)``. However, the value stored in the moving average of the standard-deviation is calculated via the unbiased estimator, equivalent to ``torch.var(input, unbiased=True)``. Also by default, during training this layer keeps running estimates of its computed mean and variance, which are then used for normalization during evaluation. The running estimates are kept with a default :attr:`momentum` of 0.1. If :attr:`track_running_stats` is set to ``False``, this layer then does not keep running estimates, and batch statistics are instead used during evaluation time as well. .. note:: This :attr:`momentum` argument is different from one used in optimizer classes and the conventional notion of momentum. Mathematically, the update rule for running statistics here is :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`, where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the new observed value. Because the Batch Normalization is done over the `C` dimension, computing statistics on `(N, L)` slices, it's common terminology to call this Temporal Batch Normalization. Args: num_features: number of features or channels :math:`C` of the input eps: a value added to the denominator for numerical stability. Default: 1e-5 momentum: the value used for the running_mean and running_var computation. Can be set to ``None`` for cumulative moving average (i.e. simple average). Default: 0.1 affine: a boolean value that when set to ``True``, this module has learnable affine parameters. Default: ``True`` track_running_stats: a boolean value that when set to ``True``, this module tracks the running mean and variance, and when set to ``False``, this module does not track such statistics, and initializes statistics buffers :attr:`running_mean` and :attr:`running_var` as ``None``. When these buffers are ``None``, this module always uses batch statistics. in both training and eval modes. Default: ``True`` Shape: - Input: :math:`(N, C)` or :math:`(N, C, L)`, where :math:`N` is the batch size, :math:`C` is the number of features or channels, and :math:`L` is the sequence length - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input) Examples:: >>> # With Learnable Parameters >>> m = nn.BatchNorm1d(100) >>> # Without Learnable Parameters >>> m = nn.BatchNorm1d(100, affine=False) >>> input = torch.randn(20, 100) >>> output = m(input) c|jdk7r1|jdk7rtd|jdyyNr zexpected 2D or 3D input (got D input)dim ValueErrorrGs r9rIzBatchNorm1d._check_input_dim2D 99;!  q 0/ }HE !1 r@Nr_r`rarbrIrKr@r9rrsBHr@rceZdZdZeZdZy)ra6A :class:`torch.nn.BatchNorm1d` module with lazy initialization of the ``num_features`` argument of the :class:`BatchNorm1d` that is inferred from the ``input.size(1)``. The attributes that will be lazily initialized are `weight`, `bias`, `running_mean` and `running_var`. Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation on lazy modules and their limitations. Args: eps: a value added to the denominator for numerical stability. Default: 1e-5 momentum: the value used for the running_mean and running_var computation. Can be set to ``None`` for cumulative moving average (i.e. simple average). Default: 0.1 affine: a boolean value that when set to ``True``, this module has learnable affine parameters. Default: ``True`` track_running_stats: a boolean value that when set to ``True``, this module tracks the running mean and variance, and when set to ``False``, this module does not track such statistics, and initializes statistics buffers :attr:`running_mean` and :attr:`running_var` as ``None``. When these buffers are ``None``, this module always uses batch statistics. in both training and eval modes. Default: ``True`` c|jdk7r1|jdk7rtd|jdyyrrrGs r9rIz LazyBatchNorm1d._check_input_dimUrr@N)r_r`rarbr cls_to_becomerIrKr@r9rr9s2 Mr@rceZdZdZdZy)ra Applies Batch Normalization over a 4D input (a mini-batch of 2D inputs with additional channel dimension) as described in the paper `Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift `__ . .. math:: y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta The mean and standard-deviation are calculated per-dimension over the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors of size `C` (where `C` is the input size). By default, the elements of :math:`\gamma` are set to 1 and the elements of :math:`\beta` are set to 0. At train time in the forward pass, the standard-deviation is calculated via the biased estimator, equivalent to ``torch.var(input, unbiased=False)``. However, the value stored in the moving average of the standard-deviation is calculated via the unbiased estimator, equivalent to ``torch.var(input, unbiased=True)``. Also by default, during training this layer keeps running estimates of its computed mean and variance, which are then used for normalization during evaluation. The running estimates are kept with a default :attr:`momentum` of 0.1. If :attr:`track_running_stats` is set to ``False``, this layer then does not keep running estimates, and batch statistics are instead used during evaluation time as well. .. note:: This :attr:`momentum` argument is different from one used in optimizer classes and the conventional notion of momentum. Mathematically, the update rule for running statistics here is :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`, where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the new observed value. Because the Batch Normalization is done over the `C` dimension, computing statistics on `(N, H, W)` slices, it's common terminology to call this Spatial Batch Normalization. Args: num_features: :math:`C` from an expected input of size :math:`(N, C, H, W)` eps: a value added to the denominator for numerical stability. Default: 1e-5 momentum: the value used for the running_mean and running_var computation. Can be set to ``None`` for cumulative moving average (i.e. simple average). Default: 0.1 affine: a boolean value that when set to ``True``, this module has learnable affine parameters. Default: ``True`` track_running_stats: a boolean value that when set to ``True``, this module tracks the running mean and variance, and when set to ``False``, this module does not track such statistics, and initializes statistics buffers :attr:`running_mean` and :attr:`running_var` as ``None``. When these buffers are ``None``, this module always uses batch statistics. in both training and eval modes. Default: ``True`` Shape: - Input: :math:`(N, C, H, W)` - Output: :math:`(N, C, H, W)` (same shape as input) Examples:: >>> # With Learnable Parameters >>> m = nn.BatchNorm2d(100) >>> # Without Learnable Parameters >>> m = nn.BatchNorm2d(100, affine=False) >>> input = torch.randn(20, 100, 35, 45) >>> output = m(input) cd|jdk7rtd|jdyNzexpected 4D input (got rrrGs r9rIzBatchNorm2d._check_input_dim0 99;! 6uyy{m8LM M r@NrrKr@r9rr\sCJNr@rceZdZdZeZdZy)ra6A :class:`torch.nn.BatchNorm2d` module with lazy initialization of the ``num_features`` argument of the :class:`BatchNorm2d` that is inferred from the ``input.size(1)``. The attributes that will be lazily initialized are `weight`, `bias`, `running_mean` and `running_var`. Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation on lazy modules and their limitations. Args: eps: a value added to the denominator for numerical stability. Default: 1e-5 momentum: the value used for the running_mean and running_var computation. Can be set to ``None`` for cumulative moving average (i.e. simple average). Default: 0.1 affine: a boolean value that when set to ``True``, this module has learnable affine parameters. Default: ``True`` track_running_stats: a boolean value that when set to ``True``, this module tracks the running mean and variance, and when set to ``False``, this module does not track such statistics, and initializes statistics buffers :attr:`running_mean` and :attr:`running_var` as ``None``. When these buffers are ``None``, this module always uses batch statistics. in both training and eval modes. Default: ``True`` cd|jdk7rtd|jdyrrrGs r9rIz LazyBatchNorm2d._check_input_dimrr@N)r_r`rarbrrrIrKr@r9rr2 MNr@rceZdZdZdZy)ra Applies Batch Normalization over a 5D input (a mini-batch of 3D inputs with additional channel dimension) as described in the paper `Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift `__ . .. math:: y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta The mean and standard-deviation are calculated per-dimension over the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors of size `C` (where `C` is the input size). By default, the elements of :math:`\gamma` are set to 1 and the elements of :math:`\beta` are set to 0. At train time in the forward pass, the standard-deviation is calculated via the biased estimator, equivalent to ``torch.var(input, unbiased=False)``. However, the value stored in the moving average of the standard-deviation is calculated via the unbiased estimator, equivalent to ``torch.var(input, unbiased=True)``. Also by default, during training this layer keeps running estimates of its computed mean and variance, which are then used for normalization during evaluation. The running estimates are kept with a default :attr:`momentum` of 0.1. If :attr:`track_running_stats` is set to ``False``, this layer then does not keep running estimates, and batch statistics are instead used during evaluation time as well. .. note:: This :attr:`momentum` argument is different from one used in optimizer classes and the conventional notion of momentum. Mathematically, the update rule for running statistics here is :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`, where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the new observed value. Because the Batch Normalization is done over the `C` dimension, computing statistics on `(N, D, H, W)` slices, it's common terminology to call this Volumetric Batch Normalization or Spatio-temporal Batch Normalization. Args: num_features: :math:`C` from an expected input of size :math:`(N, C, D, H, W)` eps: a value added to the denominator for numerical stability. Default: 1e-5 momentum: the value used for the running_mean and running_var computation. Can be set to ``None`` for cumulative moving average (i.e. simple average). Default: 0.1 affine: a boolean value that when set to ``True``, this module has learnable affine parameters. Default: ``True`` track_running_stats: a boolean value that when set to ``True``, this module tracks the running mean and variance, and when set to ``False``, this module does not track such statistics, and initializes statistics buffers :attr:`running_mean` and :attr:`running_var` as ``None``. When these buffers are ``None``, this module always uses batch statistics. in both training and eval modes. Default: ``True`` Shape: - Input: :math:`(N, C, D, H, W)` - Output: :math:`(N, C, D, H, W)` (same shape as input) Examples:: >>> # With Learnable Parameters >>> m = nn.BatchNorm3d(100) >>> # Without Learnable Parameters >>> m = nn.BatchNorm3d(100, affine=False) >>> input = torch.randn(20, 100, 35, 45, 10) >>> output = m(input) cd|jdk7rtd|jdyNzexpected 5D input (got rrrGs r9rIzBatchNorm3d._check_input_dimrr@NrrKr@r9rrsDLNr@rceZdZdZeZdZy)ra6A :class:`torch.nn.BatchNorm3d` module with lazy initialization of the ``num_features`` argument of the :class:`BatchNorm3d` that is inferred from the ``input.size(1)``. The attributes that will be lazily initialized are `weight`, `bias`, `running_mean` and `running_var`. Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation on lazy modules and their limitations. Args: eps: a value added to the denominator for numerical stability. Default: 1e-5 momentum: the value used for the running_mean and running_var computation. Can be set to ``None`` for cumulative moving average (i.e. simple average). Default: 0.1 affine: a boolean value that when set to ``True``, this module has learnable affine parameters. Default: ``True`` track_running_stats: a boolean value that when set to ``True``, this module tracks the running mean and variance, and when set to ``False``, this module does not track such statistics, and initializes statistics buffers :attr:`running_mean` and :attr:`running_var` as ``None``. When these buffers are ``None``, this module always uses batch statistics. in both training and eval modes. Default: ``True`` cd|jdk7rtd|jdyrrrGs r9rIz LazyBatchNorm3d._check_input_dim0rr@N)r_r`rarbrrrIrKr@r9rrrr@rceZdZdZ ddedededededeed dffd Z d Z d Z d e d e fdZ eddZxZS)r aApplies Batch Normalization over a N-Dimensional input (a mini-batch of [N-2]D inputs with additional channel dimension) as described in the paper `Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift `__ . .. math:: y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta The mean and standard-deviation are calculated per-dimension over all mini-batches of the same process groups. :math:`\gamma` and :math:`\beta` are learnable parameter vectors of size `C` (where `C` is the input size). By default, the elements of :math:`\gamma` are sampled from :math:`\mathcal{U}(0, 1)` and the elements of :math:`\beta` are set to 0. The standard-deviation is calculated via the biased estimator, equivalent to `torch.var(input, unbiased=False)`. Also by default, during training this layer keeps running estimates of its computed mean and variance, which are then used for normalization during evaluation. The running estimates are kept with a default :attr:`momentum` of 0.1. If :attr:`track_running_stats` is set to ``False``, this layer then does not keep running estimates, and batch statistics are instead used during evaluation time as well. .. note:: This :attr:`momentum` argument is different from one used in optimizer classes and the conventional notion of momentum. Mathematically, the update rule for running statistics here is :math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`, where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the new observed value. Because the Batch Normalization is done for each channel in the ``C`` dimension, computing statistics on ``(N, +)`` slices, it's common terminology to call this Volumetric Batch Normalization or Spatio-temporal Batch Normalization. Currently :class:`SyncBatchNorm` only supports :class:`~torch.nn.DistributedDataParallel` (DDP) with single GPU per process. Use :meth:`torch.nn.SyncBatchNorm.convert_sync_batchnorm()` to convert :attr:`BatchNorm*D` layer to :class:`SyncBatchNorm` before wrapping Network with DDP. Args: num_features: :math:`C` from an expected input of size :math:`(N, C, +)` eps: a value added to the denominator for numerical stability. Default: ``1e-5`` momentum: the value used for the running_mean and running_var computation. Can be set to ``None`` for cumulative moving average (i.e. simple average). Default: 0.1 affine: a boolean value that when set to ``True``, this module has learnable affine parameters. Default: ``True`` track_running_stats: a boolean value that when set to ``True``, this module tracks the running mean and variance, and when set to ``False``, this module does not track such statistics, and initializes statistics buffers :attr:`running_mean` and :attr:`running_var` as ``None``. When these buffers are ``None``, this module always uses batch statistics. in both training and eval modes. Default: ``True`` process_group: synchronization of stats happen within each process group individually. Default behavior is synchronization across the whole world Shape: - Input: :math:`(N, C, +)` - Output: :math:`(N, C, +)` (same shape as input) .. note:: Synchronization of batchnorm statistics occurs only while training, i.e. synchronization is disabled when ``model.eval()`` is set or if ``self.training`` is otherwise ``False``. Examples:: >>> # xdoctest: +SKIP >>> # With Learnable Parameters >>> m = nn.SyncBatchNorm(100) >>> # 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] >>> # Without Learnable Parameters >>> m = nn.BatchNorm3d(100, affine=False, process_group=process_group) >>> input = torch.randn(20, 100, 35, 45, 10) >>> output = m(input) >>> # network is nn.BatchNorm layer >>> sync_bn_network = nn.SyncBatchNorm.convert_sync_batchnorm(network, process_group) >>> # only single gpu per process is currently supported >>> ddp_sync_bn_network = torch.nn.parallel.DistributedDataParallel( >>> sync_bn_network, >>> device_ids=[args.local_rank], >>> output_device=args.local_rank) Nrrrrr process_grouprc F||d} t ||||||fi| ||_yrn)r(r)r) r4rrrrrrr r!r5r8s r9r)zSyncBatchNorm.__init__s:%+U;  #x1D HV +r@cd|jdkrtd|jdy)Nr z expected at least 2D input (got rrrGs r9rIzSyncBatchNorm._check_input_dims3 99;?2599;-xH  r@cB|jddk(r tdy)Nr rz9SyncBatchNorm number of input channels should be non-zero)sizerrGs r9_check_non_zero_input_channelsz,SyncBatchNorm._check_non_zero_input_channelss' ::a=A K  r@rHc z|j||j||jd}n |j}|jrk|jr_|j J|j j d|jd|j jz }n |j} |jrd}n|jduxr|jdu} |jr |jr |jnd}|jr |jr |jnd}|xrL|jxr>tjjxrtjj}|r|jjdtj j#fvr*t%dtj j#tjj&j(}|j*r |j*}tjj-|}|dkD}|s:t/j0||||j2|j4|||j6S|sJt9j:||j2|j4|||j6| S)Nrpr rqTcudaz4SyncBatchNorm expected input tensor to be on GPU or )rIrrrrrr&rsitemr$r%r* distributed is_availableis_initializedr type_C_get_privateuse1_backend_namergroupWORLDrget_world_sizertrur"r#rsync_batch_normapply) r4rHrvrwr$r% need_syncr world_sizes r9rxzSyncBatchNorm.forwardsy e$ ++E2 == ), &)- & ==T55++7 77  $ $ ) )! ,}}$-043K3K3P3P3R-R*-1]]*  ==K,,4T4;K;Kt;SK &*]]d6N6ND  TX %)MMT5M5MD  SW  !]T]]]&&335]:?:K:K:Z:Z:\  ||  1W1W1Y(ZZ !W$)HH$J$J$L#M"OPP"--3399M!! $ 2 2 **99-HJ"QI<<  *   ;"((  *  r@c|}t|tjjjj rtjj |j|j|j|j|j|}|jr?tj5|j|_ |j|_ddd|j|_|j |_|j"|_|j$|_t'|dr|j(|_|j+D]'\}}|j-||j/||)~|S#1swYxYw)a{Helper function to convert all :attr:`BatchNorm*D` layers in the model to :class:`torch.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 >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) >>> 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. >>> # xdoctest: +SKIP("distributed") >>> 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 = torch.nn.SyncBatchNorm.convert_sync_batchnorm(module, process_group) Nqconfig)rr*nnmodules batchnormrlr rrrrrno_gradr"r#r$r%r&rrhasattrrnamed_children add_moduleconvert_sync_batchnorm)clsmoduler module_outputnamechilds r9rz$SyncBatchNorm.convert_sync_batchnorm s/J fehh..88CC D!HH22##  ** M}}]]_+1==M()/M&%*0)<)rs UU8! /i i X> > B,$OY,$^I*IX mZ FHN*HNVNmZNBIN*INXNmZNBTJTr@