Adjusted ActNorm to work as described in the paper (#167)

* Adjusted ActNorm to work as described in the paper

* Fix off-by-one

* Fix log jacobian computation

FrEIA/modules/invertible_resnet.py

@@ -30,7 +30,8 @@ def __init__(self, dims_in, dims_c=None, init_data: torch.Tensor = None):
 
         self.register_buffer("is_initialized", torch.tensor(False))
 
-        dims = next(iter(dims_in))
+        dims = list(next(iter(dims_in)))
+        dims[1:] = [1] * len(dims[1:])
         self.log_scale = nn.Parameter(torch.empty(1, *dims))
         self.loc = nn.Parameter(torch.empty(1, *dims))
 
@@ -42,9 +43,24 @@ def scale(self):
         return torch.exp(self.log_scale)
 
     def initialize(self, batch: torch.Tensor):
+        if batch.ndim != self.log_scale.ndim:
+            raise ValueError(f"Expected batch of dimension {self.log_scale.ndim}, but got {batch.ndim}.")
+
+        # we draw the mean and std over all dimensions except the channel dimension
+        dims = [0] + list(range(2, batch.ndim))
+
+        loc = torch.mean(batch, dim=dims, keepdim=True)
+        scale = torch.std(batch, dim=dims, keepdim=True)
+
+        # check for zero std
+        if torch.any(torch.isclose(scale, torch.tensor(0.0))):
+            raise ValueError("Failed to initialize ActNorm: One or more channels have zero standard deviation.")
+
+        # slice here to avoid silent device move
+        self.log_scale.data[:] = torch.log(scale)
+        self.loc.data[:] = loc
+
         self.is_initialized.data = torch.tensor(True)
-        self.log_scale.data = torch.log(torch.std(batch, dim=0, keepdim=True))
-        self.loc.data = torch.mean(batch, dim=0, keepdim=True)
 
     def output_dims(self, input_dims):
         assert len(input_dims) == 1, "Can only use one input"
@@ -61,10 +77,10 @@ def forward(self, x, c=None, rev=False, jac=True):
 
         if not rev:
             out = (x - self.loc) / self.scale
-            log_jac_det = -utils.sum_except_batch(self.log_scale)
+            log_jac_det = -utils.sum_except_batch(self.log_scale) * torch.prod(torch.tensor(x.shape[2:]).float())
         else:
             out = self.scale * x + self.loc
-            log_jac_det = utils.sum_except_batch(self.log_scale)
+            log_jac_det = utils.sum_except_batch(self.log_scale) * torch.prod(torch.tensor(x.shape[2:]).float())
 
         return (out,), log_jac_det
 

tests/test_invertible_resnet.py

@@ -36,8 +36,12 @@ def test_conv(self):
         self.assertStandardMoments(y_)
 
     def assertStandardMoments(self, data):
-        self.assertTrue(torch.allclose(torch.mean(data, dim=0), torch.zeros(data.shape[-1]), atol=1e-7))
-        self.assertTrue(torch.allclose(torch.std(data, dim=0), torch.ones(data.shape[-1])))
+        dims = [0] + list(range(2, data.ndim))
+        mean = torch.mean(data, dim=dims)
+        std = torch.std(data, dim=dims)
+
+        self.assertTrue(torch.allclose(mean, torch.zeros_like(mean), atol=1e-7))
+        self.assertTrue(torch.allclose(std,  torch.ones_like(std)))
 
 
 class IResNetTest(unittest.TestCase):