Revert "use object.__setattr__ in DualModule and DeployableModule to avoid unexpected key in self._modules"

src/onediff/infer_compiler/with_oneflow_compile.py

@@ -19,12 +19,8 @@
 class DualModule(torch.nn.Module):
     def __init__(self, torch_module, oneflow_module):
         torch.nn.Module.__init__(self)
-        # avoid to set key '_torch_module' in self._modules
-        object.__setattr__(self, '_torch_module', torch_module)
-        object.__setattr__(self, '_oneflow_module', oneflow_module)
-        self._modules.update(**torch_module._modules)
-        self._parameters.update(**torch_module._parameters)
-        self._buffers.update(**torch_module._buffers)
+        self._torch_module = torch_module
+        self._oneflow_module = oneflow_module
 
     @property
     def oneflow_module(self):
@@ -237,14 +233,9 @@ def __init__(
         graph_device=None,
     ):
         torch.nn.Module.__init__(self)
-        mixed_dual_module = get_mixed_dual_module(torch_module.__class__)(torch_module, oneflow_module)
-        # avoid to set '_deployable_module_model' in self._modules
-        object.__setattr__(self, '_deployable_module_model', mixed_dual_module)
-
-        self._modules.update(**self._deployable_module_model._modules)
-        self._parameters.update(**self._deployable_module_model._parameters)
-        self._buffers.update(**self._deployable_module_model._buffers)
-
+        self._deployable_module_model = get_mixed_dual_module(torch_module.__class__)(
+            torch_module, oneflow_module
+        )
         self._deployable_module_use_graph = use_graph
         self._deployable_module_options = options
         self._deployable_module_dpl_graph = None
@@ -417,6 +408,19 @@ def get_oneflow_graph(model, size=9, all_dynamic=False):
     return g
 
 
+def state_dict_hook(module, state_dict, prefix, local_metadata):
+    pytorch_key_prefix = "_deployable_module_model._torch_module."
+    new_state_dict = type(state_dict)()
+    for k, v in state_dict.items():
+        # _deployable_module_model._torch_module.out.2.weight => out.2.weight
+        if k.startswith(pytorch_key_prefix):
+            new_k = k[len(pytorch_key_prefix) :]
+            new_state_dict[new_k] = v
+        else:
+            new_state_dict[k] = v
+    return new_state_dict
+
+
 # Return a DeployableModule that using module_cls as it's parent class.
 def get_mixed_deployable_module(module_cls):
     class MixedDeployableModule(DeployableModule, module_cls):
@@ -473,5 +477,6 @@ def wrap_module(module):
     model = wrap_module(torch_module)
     assert isinstance(model, DeployableModule)
     assert isinstance(model, torch_module.__class__)
+    model._register_state_dict_hook(state_dict_hook)
 
     return model