Repository: microsoft/nni

Issue:
Using torch.jit.export, AttributeError: 'RecursiveScriptModule' object has no attribute '_name' ​
**Describe the issue**:
I am trying to package my model into a jit package. In the process, I want to include the predict function (SSD model) in the jit package for which I am using a function decorator called `@torch.jit.export`. On decorating the function with this, I am getting the following error. 

If I exclude this decorator, then I do not face any issues.

```python
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
/tmp/ipykernel_416/3006178987.py in <module>
      1 pruner_config = config.pruner
      2 pruner = instantiate(pruner_config, model=model)
----> 3 pruner.setup()

/workspace/pest-monitoring-new/compression/pruner.py in setup(self)
     45 
     46         # Setup pruner
---> 47         self._set_pruner()
     48 
     49     def compress(self):

/workspace/pest-monitoring-new/compression/pruner.py in _set_pruner(self)
    129     def _set_pruner(self):
    130         self.model.eval()
--> 131         self.pruner = L1FilterPruner_(
    132             self.model,
    133             oc.to_container(self.config_list),

/usr/local/lib/python3.8/dist-packages/nni/algorithms/compression/pytorch/pruning/one_shot_pruner.py in __init__(self, model, config_list, dependency_aware, dummy_input)
    102 
    103     def __init__(self, model, config_list, dependency_aware=False, dummy_input=None):
--> 104         super().__init__(model, config_list, pruning_algorithm='l1', dependency_aware=dependency_aware,
    105                          dummy_input=dummy_input)
    106 

/usr/local/lib/python3.8/dist-packages/nni/algorithms/compression/pytorch/pruning/one_shot_pruner.py in __init__(self, model, config_list, pruning_algorithm, dependency_aware, dummy_input, **algo_kwargs)
     38             Additional parameters passed to pruning algorithm masker class
     39         """
---> 40         super().__init__(model, config_list, None, pruning_algorithm, dependency_aware, dummy_input, **algo_kwargs)
     41 
     42     def validate_config(self, model, config_list):

/usr/local/lib/python3.8/dist-packages/nni/algorithms/compression/pytorch/pruning/dependency_aware_pruner.py in __init__(self, model, config_list, optimizer, pruning_algorithm, dependency_aware, dummy_input, **algo_kwargs)
     43             # to trace the model, we need to unwrap the wrappers
     44             self._unwrap_model()
---> 45             self.graph = TorchModuleGraph(model, dummy_input)
     46             self._wrap_model()
     47             self.channel_depen = ChannelDependency(model, dummy_input, traced_model=self.graph.trace)

/usr/local/lib/python3.8/dist-packages/nni/common/graph_utils.py in __init__(self, model, dummy_input, traced_model)
    253         self.global_count = 0
    254         self.reused_module = set()
--> 255         self.name_to_node, self.input_to_node, self.output_to_node = self._build_graph()
    256         self._extract_auxiliary_info()
    257 

/usr/local/lib/python3.8/dist-packages/nni/common/graph_utils.py in _build_graph(self)
    709 
    710         self.leaf_modules = self._extract_leaf_modules()
--> 711         module_to_type = {name: parse_traced_name(
    712             module._name) for name, module in self.trace.named_modules()}
    713 

/usr/local/lib/python3.8/dist-packages/nni/common/graph_utils.py in <dictcomp>(.0)
    710         self.leaf_modules = self._extract_leaf_modules()
    711         module_to_type = {name: parse_traced_name(
--> 712             module._name) for name, module in self.trace.named_modules()}
    713 
    714         # associate module name with their trace graph nodes

/usr/local/lib/python3.8/dist-packages/torch/jit/_script.py in __getattr__(self, attr)
    559                 return script_method
    560 
--> 561             return super(RecursiveScriptModule, self).__getattr__(attr)
    562 
    563         def __setattr__(self, attr, value):

/usr/local/lib/python3.8/dist-packages/torch/jit/_script.py in __getattr__(self, attr)
    289         def __getattr__(self, attr):
    290             if "_actual_script_module" not in self.__dict__:
--> 291                 return super(ScriptModule, self).__getattr__(attr)
    292             return getattr(self._actual_script_module, attr)
    293 

/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py in __getattr__(self, name)
    945             if name in modules:
    946                 return modules[name]
--> 947         raise AttributeError("'{}' object has no attribute '{}'".format(
    948             type(self).__name__, name))
    949 

AttributeError: 'RecursiveScriptModule' object has no attribute '_name'
```

Any idea on how this can be resolved.


Discussion:
Hello @shenoynikhil98 , could you show how you use this predict function? Can wrap it as a nn.Module can work?
Hi @J-shang,
I tested this with a simpler example, Say I have two models,
```python
class TestModule1(torch.nn.Module):
    def forward(self, x):
        return x + 10

    def helper(self, x):
        return x + 5

class TestModule2(torch.nn.Module):
    def forward(self, x):
        return x + 10
    
    @torch.jit.export
    def helper(self, x):
        return x + 5    

model1 = TestModule1()
model2 = TestModule2()

model1_jit = torch.jit.trace(model1, torch.randn(1))
model2_jit = torch.jit.trace(model2, torch.randn(1))
```
There is a [line (711)](https://github.com/microsoft/nni/blob/8c2f717d830cc9b0da10eb41c3c0bcc01e022b61/nni/common/graph_utils.py#L711) in `nni/common/graph_utils.py` where 
```python
module_to_type = {name: parse_traced_name(module._name) for name, module in self.trace.named_modules()}
```

If I try to run a similar line, for `module1_jit` and `module2_jit`, it fails for the second one as it is a `RecursiveScriptModule`. Can you check this line of code in `nni/common/graph_utils.py`?
```python
[module._name for name, module in model2_jit.named_modules()]

AttributeError                            Traceback (most recent call last)
/tmp/ipykernel_45/2244082312.py in <module>
----> 1 [module._name for name, module in model2_jit.named_modules()]

/tmp/ipykernel_45/2244082312.py in <listcomp>(.0)
----> 1 [module._name for name, module in model2_jit.named_modules()]

/usr/local/lib/python3.8/dist-packages/torch/jit/_script.py in __getattr__(self, attr)
    559                 return script_method
    560 
--> 561             return super(RecursiveScriptModule, self).__getattr__(attr)
    562 
    563         def __setattr__(self, attr, value):

/usr/local/lib/python3.8/dist-packages/torch/jit/_script.py in __getattr__(self, attr)
    289         def __getattr__(self, attr):
    290             if "_actual_script_module" not in self.__dict__:
--> 291                 return super(ScriptModule, self).__getattr__(attr)
    292             return getattr(self._actual_script_module, attr)
    293 

/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py in __getattr__(self, name)
    945             if name in modules:
    946                 return modules[name]
--> 947         raise AttributeError("'{}' object has no attribute '{}'".format(
    948             type(self).__name__, name))
    949 

AttributeError: 'RecursiveScriptModule' object has no attribute '_name'

Resolving patch:
diff --git a/nni/common/graph_utils.py b/nni/common/graph_utils.py
--- a/nni/common/graph_utils.py
+++ b/nni/common/graph_utils.py
@@ -57,7 +57,7 @@ def __init__(self, model=None, dummy_input=None, traced_model=None):
         assert torch.__version__ >= '1.3.1'
         # check if the input is legal
         if traced_model is not None:
-            assert isinstance(traced_model, torch.jit.TopLevelTracedModule)
+            assert isinstance(traced_model, torch.jit.TopLevelTracedModule) or isinstance(traced_model, torch.jit.RecursiveScriptModule)
             self.trace = traced_model
             # it's ok if the graph is already unpacked
             torch._C._jit_pass_inline(self.trace.graph)
@@ -709,7 +709,7 @@ def _build_graph(self):
 
         self.leaf_modules = self._extract_leaf_modules()
         module_to_type = {name: parse_traced_name(
-            module._name) for name, module in self.trace.named_modules()}
+            module._name if hasattr(module, '_name') else module.original_name) for name, module in self.trace.named_modules()}
 
         # associate module name with their trace graph nodes
         for node in graph.nodes():


Test patch: