less naive SyncBN
Summary:
Support different #images per worker, and 0 batch size. 5~10% slower than old implementation.
If needed, it should be used by calling the module manually - currently it cannot be controlled by a config.
@@ -169,28 +169,68 @@ class NaiveSyncBatchNorm(BatchNorm2d):
Use this implementation before `nn.SyncBatchNorm` is fixed.
It is slower than `nn.SyncBatchNorm`.
+
+ Note:
+ There isn't a single definition of Sync BatchNorm.
+
+ When ``stats_mode==""``, this module computes overall statistics by using
+ statistics of each worker with equal weight. The result is true statistics
+ of all samples (as if they are all on one worker) only when all workers
+ have the same (N, H, W).
+
+ When ``stats_mode=="N"``, this module computes overall statistics by weighting
+ the statistics of each worker by their ``N``. The result is true statistics
+ of all samples (as if they are all on one worker) only when all workers
+ have the same (H, W). It is slower than ``stats_mode==""``.
+
+ Even though the result of this module may not be the true statistics of all samples,
+ it may still be reasonable because it might be preferrable to assign equal weights
+ to all workers, regardless of their (H, W) dimension, instead of putting larger weight
+ on larger images. From preliminary experiments, little difference is found between such
+ a simplified implementation and an accurate computation of overall mean & variance.
"""
+ def __init__(self, *args, stats_mode="", **kwargs):
+ super().__init__(*args, **kwargs)
+ assert stats_mode in ["", "N"]
+ self._stats_mode = stats_mode
+
def forward(self, input):
if comm.get_world_size() == 1 or not self.training:
return super().forward(input)
- assert input.shape[0] > 0, "SyncBatchNorm does not support empty inputs"
- C = input.shape[1]
+ B, C = input.shape[0], input.shape[1]
+
mean = torch.mean(input, dim=[0, 2, 3])
meansqr = torch.mean(input * input, dim=[0, 2, 3])
+ if self._stats_mode == "":
vec = torch.cat([mean, meansqr], dim=0)
vec = AllReduce.apply(vec) * (1.0 / dist.get_world_size())
-
mean, meansqr = torch.split(vec, C)
- var = meansqr - mean * mean
- self.running_mean += self.momentum * (mean.detach() - self.running_mean)
- self.running_var += self.momentum * (var.detach() - self.running_var)
+ momentum = self.momentum
+ else:
+ if B == 0:
+ vec = torch.zeros([2 * C + 1], device=mean.device, dtype=mean.dtype)
+ else:
+ vec = torch.cat(
+ [mean, meansqr, torch.ones([1], device=mean.device, dtype=mean.dtype)], dim=0
+ )
+ vec = AllReduce.apply(vec * B)
+ total_batch = vec[-1]
+ momentum = (
+ total_batch.detach().clamp(max=1) * self.momentum
+ ) # no update if total_batch is 0
+ mean, meansqr, _ = torch.split(vec / total_batch, C)
+
+ 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)
return input * scale + bias
