Add kMaX-DeepLab

kmax_deeplab/modeling/pixel_decoder/kmax_pixel_decoder.py

@@ -1,6 +1,7 @@
 # Reference: https://github.com/google-research/deeplab2/blob/main/model/pixel_decoder/kmax.py
 # Modified by Qihang Yu
 
+from turtle import forward
 from typing import Dict, List
 
 import torch
@@ -28,13 +29,17 @@ def get_activation(name):
     elif name == 'gelu':
         return nn.GELU()
 
+class SyncBNCPU(nn.SyncBatchNorm):
+    def forward(self, input):
+        self.eval()
+        return super().forward(input)
 
 def get_norm(name, channels):
     if name is None or name.lower() == 'none':
         return nn.Identity()
 
     if name.lower() == 'syncbn':
-        return nn.SyncBatchNorm(channels, eps=1e-3, momentum=0.01)
+        return SyncBNCPU(channels, eps=1e-3, momentum=0.01)
 
 
 class ConvBN(nn.Module):