Torch.fx replace modules

nscotto: <p>Hello,</p>
<p>In the example for the replace pattern of torch.fx, function or replaced (<code>torch.add</code> by <code>torch.mul</code>).<br>
This is very clear, however it is not clear to me if it is possible to replace modules as well, and if so, how to do it.</p>
<p>The following example failed for me with the error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "test.py", line 43, in &lt;module&gt;
    torch.fx.subgraph_rewriter.replace_pattern(gm, pattern, replacement)
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/subgraph_rewriter.py", line 201, in replace_pattern
    pattern_graph = symbolic_trace(pattern).graph
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/symbolic_trace.py", line 606, in symbolic_trace
    graph = tracer.trace(root, concrete_args)
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/symbolic_trace.py", line 355, in trace
    self.create_node('output', 'output', (self.create_arg(fn(*args)),), {},
  File "test.py", line 33, in pattern
    val1 = F.relu(lin1(x))
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/symbolic_trace.py", line 344, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/symbolic_trace.py", line 221, in call_module
    module_qualified_name = self.path_of_module(m)
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/symbolic_trace.py", line 191, in path_of_module
    raise NameError('module is not installed as a submodule')
NameError: module is not installed as a submodule
</code></pre>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.fx


class TwoLayerSubNet(nn.Module):
    def __init__(self, D_in=100, H=50, D_out=10):
        super(TwoLayerSubNet, self).__init__()
        self.linear1 = nn.Linear(D_in, H)
        self.linear2 = nn.Linear(H, D_out)

    def forward(self, x):
            h_relu = F.relu(self.linear1(x))
            y_pred = self.linear2(h_relu)
            return y_pred


class TwoLayerNet(nn.Module):
    def __init__(self, D_in=100, H=50, D_int=20, D_out=10):
        super(TwoLayerNet, self).__init__()
        self.sub1 = TwoLayerSubNet(D_in, H, D_int)
        self.sub2 = TwoLayerSubNet(D_int, H, D_out)

    def forward(self, x):
            h_sub1 = self.sub1(x)
            y_pred = self.sub2(x)
            return y_pred

def pattern(x):
    lin1 = torch.nn.Linear(30, 20)
    lin2 = torch.nn.Linear(20, 10)
    val1 = F.relu(lin1(x))
    return F.relu(lin2(val1))


def replacement(x):
    lin1 = torch.nn.Linear(30, 10)
    return F.relu(lin1(x))

m = TwoLayerNet()
gm = torch.fx.symbolic_trace(m)
torch.fx.subgraph_rewriter.replace_pattern(gm, pattern, replacement)
</code></pre>
<p>Am I missing something?</p>
<p>Thanks</p>

MartinZhang: <p>Do you solve it？I have the same problem in resnet18.</p>

Egor_Shvecov: <p>I am also curious and I still have not figured out how to do that.</p>
<p>Right now I am looking at test code here: <a href="https://github.com/pytorch/pytorch/blob/ebf7a4f843ac2c34f0f0765c57f8c0bb8a194686/test/test_fx.py#L1863" class="inline-onebox" rel="noopener nofollow ugc">pytorch/test_fx.py at ebf7a4f843ac2c34f0f0765c57f8c0bb8a194686 · pytorch/pytorch · GitHub</a></p>
<p>There is some functionality to add and delete some atomic functions, however, I am not sure if it is possible to replace bigger classes.</p>

clark3614: <p>You should change the name, you MUST can’t write nn.relu or F.relu in foward function, use self.act = nn.relu insdead.</p>