in mobile_cv/arch/layers/batch_norm.py [0:0]
def forward(self, input):
if get_world_size() == 1 or not self.training:
return super().forward(input)
B, C = input.shape[0], input.shape[1]
half_input = input.dtype == torch.float16
if half_input:
# fp16 does not have good enough numerics for the reduction here
input = input.float()
mean = torch.mean(input, dim=[0, 2, 3])
meansqr = torch.mean(input * input, dim=[0, 2, 3])
if self._stats_mode == "":
assert (
B > 0
), 'SyncBatchNorm(stats_mode="") does not support zero batch size.'
vec = torch.cat([mean, meansqr], dim=0)
vec = differentiable_all_reduce(vec) * (1.0 / dist.get_world_size())
mean, meansqr = torch.split(vec, C)
momentum = self.momentum
else:
if B == 0:
vec = torch.zeros([2 * C + 1], device=mean.device, dtype=mean.dtype)
vec = vec + input.sum() # make sure there is gradient w.r.t input
else:
vec = torch.cat(
[
mean,
meansqr,
torch.ones([1], device=mean.device, dtype=mean.dtype),
],
dim=0,
)
vec = differentiable_all_reduce(vec * B)
total_batch = vec[-1].detach()
momentum = (
total_batch.clamp(max=1) * self.momentum
) # no update if total_batch is 0
mean, meansqr, _ = torch.split(
vec / total_batch.clamp(min=1), C
) # avoid div-by-zero
var = meansqr - mean * mean
invstd = torch.rsqrt(var + self.eps)
scale = self.weight * invstd
bias = self.bias - mean * scale
scale = scale.reshape(1, -1, 1, 1)
bias = bias.reshape(1, -1, 1, 1)
self.running_mean += momentum * (mean.detach() - self.running_mean)
self.running_var += momentum * (var.detach() - self.running_var)
ret = input * scale + bias
if half_input:
ret = ret.half()
return ret