(self, input)
| 641 | ) |
| 642 | |
| 643 | def forward(self, input): |
| 644 | # currently only GPU input is supported |
| 645 | if not input.is_cuda: |
| 646 | raise ValueError("SyncBatchNorm expected input tensor to be on GPU") |
| 647 | |
| 648 | self._check_input_dim(input) |
| 649 | self._check_non_zero_input_channels(input) |
| 650 | |
| 651 | if self.momentum is None: |
| 652 | exponential_average_factor = 0.0 |
| 653 | else: |
| 654 | exponential_average_factor = self.momentum |
| 655 | |
| 656 | if self.training and self.track_running_stats: |
| 657 | assert self.num_batches_tracked is not None |
| 658 | self.num_batches_tracked.add_(1) |
| 659 | if self.momentum is None: # use cumulative moving average |
| 660 | exponential_average_factor = 1.0 / self.num_batches_tracked.item() |
| 661 | else: # use exponential moving average |
| 662 | exponential_average_factor = self.momentum |
| 663 | |
| 664 | r""" |
| 665 | Decide whether the mini-batch stats should be used for normalization rather than the buffers. |
| 666 | Mini-batch stats are used in training mode, and in eval mode when buffers are None. |
| 667 | """ |
| 668 | if self.training: |
| 669 | bn_training = True |
| 670 | else: |
| 671 | bn_training = (self.running_mean is None) and (self.running_var is None) |
| 672 | |
| 673 | # Don't sync batchnorm stats in inference mode (model.eval()). |
| 674 | need_sync = bn_training and self.training |
| 675 | if need_sync: |
| 676 | need_sync = flow.env.get_world_size() > 1 |
| 677 | |
| 678 | # # fallback to framework BN when synchronization is not necessary |
| 679 | if not need_sync: |
| 680 | return flow._C.normalization( |
| 681 | input, |
| 682 | self.running_mean, |
| 683 | self.running_var, |
| 684 | self.weight, |
| 685 | self.bias, |
| 686 | axis=self.channel_axis, |
| 687 | epsilon=self.eps, |
| 688 | momentum=exponential_average_factor, |
| 689 | is_training=bn_training, |
| 690 | ) |
| 691 | else: |
| 692 | assert bn_training |
| 693 | global global_eps |
| 694 | global global_momentum |
| 695 | global global_world_size |
| 696 | global global_axis |
| 697 | global_eps = self.eps |
| 698 | global_momentum = exponential_average_factor |
| 699 | global_world_size = flow.env.get_world_size() |
| 700 | global_axis = self.channel_axis |
nothing calls this directly
no test coverage detected