Merge pull request #6 from tfjgeorge/deterministic

adds seed for deterministic tests

tests/test_jacobian.py

@@ -45,6 +45,12 @@ def get_output_vector(loader, function):
         return torch.cat(outputs)
 
 
+@pytest.fixture(autouse=True)
+def make_test_deterministic():
+    torch.manual_seed(1234)
+    yield
+
+
 def test_jacobian_pushforward_dense_linear():
     for get_task in linear_tasks:
         loader, lc, parameters, model, function, n_output = get_task()
@@ -84,8 +90,10 @@ def test_jacobian_pushforward_dense_nonlinear():
         update_model(parameters, dw.get_flat_representation())
         output_after = get_output_vector(loader, function)
 
+        # This is non linear, so we don't expect the finite difference
+        # estimate to be very accurate. We use a larger eps value
         check_tensors(output_after - output_before,
-                      doutput_lin.get_flat_representation().t(), eps=1e-2)
+                      doutput_lin.get_flat_representation().t(), eps=5e-2)
 
 
 def test_jacobian_pushforward_implicit():

tests/test_jacobian_ekfac.py

@@ -4,6 +4,13 @@
 from tasks import get_fullyconnect_task, get_conv_task, device
 from nngeometry.object.vector import random_pvector
 from utils import check_ratio, check_tensors
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def make_test_deterministic():
+    torch.manual_seed(1234)
+    yield
 
 
 def test_pspace_ekfac_vs_kfac():

tests/test_jacobian_kfac.py

@@ -11,6 +11,7 @@
 from utils import check_ratio, check_tensors, angle
 from tasks import get_fullyconnect_task, get_mnist, get_conv_task
 import os
+import pytest
 
 default_datapath = 'tmp'
 if 'SLURM_TMPDIR' in os.environ:
@@ -129,6 +130,12 @@ def output_fn(input, target):
             output_fn, 10)
 
 
+@pytest.fixture(autouse=True)
+def make_test_deterministic():
+    torch.manual_seed(1234)
+    yield
+
+
 def test_jacobian_kfac_vs_pblockdiag():
     """
     Compares blockdiag and kfac representation on datasets/architectures

tests/test_metrics.py

@@ -21,8 +21,14 @@
 
 nonlinear_tasks = [get_conv_gn_task, get_fullyconnect_task, get_conv_task]
 
-
 import numpy as np
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def make_test_deterministic():
+    torch.manual_seed(1234)
+    yield
 
 
 def test_FIM_MC_vs_linearization():

tests/test_vector.py

@@ -5,6 +5,7 @@
 import torch.nn as nn
 import torch.nn.functional as tF
 from utils import check_ratio
+import pytest
 
 
 class ConvNet(nn.Module):
@@ -27,6 +28,12 @@ def forward(self, x):
         return tF.log_softmax(x, dim=1)
 
 
+@pytest.fixture(autouse=True)
+def make_test_deterministic():
+    torch.manual_seed(1234)
+    yield
+
+
 def test_from_dict_to_pvector():
     eps = 1e-8
     model = ConvNet()