Merge pull request #49 from KevinMusgrave/dev

Added test_equivalent_adapter to hook tests

Author: KevinMusgrave

src/pytorch_adapt/adapters/adabn.py

@@ -8,9 +8,9 @@
 class AdaBN(BaseAdapter):
     hook_cls = AdaBNHook
 
-    def __init__(self, inference_fn=None, **kwargs):
+    def __init__(self, *args, inference_fn=None, **kwargs):
         inference_fn = c_f.default(inference_fn, adabn_fn)
-        super().__init__(inference_fn=inference_fn, **kwargs)
+        super().__init__(*args, inference_fn=inference_fn, **kwargs)
 
     def init_hook(self, hook_kwargs):
         self.hook = self.hook_cls(**hook_kwargs)

src/pytorch_adapt/adapters/adda.py

@@ -23,9 +23,9 @@ class ADDA(BaseAdapter):
 
     hook_cls = ADDAHook
 
-    def __init__(self, inference_fn=None, **kwargs):
+    def __init__(self, *args, inference_fn=None, **kwargs):
         inference_fn = c_f.default(inference_fn, adda_fn)
-        super().__init__(inference_fn=inference_fn, **kwargs)
+        super().__init__(*args, inference_fn=inference_fn, **kwargs)
 
     def get_default_containers(self) -> MultipleContainers:
         """

src/pytorch_adapt/adapters/aligner.py

@@ -38,9 +38,9 @@ class RTN(Aligner):
 
     hook_cls = RTNHook
 
-    def __init__(self, inference_fn=None, **kwargs):
+    def __init__(self, *args, inference_fn=None, **kwargs):
         inference_fn = c_f.default(inference_fn, rtn_fn)
-        super().__init__(inference_fn=inference_fn, **kwargs)
+        super().__init__(*args, inference_fn=inference_fn, **kwargs)
 
     def get_key_enforcer(self) -> KeyEnforcer:
         ke = super().get_key_enforcer()

src/pytorch_adapt/adapters/mcd.py

@@ -22,9 +22,9 @@ class MCD(BaseGCAdapter):
 
     hook_cls = MCDHook
 
-    def __init__(self, inference_fn=None, **kwargs):
+    def __init__(self, *args, inference_fn=None, **kwargs):
         inference_fn = c_f.default(inference_fn, mcd_fn)
-        super().__init__(inference_fn=inference_fn, **kwargs)
+        super().__init__(*args, inference_fn=inference_fn, **kwargs)
 
     def init_hook(self, hook_kwargs):
         self.hook = self.hook_cls(

src/pytorch_adapt/adapters/symnets.py

@@ -20,9 +20,9 @@ class SymNets(BaseGCAdapter):
 
     hook_cls = SymNetsHook
 
-    def __init__(self, inference_fn=None, **kwargs):
+    def __init__(self, *args, inference_fn=None, **kwargs):
         inference_fn = c_f.default(inference_fn, symnets_fn)
-        super().__init__(inference_fn=inference_fn, **kwargs)
+        super().__init__(*args, inference_fn=inference_fn, **kwargs)
 
     def init_hook(self, hook_kwargs):
         self.hook = self.hook_cls(

tests/hooks/test_adabn.py

@@ -3,13 +3,22 @@
 
 import torch
 
+from pytorch_adapt.adapters import AdaBN
+from pytorch_adapt.containers import Models
 from pytorch_adapt.hooks import AdaBNHook, validate_hook
 from pytorch_adapt.layers import AdaptiveBatchNorm2d
 from pytorch_adapt.layers.adaptive_batch_norm import set_curr_domain
 
 from .utils import assertRequiresGrad
 
 
+def test_equivalent_adapter(G, C, data):
+    models = Models({"G": copy.deepcopy(G), "C": copy.deepcopy(C)})
+    adapter = AdaBN(models)
+    adapter.training_step(data)
+    return models
+
+
 class Net(torch.nn.Module):
     def __init__(self, in_size, out_size):
         super().__init__()
@@ -53,17 +62,22 @@ def test_adabn_hook(self):
         assertRequiresGrad(self, outputs)
         self.assertTrue(len(losses) == 0)
 
+        adapter_models = test_equivalent_adapter(originalG, originalC, data)
+
         originalG.net[0].bns[0](src_imgs)
         originalG.net[0].bns[1](target_imgs)
 
         for i in range(2):
-            self.assertTrue(
-                torch.equal(
-                    G.net[0].bns[i].running_mean, originalG.net[0].bns[i].running_mean
+            for M in [models, adapter_models]:
+                self.assertTrue(
+                    torch.equal(
+                        M["G"].net[0].bns[i].running_mean,
+                        originalG.net[0].bns[i].running_mean,
+                    )
                 )
-            )
-            self.assertTrue(
-                torch.equal(
-                    G.net[0].bns[i].running_var, originalG.net[0].bns[i].running_var
+                self.assertTrue(
+                    torch.equal(
+                        M["G"].net[0].bns[i].running_var,
+                        originalG.net[0].bns[i].running_var,
+                    )
                 )
-            )

tests/hooks/test_adda.py

@@ -5,18 +5,34 @@
 import torch
 import torch.nn.functional as F
 
+from pytorch_adapt.adapters import ADDA
+from pytorch_adapt.containers import Models, Optimizers
 from pytorch_adapt.hooks import ADDAHook, BSPHook, validate_hook
 from pytorch_adapt.utils import common_functions as c_f
 
 from .utils import (
     Net,
+    assert_equal_models,
     assertRequiresGrad,
+    get_opt_tuple,
     get_opts,
     post_g_hook_update_keys,
     post_g_hook_update_total_loss,
 )
 
 
+def test_equivalent_adapter(G, D, data, post_g, threshold):
+    models = Models(
+        {"G": copy.deepcopy(G), "D": copy.deepcopy(D), "C": torch.nn.Identity()}
+    )
+    optimizers = Optimizers(get_opt_tuple())
+    adapter = ADDA(
+        models, optimizers, hook_kwargs={"post_g": post_g, "threshold": threshold}
+    )
+    adapter.training_step(data)
+    return models
+
+
 def get_models_and_data():
     src_domain = torch.randint(0, 2, size=(100,)).float()
     target_domain = torch.randint(0, 2, size=(100,)).float()
@@ -30,7 +46,7 @@ def get_models_and_data():
 class TestADDA(unittest.TestCase):
     def test_adda(self):
         torch.manual_seed(922)
-        for post_g in [None, BSPHook(domains=["target"])]:
+        for post_g in [None, [BSPHook(domains=["target"])]]:
             for threshold in np.linspace(0, 1, 10):
                 (
                     G,
@@ -46,9 +62,8 @@ def test_adda(self):
                 originalT = copy.deepcopy(T)
                 d_opts = get_opts(D)
                 g_opts = get_opts(T)
-                post_g_ = [post_g] if post_g is not None else post_g
                 h = ADDAHook(
-                    d_opts=d_opts, g_opts=g_opts, threshold=threshold, post_g=post_g_
+                    d_opts=d_opts, g_opts=g_opts, threshold=threshold, post_g=post_g
                 )
                 models = {"G": G, "D": D, "T": T}
                 data = {
@@ -78,6 +93,10 @@ def test_adda(self):
                 )
                 self.assertTrue(losses["g_loss"].keys() == g_loss_keys)
 
+                adapter_models = test_equivalent_adapter(
+                    originalG, originalD, data, post_g, threshold
+                )
+
                 d_opts = get_opts(originalD)[0]
                 g_opts = get_opts(originalT)[0]
                 originalG.eval()
@@ -142,9 +161,12 @@ def test_adda(self):
                 # can't use model_counts for conditional part
                 self.assertTrue(D.count == d_count)
 
-                for x, y in [(G, originalG), (T, originalT), (D, originalD)]:
-                    self.assertTrue(
-                        c_f.state_dicts_are_equal(
-                            x.state_dict(), y.state_dict(), rtol=1e-3
-                        )
-                    )
+                assert_equal_models(
+                    self, (G, adapter_models["G"], originalG), rtol=1e-3
+                )
+                assert_equal_models(
+                    self, (T, adapter_models["T"], originalT), rtol=1e-3
+                )
+                assert_equal_models(
+                    self, (D, adapter_models["D"], originalD), rtol=1e-3
+                )

tests/hooks/test_aligners.py

@@ -5,11 +5,30 @@
 import torch
 import torch.nn.functional as F
 
+from pytorch_adapt.adapters import Aligner
+from pytorch_adapt.containers import Models, Optimizers
 from pytorch_adapt.hooks import AlignerPlusCHook, JointAlignerHook, validate_hook
 from pytorch_adapt.layers import CORALLoss, MMDLoss
-from pytorch_adapt.utils import common_functions as c_f
 
-from .utils import assertRequiresGrad, get_models_and_data, get_opts
+from .utils import (
+    assert_equal_models,
+    assertRequiresGrad,
+    get_models_and_data,
+    get_opt_tuple,
+    get_opts,
+)
+
+
+def test_equivalent_adapter(G, C, data, aligner_hook, loss_fn):
+    models = Models({"G": copy.deepcopy(G), "C": copy.deepcopy(C)})
+    optimizers = Optimizers(get_opt_tuple())
+    adapter = Aligner(
+        models,
+        optimizers,
+        hook_kwargs={"aligner_hook": aligner_hook, "loss_fn": loss_fn},
+    )
+    adapter.training_step(data)
+    return models
 
 
 class TestAligners(unittest.TestCase):
@@ -74,6 +93,10 @@ def test_aligner_plus_classifier_hook(self):
                     G.count == model_counts["G"] == C.count == model_counts["C"] == 2
                 )
 
+                adapter_models = test_equivalent_adapter(
+                    originalG, originalC, data, aligner_hook, loss_fn()
+                )
+
                 opts = get_opts(originalG, originalC)
 
                 src_features = originalG(src_imgs)
@@ -120,9 +143,9 @@ def test_aligner_plus_classifier_hook(self):
                 total_loss.backward()
                 [x.step() for x in opts]
 
-                for x, y in [(G, originalG), (C, originalC)]:
-                    self.assertTrue(
-                        c_f.state_dicts_are_equal(
-                            x.state_dict(), y.state_dict(), rtol=1e-6
-                        )
-                    )
+                assert_equal_models(
+                    self, (G, adapter_models["G"], originalG), rtol=1e-6
+                )
+                assert_equal_models(
+                    self, (C, adapter_models["C"], originalC), rtol=1e-6
+                )

tests/hooks/test_cdan.py

@@ -5,6 +5,8 @@
 import numpy as np
 import torch
 
+from pytorch_adapt.adapters import CDANE
+from pytorch_adapt.containers import Misc, Models, Optimizers
 from pytorch_adapt.hooks import (
     AFNHook,
     BNMHook,
@@ -16,18 +18,41 @@
     validate_hook,
 )
 from pytorch_adapt.layers import RandomizedDotProduct
-from pytorch_adapt.utils import common_functions as c_f
 
 from .utils import (
+    assert_equal_models,
     assertRequiresGrad,
     get_entropy_weights,
     get_models_and_data,
+    get_opt_tuple,
     get_opts,
     post_g_hook_update_keys,
     post_g_hook_update_total_loss,
 )
 
 
+def test_equivalent_adapter(
+    G, D, C, feature_combiner, data, detach_reducer, post_g, softmax
+):
+    models = Models(
+        {"G": copy.deepcopy(G), "D": copy.deepcopy(D), "C": copy.deepcopy(C)}
+    )
+    misc = Misc({"feature_combiner": copy.deepcopy(feature_combiner)})
+    optimizers = Optimizers(get_opt_tuple())
+    adapter = CDANE(
+        models,
+        optimizers,
+        misc=misc,
+        hook_kwargs={
+            "detach_entropy_reducer": detach_reducer,
+            "post_g": post_g,
+            "softmax": softmax,
+        },
+    )
+    adapter.training_step(data)
+    return models
+
+
 def get_correct_domain_losses(
     G,
     C,
@@ -159,7 +184,7 @@ def test_cdan_domain_hooks(self):
     def test_cdan_hook(self):
         torch.manual_seed(985)
         for detach_reducer in [False, True]:
-            for post_g in [None, BSPHook(), BNMHook(), MCCHook(), AFNHook()]:
+            for post_g in [None, [BSPHook()], [BNMHook()], [MCCHook()], [AFNHook()]]:
                 softmax = True
                 fc_out_size = 16
                 (
@@ -189,13 +214,12 @@ def test_cdan_hook(self):
                     "src_domain": src_domain,
                     "target_domain": target_domain,
                 }
-                post_g_ = [post_g] if post_g is not None else post_g
                 hook = CDANEHook(
                     detach_entropy_reducer=detach_reducer,
                     d_opts=d_opts,
                     g_opts=g_opts,
                     softmax=softmax,
-                    post_g=post_g_,
+                    post_g=post_g,
                 )
                 model_counts = validate_hook(hook, list(data.keys()))
                 outputs, losses = hook({**models, **data})
@@ -236,6 +260,17 @@ def test_cdan_hook(self):
                 )
                 self.assertTrue(D.count == model_counts["D"] == 4)
 
+                adapter_models = test_equivalent_adapter(
+                    originalG,
+                    originalD,
+                    originalC,
+                    originalFeatureCombiner,
+                    data,
+                    detach_reducer,
+                    post_g,
+                    softmax,
+                )
+
                 d_opts = get_opts(originalD)
                 g_opts = get_opts(originalG, originalC)
 
@@ -338,9 +373,12 @@ def test_cdan_hook(self):
                 total_loss.backward()
                 [x.step() for x in g_opts]
 
-                for x, y in [(G, originalG), (C, originalC), (D, originalD)]:
-                    self.assertTrue(
-                        c_f.state_dicts_are_equal(
-                            x.state_dict(), y.state_dict(), rtol=1e-2
-                        )
-                    )
+                assert_equal_models(
+                    self, (G, adapter_models["G"], originalG), rtol=1e-2
+                )
+                assert_equal_models(
+                    self, (C, adapter_models["C"], originalC), rtol=1e-2
+                )
+                assert_equal_models(
+                    self, (D, adapter_models["D"], originalD), rtol=1e-2
+                )