fix metaclass issue in patch.py & now compatible with PyTorch 1.9.0

Author: KiddoZhu

torchdrug/core/__init__.py

@@ -1,4 +1,4 @@
-from .core import _MetaContainer, Registry, Configurable
+from .core import _MetaContainer, Registry, Configurable, make_configurable
 from .engine import Engine
 from .meter import Meter
 

torchdrug/core/core.py

@@ -187,7 +187,6 @@ def wrapper(obj):
             for key in keys[:-1]:
                 entry = entry[key]
             if keys[-1] in entry:
-                return obj
                 raise KeyError("`%s` has already been registered by %s" % (name, entry[keys[-1]]))
 
             entry[keys[-1]] = obj
@@ -304,7 +303,7 @@ def get_function(method):
 class Configurable(metaclass=_Configurable):
     """
     Class for load/save configuration.
-    It will automatically record every argument passed to ``__init__`` function.
+    It will automatically record every argument passed to the ``__init__`` function.
 
     This class is inspired by :meth:`state_dict()` in PyTorch, but designed for hyperparameters.
 
@@ -334,3 +333,26 @@ class Configurable(metaclass=_Configurable):
     >>> gcn = Configurable.load_config_dict(config)
     """
     pass
+
+
+def make_configurable(cls, module=None, ignore_args=()):
+    """
+    Make a configurable class out of an existing class.
+    The configurable class will automatically record every argument passed to its ``__init__`` function.
+
+    Parameters:
+        cls (type): input class
+        module (str, optional): bind the output class to this module.
+            By default, bind to the original module of the input class.
+        ignore_args (set of str, optional): arguments to ignore in the ``__init__`` function
+    """
+    ignore_args = set(ignore_args)
+    module = module or cls.__module__
+    Metaclass = type(cls)
+    if issubclass(Metaclass, _Configurable): # already a configurable class
+        return cls
+    if Metaclass != type: # already have a meta class
+        MetaClass = type(_Configurable.__name__, (Metaclass, _Configurable), {})
+    else:
+        MetaClass = _Configurable
+    return MetaClass(cls.__name__, (cls,), {"_ignore_args": ignore_args, "__module__": module})

torchdrug/patch.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import inspect
 
 import torch
@@ -12,16 +13,14 @@
 from torchdrug import core, data
 from torchdrug.core import Registry as R
 
-
-def make_configurable(cls, ignore_args=()):
-    ignore_args = set(ignore_args)
-    return type(cls.__name__, (cls, core.Configurable), {"_ignore_args": ignore_args})
+module = sys.modules[__name__]
 
 
 class PatchedModule(nn.Module):
 
     def __init__(self):
         super(PatchedModule, self).__init__()
+        # TODO: these hooks are bugged.
         # self._register_state_dict_hook(PatchedModule.graph_state_dict)
         # self._register_load_state_dict_pre_hook(PatchedModule.load_graph_state_dict)
 
@@ -54,7 +53,6 @@ def load_graph_state_dict(cls, state_dict, prefix, local_metadata, strict, missi
                 print("successfully assigned %s" % (prefix + name))
             except:
                 error_msgs.append("Can't construct Graph `%s` from tensors in the state dict" % key)
-        raise ValueError
         return state_dict
 
     @property
@@ -137,22 +135,22 @@ def _get_build_directory(name, verbose):
 for name, cls in inspect.getmembers(optim):
     if inspect.isclass(cls) and issubclass(cls, Optimizer):
         setattr(optim, "_%s" % name, cls)
-        cls = make_configurable(cls, ignore_args=("params",))
+        cls = core.make_configurable(cls, ignore_args=("params",))
         cls = R.register("optim.%s" % name)(cls)
         setattr(optim, name, cls)
 
 Scheduler = scheduler._LRScheduler
 for name, cls in inspect.getmembers(scheduler):
     if inspect.isclass(cls) and issubclass(cls, Scheduler):
         setattr(scheduler, "_%s" % name, cls)
-        cls = make_configurable(cls, ignore_args=("optimizer",))
+        cls = core.make_configurable(cls, ignore_args=("optimizer",))
         cls = R.register("scheduler.%s" % name)(cls)
         setattr(optim, name, cls)
 
 Dataset = dataset.Dataset
 for name, cls in inspect.getmembers(dataset):
     if inspect.isclass(cls) and issubclass(cls, Dataset):
         setattr(dataset, "_%s" % name, cls)
-        cls = make_configurable(cls)
+        cls = core.make_configurable(cls)
         cls = R.register("dataset.%s" % name)(cls)
         setattr(dataset, name, cls)