Merge pull request #52 from KevinMusgrave/dev

v0.0.61

Author: KevinMusgrave

src/pytorch_adapt/__init__.py

@@ -1 +1 @@
-__version__ = "0.0.60"
+__version__ = "0.0.61"

src/pytorch_adapt/hooks/adabn.py

@@ -18,7 +18,8 @@ def add_if_new(
             inputs,
             model_name,
             in_keys,
-            other_args=[domain],
+            other_args={"domain": domain},
+            logger=self.logger,
         )
 
 

src/pytorch_adapt/hooks/base.py

@@ -5,6 +5,7 @@
 import torch
 
 from ..utils import common_functions as c_f
+from .logger import HookLogger
 
 
 class BaseHook(ABC):
@@ -41,13 +42,16 @@ def __init__(
         self.out_suffix = out_suffix
         self.key_map = c_f.default(key_map, {})
         self.in_keys = []
+        self.logger = HookLogger(c_f.cls_name(self))
 
     def __call__(self, inputs, losses=None):
+        self.logger("__call__")
         losses = c_f.default(losses, {})
         try:
             inputs = c_f.map_keys(inputs, self.key_map)
             x = self.call(inputs, losses)
             if isinstance(x, (bool, np.bool_)):
+                self.logger.reset()
                 return x
             elif isinstance(x, tuple):
                 outputs, losses = x
@@ -56,14 +60,15 @@ def __call__(self, inputs, losses=None):
                 outputs = wrap_keys(outputs, self.out_prefix, self.out_suffix)
                 losses = wrap_keys(losses, self.loss_prefix, self.loss_suffix)
                 self.check_losses_and_outputs(outputs, losses, inputs)
+                self.logger.reset()
                 return outputs, losses
             else:
                 raise TypeError(
                     f"Output is of type {type(x)}, but should be bool or tuple"
                 )
         except Exception as e:
-            if not isinstance(e, KeyError):
-                c_f.append_error_message(e, self.str_for_error_msg(n=1))
+            c_f.add_error_message(e, f"in {self.logger.str}\n", prepend=True)
+            self.logger.reset()
             raise
 
     @abstractmethod
@@ -124,13 +129,6 @@ def children_repr(self):
         all_modules = c_f.attrs_of_type(self, torch.nn.Module)
         return c_f.assert_dicts_are_disjoint(all_hooks, all_modules)
 
-    def str_for_error_msg(self, x=None, n=None):
-        e = str(self if x is None else x)
-        if n is not None:
-            e = "\n".join(e.split("\n")[:n])
-            e += "\n...\n"
-        return f"\nERROR occuring in:\n{e}"
-
     def check_losses_and_outputs(self, outputs, losses, inputs):
         check_keys_are_present(self, self.loss_keys, [losses], "loss_keys", "losses")
         check_keys_are_present(

src/pytorch_adapt/hooks/domain.py

@@ -100,6 +100,7 @@ def call(self, inputs, losses):
         outputs = self.hook(inputs, losses)[0]
         labels = self.extract_domain_labels(inputs)
         for domain_name, labels in labels.items():
+            self.logger(f"Computing loss for {domain_name} domain")
             [dlogits] = c_f.extract(
                 [outputs, inputs],
                 c_f.filter(self.hook.out_keys, f"_dlogits$", [f"^{domain_name}"]),
@@ -113,6 +114,7 @@ def call(self, inputs, losses):
         return outputs, losses
 
     def extract_domain_labels(self, inputs):
+        self.logger("Expecting 'src_domain' and 'target_domain' in inputs")
         [src_domain, target_domain] = c_f.extract(
             inputs, ["src_domain", "target_domain"]
         )

src/pytorch_adapt/hooks/features.py

@@ -73,6 +73,7 @@ def call(self, inputs, losses):
         """"""
         outputs = {}
         for domain in self.domains:
+            self.logger(f"Getting {domain}")
             detach = self.check_grad_mode(domain)
             func = self.mode_detached if detach else self.mode_with_grad
             in_keys = c_f.filter(self.in_keys, f"^{domain}")
@@ -129,7 +130,15 @@ def mode_detached(self, inputs, outputs, domain, in_keys):
     def add_if_new(
         self, outputs, full_key, output_vals, inputs, model_name, in_keys, domain
     ):
-        c_f.add_if_new(outputs, full_key, output_vals, inputs, model_name, in_keys)
+        c_f.add_if_new(
+            outputs,
+            full_key,
+            output_vals,
+            inputs,
+            model_name,
+            in_keys,
+            logger=self.logger,
+        )
 
     def create_keys(self, domain, suffix, starting_keys=None, detach=False):
         if starting_keys is None:

src/pytorch_adapt/hooks/gvb.py

@@ -30,7 +30,8 @@ def add_if_new(
             inputs,
             model_name,
             in_keys,
-            other_args=[True],
+            other_args={"return_bridge": True},
+            logger=self.logger,
         )
 
 

src/pytorch_adapt/hooks/logger.py

@@ -0,0 +1,12 @@
+class HookLogger:
+    def __init__(self, name):
+        self.name = name
+        self.reset()
+
+    def __call__(self, x):
+        if self.str:
+            self.str += "\n"
+        self.str += f"{self.name}: {x}"
+
+    def reset(self):
+        self.str = ""

src/pytorch_adapt/utils/common_functions.py

@@ -1,5 +1,6 @@
 import errno
 import glob
+import inspect
 import itertools
 import json
 import logging
@@ -57,27 +58,45 @@ def dict_pop_lazy(x, key, *args, **kwargs):
     return default(y, *args, **kwargs)
 
 
-def add_if_new(d, key, x, kwargs, model_name, in_keys, other_args=None):
-    # if key is list then assume model returns multiple args
+def try_use_model(model, model_name, input_vals):
+    try:
+        return default(None, model, input_vals, is_none)
+    except TypeError as e:
+        add_error_message(
+            e,
+            f"\n{model_name}.forward() signature is {inspect.signature(model.forward)}",
+        )
+        raise
+
+
+def assign_to_output(d, key, x, new_x, condition):
+    if len(x) > 1:
+        if not is_list_or_tuple(new_x) or len(new_x) != len(x):
+            raise TypeError(
+                "if input x and key are lists, then output of model must be a list of the same length"
+            )
+        for i in range(len(x)):
+            if condition(x[i]):
+                d[key[i]] = new_x[i]
+    else:
+        d[key[0]] = new_x
+
+
+def add_if_new(d, key, x, kwargs, model_name, in_keys, other_args=None, logger=None):
+    other_args = default(other_args, {})
+    if logger:
+        logger(f"Getting output: {key}")
+        logger(
+            f"Using model {model_name} with inputs: {', '.join(in_keys + list(other_args.keys()))}"
+        )
     if not is_list_or_tuple(key) or not is_list_or_tuple(x):
         raise TypeError("key and x must both be a list or tuple")
     condition = is_none
     if any(condition(y) for y in x):
         model = kwargs[model_name]
-        input_vals = [kwargs[k] for k in in_keys]
-        if other_args is not None:
-            input_vals += other_args
-        new_x = default(None, model, input_vals, is_none)
-        if len(x) > 1:
-            if not is_list_or_tuple(new_x) or len(new_x) != len(x):
-                raise TypeError(
-                    "if input x and key are lists, then output of model must be a list of the same length"
-                )
-            for i in range(len(x)):
-                if condition(x[i]):
-                    d[key[i]] = new_x[i]
-        else:
-            d[key[0]] = new_x
+        input_vals = [kwargs[k] for k in in_keys] + list(other_args.values())
+        new_x = try_use_model(model, model_name, input_vals)
+        assign_to_output(d, key, x, new_x, condition)
 
 
 def class_default(cls, x, default):
@@ -399,9 +418,19 @@ def attrs_of_type(cls, obj):
     return {k: v for k, v in attrs.items() if isinstance(v, obj)}
 
 
-def append_error_message(e, msg):
+# https://stackoverflow.com/a/70114007/16941290
+class ErrorMsgWithNewLines(str):
+    def __repr__(self):
+        return str(self)
+
+
+def add_error_message(e, msg, prepend=False):
     if len(e.args) >= 1:
-        e.args = (e.args[0] + msg,) + e.args[1:]
+        if prepend:
+            x = (ErrorMsgWithNewLines(msg + e.args[0]),)
+        else:
+            x = (ErrorMsgWithNewLines(e.args[0] + msg),)
+        e.args = x + e.args[1:]
 
 
 def requires_grad(x, does=True):

src/pytorch_adapt/weighters/base_weighter.py

@@ -19,7 +19,7 @@ def weight_losses(reduction, weights, scale, loss_dict):
             losses.append(loss)
             components[k] = loss.item()
         except Exception as e:
-            c_f.append_error_message(e, f"\nError occuring with loss key = {k}")
+            c_f.add_error_message(e, f"\nError occuring with loss key = {k}")
             raise
 
     total = reduction(losses)

tests/hooks/test_logger.py

@@ -0,0 +1,80 @@
+import unittest
+
+import torch
+
+from pytorch_adapt.hooks import DANNHook, GVBHook
+from pytorch_adapt.hooks.logger import HookLogger
+
+
+class TestHookLogger(unittest.TestCase):
+    def test_basic(self):
+        test_str = "test test"
+        correct_out = f"{__name__}: {test_str}"
+        logger = HookLogger(__name__)
+        logger(test_str)
+        self.assertTrue(logger.str == correct_out)
+        logger.reset()
+        self.assertTrue(logger.str == "")
+        logger(test_str)
+        self.assertTrue(logger.str == correct_out)
+
+    def test_key_errors_dann(self):
+        hook = DANNHook(opts=[])
+        data = {
+            "src_imgs": torch.randn(32, 32),
+        }
+        models = {
+            "G": torch.nn.Linear(32, 10),
+            "C": torch.nn.Linear(10, 2),
+            "D": torch.nn.Sequential(torch.nn.Linear(10, 1), torch.nn.Flatten(0)),
+        }
+        with self.assertRaises(KeyError) as cm:
+            hook({**data, **models})
+
+        correct_str = (
+            "in DANNHook: __call__"
+            "\nin ChainHook: __call__"
+            "\nin OptimizerHook: __call__"
+            "\nin ChainHook: __call__"
+            "\nin FeaturesForDomainLossHook: __call__"
+            "\nin FeaturesHook: __call__"
+            "\nFeaturesHook: Getting src"
+            "\nFeaturesHook: Getting output: ['src_imgs_features']"
+            "\nFeaturesHook: Using model G with inputs: src_imgs"
+            "\nFeaturesHook: Getting target"
+            "\nFeaturesHook: Getting output: ['target_imgs_features']"
+            "\nFeaturesHook: Using model G with inputs: target_imgs"
+            "\ntarget_imgs"
+        )
+
+        self.assertTrue(str(cm.exception) == correct_str)
+        self.assertTrue(hook.logger.str == "")
+
+    def test_type_error_gvb(self):
+        hook = GVBHook(opts=[])
+        data = {"src_imgs": torch.randn(32, 32), "target_imgs": torch.randn(32, 32)}
+        models = {
+            "G": torch.nn.Linear(32, 10),
+            "C": torch.nn.Linear(10, 2),
+            "D": torch.nn.Sequential(torch.nn.Linear(10, 1), torch.nn.Flatten(0)),
+        }
+
+        with self.assertRaises(TypeError) as cm:
+            hook({**data, **models})
+
+        correct_str = (
+            "in GVBHook: __call__"
+            "\nin ChainHook: __call__"
+            "\nin OptimizerHook: __call__"
+            "\nin ChainHook: __call__"
+            "\nin FeaturesLogitsAndGBridge: __call__"
+            "\nin GBridgeAndLogitsHook: __call__"
+            "\nGBridgeAndLogitsHook: Getting src"
+            "\nGBridgeAndLogitsHook: Getting output: ['src_imgs_features_logits', 'src_imgs_features_gbridge']"
+            "\nGBridgeAndLogitsHook: Using model C with inputs: src_imgs_features, return_bridge"
+            "\nforward() takes 2 positional arguments but 3 were given"
+            "\nC.forward() signature is (input: torch.Tensor) -> torch.Tensor"
+        )
+
+        self.assertTrue(str(cm.exception) == correct_str)
+        self.assertTrue(hook.logger.str == "")