Add kMaX-DeepLab

kmax_deeplab/modeling/pixel_decoder/kmax_pixel_decoder.py

@@ -31,8 +31,29 @@ def get_activation(name):
 
 class SyncBNCPU(nn.SyncBatchNorm):
     def forward(self, input):
-        self.eval()
-        return super().forward(input)
+        self._check_input_dim(input)
+        self._check_non_zero_input_channels(input)
+        if self.momentum is None:
+            exponential_average_factor = 0.0
+        else:
+            exponential_average_factor = self.momentum
+        bn_training = False
+
+        running_mean = self.running_mean 
+        running_var = self.running_var
+
+        # fallback to framework BN when synchronization is not necessary
+        return F.batch_norm(
+            input,
+            running_mean,
+            running_var,
+            self.weight,
+            self.bias,
+            bn_training,
+            exponential_average_factor,
+            self.eps,
+        )
+
 
 def get_norm(name, channels):
     if name is None or name.lower() == 'none':