add batchnormalization

Author: SkafteNicki

stochman/nnj.py

@@ -266,6 +266,27 @@ def __call__(self, x: Tensor, jacobian: bool = False) -> Union[Tensor, Tuple[Ten
         return val
 
 
+class BatchNorm1d(AbstractActivationJacobian, nn.BatchNorm1d):
+    # only implements jacobian during testing
+    def _jacobian(self, x: Tensor, val: Tensor) -> Tensor:
+        jac = (self.weight / (self.running_var + self.eps).sqrt()).unsqueeze(0)
+        return jac
+
+
+class BatchNorm2d(AbstractActivationJacobian, nn.BatchNorm2d):
+    # only implements jacobian during testing
+    def _jacobian(self, x: Tensor, val: Tensor) -> Tensor:
+        jac = (self.weight / (self.running_var + self.eps).sqrt()).unsqueeze(0)
+        return jac
+
+
+class BatchNorm3d(AbstractActivationJacobian, nn.BatchNorm3d):
+    # only implements jacobian during testing
+    def _jacobian(self, x: Tensor, val: Tensor) -> Tensor:
+        jac = (self.weight / (self.running_var + self.eps).sqrt()).unsqueeze(0)
+        return jac
+
+
 class Sigmoid(AbstractActivationJacobian, nn.Sigmoid):
     def _jacobian(self, x: Tensor, val: Tensor) -> Tensor:
         jac = val * (1.0 - val)
@@ -302,7 +323,7 @@ def _jacobian(self, x: Tensor, val: Tensor) -> Tensor:
 class LeakyReLU(AbstractActivationJacobian, nn.LeakyReLU):
     def _jacobian(self, x: Tensor, val: Tensor) -> Tensor:
         jac = torch.ones_like(val)
-        jac[val < 0.0] = self.negative_slope
+        jac[x < 0.0] = self.negative_slope
         return jac
 
 

tests/test_nnj.py

@@ -11,10 +11,10 @@
 _features = 5
 _dims = 6
 
-_linear_input = torch.randn(_batch_size, _features)
-_1d_conv_input = torch.randn(_batch_size, _features, _dims)
-_2d_conv_input = torch.randn(_batch_size, _features, _dims, _dims)
-_3d_conv_input = torch.randn(_batch_size, _features, _dims, _dims, _dims)
+_linear_input_shape = (_batch_size, _features)
+_1d_conv_input_shape = (_batch_size, _features, _dims)
+_2d_conv_input_shape = (_batch_size, _features, _dims, _dims)
+_3d_conv_input_shape = (_batch_size, _features, _dims, _dims, _dims)
 
 
 def _compare_jacobian(f: Callable, x: torch.Tensor) -> torch.Tensor:
@@ -28,18 +28,16 @@ def _compare_jacobian(f: Callable, x: torch.Tensor) -> torch.Tensor:
 
 
 @pytest.mark.parametrize(
-    "model, input",
+    "model, input_shape",
     [
-        (nnj.Sequential(nnj.Identity(), nnj.Identity()), _linear_input),
-        (nnj.Linear(_features, 2), _linear_input),
-        (nnj.Sequential(nnj.PosLinear(_features, 2), nnj.Reciprocal()), _linear_input),
-        (nnj.Sequential(nnj.Linear(_features, 2), nnj.Sigmoid(), nnj.ArcTanh()), _linear_input),
-        (nnj.Sequential(nnj.Linear(_features, 5), nnj.Sigmoid(), nnj.Linear(5, 2)), _linear_input),
+        (nnj.Sequential(nnj.Identity(), nnj.Identity()), _linear_input_shape),
+        (nnj.Linear(_features, 2), _linear_input_shape),
+        (nnj.Sequential(nnj.PosLinear(_features, 2), nnj.Reciprocal()), _linear_input_shape),
+        (nnj.Sequential(nnj.Linear(_features, 2), nnj.Sigmoid(), nnj.ArcTanh()), _linear_input_shape),
+        (nnj.Sequential(nnj.Linear(_features, 5), nnj.Sigmoid(), nnj.Linear(5, 2)), _linear_input_shape),
         (
-            nnj.Sequential(
-                nnj.Linear(_features, 2), nnj.Softplus(beta=100, threshold=5), nnj.Linear(2, 4), nnj.Tanh()
-            ),
-            _linear_input,
+            nnj.Sequential(nnj.Linear(_features, 2), nnj.Softplus(beta=100, threshold=5), nnj.Linear(2, 4)),
+            _linear_input_shape,
         ),
         (
             nnj.Sequential(
@@ -50,21 +48,31 @@ def _compare_jacobian(f: Callable, x: torch.Tensor) -> torch.Tensor:
                 nnj.Sqrt(),
                 nnj.Hardshrink(),
             ),
-            _linear_input,
+            _linear_input_shape,
+        ),
+        (nnj.Sequential(nnj.Linear(_features, 2), nnj.LeakyReLU()), _linear_input_shape),
+        (nnj.Sequential(nnj.Linear(_features, 2), nnj.Tanh()), _linear_input_shape),
+        (nnj.Sequential(nnj.Linear(_features, 2), nnj.OneMinusX()), _linear_input_shape),
+        (
+            nnj.Sequential(nnj.Conv1d(_features, 2, 5), nnj.ConvTranspose1d(2, _features, 5)),
+            _1d_conv_input_shape,
+        ),
+        (
+            nnj.Sequential(nnj.Conv2d(_features, 2, 5), nnj.ConvTranspose2d(2, _features, 5)),
+            _2d_conv_input_shape,
+        ),
+        (
+            nnj.Sequential(nnj.Conv3d(_features, 2, 5), nnj.ConvTranspose3d(2, _features, 5)),
+            _3d_conv_input_shape,
         ),
-        (nnj.Sequential(nnj.Linear(_features, 2), nnj.LeakyReLU()), _linear_input),
-        (nnj.Sequential(nnj.Linear(_features, 2), nnj.OneMinusX()), _linear_input),
-        (nnj.Sequential(nnj.Conv1d(_features, 2, 5), nnj.ConvTranspose1d(2, _features, 5)), _1d_conv_input),
-        (nnj.Sequential(nnj.Conv2d(_features, 2, 5), nnj.ConvTranspose2d(2, _features, 5)), _2d_conv_input),
-        (nnj.Sequential(nnj.Conv3d(_features, 2, 5), nnj.ConvTranspose3d(2, _features, 5)), _3d_conv_input),
         (
             nnj.Sequential(
                 nnj.Linear(_features, 8),
                 nnj.Sigmoid(),
                 nnj.Reshape(2, 4),
                 nnj.Conv1d(2, 1, 2),
             ),
-            _linear_input,
+            _linear_input_shape,
         ),
         (
             nnj.Sequential(
@@ -73,7 +81,7 @@ def _compare_jacobian(f: Callable, x: torch.Tensor) -> torch.Tensor:
                 nnj.Reshape(2, 4, 4),
                 nnj.Conv2d(2, 1, 2),
             ),
-            _linear_input,
+            _linear_input_shape,
         ),
         (
             nnj.Sequential(
@@ -82,7 +90,7 @@ def _compare_jacobian(f: Callable, x: torch.Tensor) -> torch.Tensor:
                 nnj.Reshape(2, 4, 4, 4),
                 nnj.Conv3d(2, 1, 2),
             ),
-            _linear_input,
+            _linear_input_shape,
         ),
         (
             nnj.Sequential(
@@ -91,7 +99,7 @@ def _compare_jacobian(f: Callable, x: torch.Tensor) -> torch.Tensor:
                 nnj.Linear(4 * 2, 5),
                 nnj.ReLU(),
             ),
-            _1d_conv_input,
+            _1d_conv_input_shape,
         ),
         (
             nnj.Sequential(
@@ -100,7 +108,7 @@ def _compare_jacobian(f: Callable, x: torch.Tensor) -> torch.Tensor:
                 nnj.Linear(4 * 4 * 2, 5),
                 nnj.ReLU(),
             ),
-            _2d_conv_input,
+            _2d_conv_input_shape,
         ),
         (
             nnj.Sequential(
@@ -109,30 +117,34 @@ def _compare_jacobian(f: Callable, x: torch.Tensor) -> torch.Tensor:
                 nnj.Linear(4 * 4 * 4 * 2, 5),
                 nnj.ReLU(),
             ),
-            _3d_conv_input,
+            _3d_conv_input_shape,
         ),
         (
             nnj.Sequential(nnj.Conv2d(_features, 2, 3), nnj.Hardtanh(), nnj.Upsample(scale_factor=2)),
-            _2d_conv_input,
+            _2d_conv_input_shape,
         ),
+        (nnj.Sequential(nnj.Conv1d(_features, 3, 3), nnj.BatchNorm1d(3)), _1d_conv_input_shape),
+        (nnj.Sequential(nnj.Conv2d(_features, 3, 3), nnj.BatchNorm2d(3)), _2d_conv_input_shape),
+        (nnj.Sequential(nnj.Conv3d(_features, 3, 3), nnj.BatchNorm3d(3)), _3d_conv_input_shape),
     ],
 )
 class TestJacobian:
     @pytest.mark.parametrize("dtype", [torch.float, torch.double])
-    def test_jacobians(self, model, input, dtype):
+    def test_jacobians(self, model, input_shape, dtype):
         """Test that the analytical jacobian of the model is consistent with finite
         order approximation
         """
-        model = deepcopy(model).to(dtype)
-        input = deepcopy(input).to(dtype)
+        model = deepcopy(model).to(dtype).eval()
+        input = torch.randn(*input_shape, dtype=dtype)
         _, jac = model(input, jacobian=True)
         jacnum = _compare_jacobian(model, input)
         assert torch.isclose(jac, jacnum, atol=1e-7).all(), "jacobians did not match"
 
     @pytest.mark.parametrize("return_jac", [True, False])
-    def test_jac_return(self, model, input, return_jac):
+    def test_jac_return(self, model, input_shape, return_jac):
         """ Test that all models returns the jacobian output if asked for it """
-        output = model(input, jacobian=return_jac)
+
+        output = model(torch.randn(*input_shape), jacobian=return_jac)
         if return_jac:
             assert len(output) == 2, "expected two outputs when jacobian=True"
             assert all(