fix syncbn bug for pytorch 1.6 (#355)

encoding/functions/syncbn.py

@@ -103,11 +103,13 @@ class syncbatchnorm_(Function):
 
         # Output
         ctx.save_for_backward(x, _ex, _exs, gamma, beta)
-        return y
+
+        ctx.mark_non_differentiable(running_mean, running_var)
+        return y, running_mean, running_var
 
     @staticmethod
     @once_differentiable
-    def backward(ctx, dz):
+    def backward(ctx, dz, _drunning_mean, _drunning_var):
         x, _ex, _exs, gamma, beta = ctx.saved_tensors
         dz = dz.contiguous()
 
@@ -243,11 +245,13 @@ class inp_syncbatchnorm_(Function):
 
         # Output
         ctx.save_for_backward(x, _ex, _exs, gamma, beta)
-        return x
+
+        ctx.mark_non_differentiable(running_mean, running_var)
+        return x, running_mean, running_var
 
     @staticmethod
     @once_differentiable
-    def backward(ctx, dz):
+    def backward(ctx, dz, _drunning_mean, _drunning_var):
         z, _ex, _exs, gamma, beta = ctx.saved_tensors
         dz = dz.contiguous()
 

encoding/nn/syncbn.py

@@ -193,13 +193,15 @@ class SyncBatchNorm(_BatchNorm):
                 "worker_queue": self.worker_queues[self.worker_ids.index(x.get_device())]
             }
         if self.inplace:
-            return inp_syncbatchnorm(x, self.weight, self.bias, self.running_mean, self.running_var,
+            y, _, _ = inp_syncbatchnorm(x, self.weight, self.bias, self.running_mean, self.running_var,
                                      extra, self.sync, self.training, self.momentum, self.eps,
-                                     self.activation, self.slope).view(input_shape)
+                                     self.activation, self.slope)
+            return y.view(input_shape)
         else:
-            return syncbatchnorm(x, self.weight, self.bias, self.running_mean, self.running_var,
+            y, _, _  = syncbatchnorm(x, self.weight, self.bias, self.running_mean, self.running_var,
                                  extra, self.sync, self.training, self.momentum, self.eps,
-                                 self.activation, self.slope).view(input_shape)
+                                 self.activation, self.slope)
+            return y.view(input_shape)
 
     def extra_repr(self):
         if self.activation == 'none':