Proxy and module calls

username: nscotto
cooked: <p>Hi,</p>
<p>Following <a href="https://discuss.pytorch.org/t/fx-insert-a-graph-instead-of-a-node/119694/4">this discussion</a> about inserting a block instead of a node, I would like to rebuild a model using <code>Proxy</code>.<br>
However I am unable to build a model with calls to modules, here is an example:</p>
<pre><code class="lang-auto">import torch.nn as nn
import torch.fx as fx

module = nn.Module()
submodule = nn.Linear(10, 1)
module.add_module('sub', submodule)

graph = fx.Graph()
raw = graph.placeholder('x')
x = fx.Proxy(raw)
y = submodule(x)
output = graph.output(y.node)
gm = fx.GraphModule(module, graph)
</code></pre>
<p>result:</p>
<pre><code class="lang-auto">Traceback (most recent call last):

  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3437, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)

  File "&lt;ipython-input-37-1d2b0a63ef08&gt;", line 1, in &lt;module&gt;
    gm = fx.GraphModule(module, graph)

  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/graph_module.py", line 194, in __init__
    self.graph = graph

  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/nn/modules/module.py", line 995, in __setattr__
    object.__setattr__(self, name, value)

  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/graph_module.py", line 217, in graph
    self.recompile()

  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/graph_module.py", line 297, in recompile
    cls.forward = _forward_from_src(self._code)

  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/graph_module.py", line 41, in _forward_from_src
    exec_with_source(src, gbls)

  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/graph_module.py", line 26, in exec_with_source
    exec(compile(src, key, 'exec'), globals)

  File "&lt;eval_with_key_10&gt;", line 3
    linear_1 = torch.nn.functional.linear(x, Parameter containing:
                                                                ^
SyntaxError: invalid syntax
</code></pre>
<p>When I print the graph I obtain this:</p>
<pre><code class="lang-auto">graph(x):
    %linear_1 : [#users=1] = call_function[target=torch.nn.functional.linear](args = (%x, Parameter containing:
tensor([[-0.1319, -0.2979,  0.2596, -0.1899,  0.1811, -0.1425,  0.1439, -0.0744,
         -0.0514,  0.1271]], requires_grad=True)), kwargs = {bias: None})
    return linear_1
</code></pre>
<p>Any help would be appreciated.<br>
Thank you.</p>

username: nscotto
cooked: <p>This seems to happen with modules containing parameters.<br>
Is it a bug?<br>
If not then how should I handle them?</p>
<p>The example in the github repository doesn’t showcase the use of proxy with these type of modules.</p>
<p>Thank you in advance.</p>

username: James_Reed
cooked: <p>Hi <a class="mention" href="/u/nscotto">@nscotto</a>,</p>
<p>To resolve modules/parameters/buffers while tracing, a <code>Proxy</code> must be initialized with an underlying <code>Tracer</code> instance. It’s kind-of hacky, but the quickest way I found to fix your code was to write it like this:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.fx as fx

module = nn.Module()
submodule = nn.Linear(10, 1)
module.add_module('sub', submodule)

graph = fx.Graph()
raw = graph.placeholder('x')
tracer = torch.fx.Tracer()
tracer.root = module
tracer.graph = graph
x = fx.Proxy(raw, tracer)
y = submodule(x)
output = graph.output(y.node)
gm = fx.GraphModule(module, graph)

print(gm.code)
</code></pre>
<p>I think a more idiomatic way of doing this would be to use a custom <code>Tracer</code> that specifies modules to not treat as leaves:</p>
<pre><code class="lang-auto">import torch.fx

class SomeModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(3, 4)

    def forward(self, x):
        return self.linear(x)

traced = torch.fx.symbolic_trace(SomeModule())
print(traced.code)
"""
def forward(self, x):
    linear = self.linear(x);  x = None
    return linear
"""

# Now we can use a custom tracer to set the linear instance as a non-leaf module
# and thus trace into it

class InlineLinearTracer(torch.fx.Tracer):
    def __init__(self, linears_to_inline):
        super().__init__()
        self.linears_to_inline = linears_to_inline

    def is_leaf_module(self, m: torch.nn.Module, module_qualified_name : str) -&gt; bool:
        if m in self.linears_to_inline:
            return False
        return super().is_leaf_module(m, module_qualified_name)


tracer = InlineLinearTracer([traced.linear])
graph = tracer.trace(traced)
inlined_gm = torch.fx.GraphModule(traced, graph)
print(inlined_gm.code)
"""
def forward(self, x):
    linear_weight = self.linear.weight
    linear_bias = self.linear.bias
    linear = torch.nn.functional.linear(x, linear_weight, bias = linear_bias);  x = linear_weight = linear_bias = None
    return linear
"""
</code></pre>

username: nscotto
cooked: <p>Thank you for the answer.<br>
Let me add a bit more context to my problem.</p>
<p>I would like to replace a block by another block (e.g. a backbone), this block eventually have leaking connections (skip connections from inside the block to outside the block).</p>
<p>Previously I was doing it by adding the new block using <code>add_module</code> as a single <code>Module</code> and then adding a node to the graph (using <code>graph.inserting_before(node)</code>), and erasing every nodes from the replaced block.</p>
<p>The problem with this approach is that it didn’t allow me to remap leaking connections because I was inserting the new block as a single module (as if it was a single layer).<br>
So I would like to insert the new block “layer by layer”/node by node, and I am looking for a solution to do it.</p>
<p>The solution proposed in <a href="https://discuss.pytorch.org/t/fx-insert-a-graph-instead-of-a-node/119694/4">this discussion</a> is to use <code>Proxy</code> to recreate the graph, but I am struggling to use proxy with modules, hence this post.<br>
While the proposed solution solve the small example I gave, I am not sure how to use it in my context, that is merging two <code>Module</code> together/replacing a block by another block.</p>

username: nscotto
cooked: <p>So right now I am trying to rebuild a module as a submodule, I think it’s a progression towards my goal of being able to replace a backbone and handle intermediary skip connections.</p>
<p>I have some code for creating a new module that mimics another module (but use it as a submodule):</p>
<p>Here is the code;</p>
<pre><code class="lang-auto">def rebuild_module_as_submodule(module: fx.GraphModule):
    module = copy.deepcopy(module)

    new_module = nn.Module()
    new_module.add_module('mod', module)

    graph = fx.Graph()
    raw1 = graph.placeholder('x')

    tracer = torch.fx.Tracer()
    tracer.root = new_module
    tracer.graph = graph

    x = fx.Proxy(raw1)

    nodes = list(module.graph.nodes)
    node2proxy = {raw1: x, nodes[0]: x}

    def nodeargs2proxy(args):
        if isinstance(args, (list, tuple)):
            return tuple(nodeargs2proxy(arg) for arg in args)
        elif isinstance(args, dict):
            return {nodeargs2proxy(key): nodeargs2proxy(val) for key, val in args.items()}
        elif isinstance(args, fx.Node):
            return node2proxy[args]
        else:
            return args

    submodules = dict(new_module.named_modules())
    
        for node in nodes[1:-1]:
        if node.op == 'call_module':
            args = nodeargs2proxy(node.args)
            kwargs = nodeargs2proxy(node.kwargs)
            node2proxy[node] = submodules[f'mod.{node.target}'](*args, **kwargs)
        elif node.op == 'call_function':
            args = nodeargs2proxy(node.args)
            kwargs = nodeargs2proxy(node.kwargs)
            node2proxy[node] = node.target(*args, **kwargs)
        elif node.op == 'call_method':
            pass
            self_obj, *args = nodeargs2proxy(node.args)
            kwargs = nodeargs2proxy(node.kwargs)
            node2proxy[node] = getattr(self_obj, node.target)(*args, **kwargs)
        else:
            raise NotImplementedError(f'Node op: op={node.op}, name={node.name}, tgt={node.target}, args={node.args}, kwargs={node.kwargs}')

    last_proxy : fx.Proxy = node2proxy[node]
    graph.output(last_proxy.node)

    return new_module, graph

def test_2():
    submodule = nn.Linear(10, 1)
    module = fx.symbolic_trace(nn.Sequential(submodule))
    new_module, graph = rebuild_module_as_submodule(module)
    print(new_module)
    print(graph)
    gm = fx.GraphModule(new_module, graph)
    print(gm)

test_2()
</code></pre>
<p>I am still getting the same error as before, even though I created an instance of a <code>Tracer</code> and assignated the module and graph I am working with.</p>
<p>Here is the result of the above code:</p>
<pre><code class="lang-auto">{'': Module(
  (mod): GraphModule(
    (0): Linear(in_features=10, out_features=1, bias=True)
  )
), 'mod': GraphModule(
  (0): Linear(in_features=10, out_features=1, bias=True)
), 'mod.0': Linear(in_features=10, out_features=1, bias=True)}
Module(
  (mod): GraphModule(
    (0): Linear(in_features=10, out_features=1, bias=True)
  )
)
graph(x):
    %linear_1 : [#users=1] = call_function[target=torch.nn.functional.linear](args = (%x, Parameter containing:
tensor([[-0.3050,  0.0150, -0.2705, -0.0750, -0.0539, -0.0999,  0.0665, -0.0039,
          0.3120, -0.2078]], requires_grad=True)), kwargs = {bias: Parameter containing:
tensor([0.1116], requires_grad=True)})
    return linear_1
Traceback (most recent call last):
  File "test.py", line 91, in &lt;module&gt;
    test_2()
  File "test.py", line 88, in test_2
    gm = fx.GraphModule(new_module, graph)
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/graph_module.py", line 194, in __init__
    self.graph = graph
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/nn/modules/module.py", line 995, in __setattr__
    object.__setattr__(self, name, value)
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/graph_module.py", line 217, in graph
    self.recompile()
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/graph_module.py", line 297, in recompile
    cls.forward = _forward_from_src(self._code)
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/graph_module.py", line 41, in _forward_from_src
    exec_with_source(src, gbls)
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/graph_module.py", line 26, in exec_with_source
    exec(compile(src, key, 'exec'), globals)
  File "&lt;eval_with_key_2&gt;", line 3
    linear_1 = torch.nn.functional.linear(x, Parameter containing:
                                                                ^
SyntaxError: invalid syntax
</code></pre>
<p>It is still this problem of parameters that are missing from the module, and I don’t know how to handle it in my case.<br>
Do you have any idea how to fix this?</p>