text: ######################################################################
  #  script name  : ConditionalBatchNorm2d.py
  #  author       : Chen Xuanhong
  #  created time : 2019/9/12 00:07
  #  modification time ：2019/9/12 00:07
  #  modified by  : Chen Xuanhong
  ######################################################################
  
  import torch.nn as nn
  import torch.nn.functional as F
  from torch.nn import init
  
  
  class ConditionalBatchNorm2d(nn.BatchNorm2d):
  
      """Conditional Batch Normalization"""
  
      def __init__(self, num_features, eps=1e-05, momentum=0.1,
                   affine=False, track_running_stats=True):
          super(ConditionalBatchNorm2d, self).__init__(
              num_features, eps, momentum, affine, track_running_stats
          )
  
      def forward(self, input, weight, bias, **kwargs):
          self._check_input_dim(input)
  
          exponential_average_factor = 0.0
  
          if self.training and self.track_running_stats:
              self.num_batches_tracked += 1
              if self.momentum is None:  # use cumulative moving average
                  exponential_average_factor = 1.0 / self.num_batches_tracked.item()
              else:  # use exponential moving average
                  exponential_average_factor = self.momentum
  
          output = F.batch_norm(input, self.running_mean, self.running_var,
                                self.weight, self.bias,
                                self.training or not self.track_running_stats,
                                exponential_average_factor, self.eps)
          if weight.dim() == 1:
              weight = weight.unsqueeze(0)
          if bias.dim() == 1:
              bias = bias.unsqueeze(0)
          size = output.size()
          weight = weight.unsqueeze(-1).unsqueeze(-1).expand(size)
          bias = bias.unsqueeze(-1).unsqueeze(-1).expand(size)
          return weight * output + bias