Error : When torch.utils.cpp_extension and torch.fx are used together

MartinZhang: <p><strong>base.cpp</strong></p>
<pre><code class="lang-auto">#include &lt;torch/extension.h&gt;

torch::Tensor base_forward(torch::Tensor x, torch:Tensor w, torch::Tensor b){
    atuo o = w * x + b;
    return o
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
    m.def("forward", &amp;base_forward, 'BASE forward');
}
</code></pre>
<p><strong>module.py</strong></p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import base_cpp
form torch.fx import Tracer

class M(nn,Module):
    def __init__(self):
        super().__init__()
            self.conv = nn.Conv2d(1, 1. 1)

    def forward(self, x):
        x = self.conv(x)
        out = base_cpp.forward(x, torch.randn(1), torch.randn(1))
        return out

module = M()

nodes = Tracer().trace(module).nodes
</code></pre>
<p><strong>Error</strong></p>
<pre><code class="lang-auto">TypeError: forward():incompatible function arguments. The following argument types are supported:
1.(argo:at::Tensor, arg1: at:Tensor, arg2: at::Tensor)-&gt;at:Tensor

Invoked with: proxy(conv), tensor([2.0171]), tensor([-1.9950])
</code></pre>

tom: <p>I think you might try your luck with <a href="https://pytorch.org/docs/stable/fx.html#non-torch-functions" rel="noopener nofollow ugc"><code>torch.fx.wrap</code> described in the non-torch functions section of the fx documentation</a>.</p>
<p>Best regards</p>
<p>Thomas</p>
<p>P.S.: So the more detailed explanation for the error is that during tracing, the function will be called not with a tensor <code>x</code> but with a proxy object that causes functions to amend to the trace. wrap is a mechanism make function calls with these proxy objects insert a record of their invocation.</p>

MartinZhang: <p>Dear friend，<br>
<strong>I realy appreciate your help.</strong> However,I still have a question about how wrap custom methods.</p>
<p>temp.py</p>
<pre><code class="lang-auto">import torch
import base_cpp
from torch.fx import symbolic_trace
import torch.fx

def pybase_cpp(x):
    return base_cpp.forward(x, torch.randn(1), torch.randn(1))  # base_cpp.forward no __code__

torch.fx.wrap('base_cpp.forward')
pybase_cpp(torch.randn(1))
traced = symbolic_trace(pybase_cpp)
</code></pre>
<p>error</p>
<pre><code class="lang-auto">KeyError: 'base_cpp.forward'
</code></pre>
<p>So. Is there any good way to wrap custom methods. Thanks very much.</p>

tom: <p>I think you need to put it in the global namespace (or at least that is a feasible workaround):</p>
<pre><code class="lang-python">import torch
import torch.utils.cpp_extension

cpp_source = """
#include &lt;torch/extension.h&gt;

torch::Tensor base_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b){
    auto o = w * x + b;
    return o;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
    m.def("forward", &amp;base_forward, "BASE forward");
}
"""

import torch
import torch.nn as nn

base_cpp = torch.utils.cpp_extension.load_inline("base_cpp", cpp_source)

from torch.fx import Tracer

base_cpp_forward = base_cpp.forward
torch.fx.wrap('base_cpp_forward')

class M(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(1, 1, 1)

    def forward(self, x):
        x = self.conv(x)
        out = base_cpp_forward(x, torch.randn(1), torch.randn(1))
        return out

module = M()

traced = torch.fx.symbolic_trace(module)
print (traced.code)
</code></pre>
<p>P.S.: Kudos for posting code, but it seems to have a lot of small typos.</p>

MartinZhang: <p><strong>It’s awesome！！！</strong></p>
<p>Thanks, my friend.  <img src="https://discuss.pytorch.org/images/emoji/apple/+1.png?v=10" title=":+1:" class="emoji" alt=":+1:"> <img src="https://discuss.pytorch.org/images/emoji/apple/+1.png?v=10" title=":+1:" class="emoji" alt=":+1:"> <img src="https://discuss.pytorch.org/images/emoji/apple/+1.png?v=10" title=":+1:" class="emoji" alt=":+1:"></p>

MartinZhang: <p>Hi, dear friend，I try to convert this .py file to the .so file</p>
<p>when I call this .so file</p>
<p>the <code>torch.fx.wrap</code> function seem unable to call</p>
<p>a.cpp : this is a ops file</p>
<pre><code class="lang-auto">#include &lt;torch/extension.h&gt;

torch::Tensor base_cpp_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b){
    auto o = w * x + b;
    return o;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
    m.def("forward", &amp;base_cpp_forward, "BASE forward");
}
</code></pre>
<p>m.py : this define the model with ops</p>
<pre><code class="lang-auto">import torch
from a import *

torch.fx.wrap('base_cpp_forward')

class M(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(1, 1, 1)

    def forward(self, x):
        x = self.conv(x)
        out = base_cpp_forward(x, torch.randn(1), torch.randn(1))
        return out
</code></pre>
<p>setup.py  : convert the .py to .so</p>
<pre><code class="lang-auto">from distutils.core import setup
from Cython.Build import cythonize

setup(name='m', ext_modules = cythonize(['m.py']))
</code></pre>
<p>t.py : test the m.so file</p>
<pre><code class="lang-auto">from m import M  # this m is .so file
from torchvision.models.feature_extraction import get_graph_node_names

m = M()
train_node = get_graph_node_names(m)
</code></pre>
<p>the error: from : _symbolic_trace.py</p>
<blockquote>
<p>assert isinstance(fn, FunctionType) = False</p>
</blockquote>
<p>when I ignore this order</p>
<pre><code class="lang-auto">TypeError: forward():incompatible function arguments. The following argument types are supported:
1.(argo:at::Tensor, arg1: at:Tensor, arg2: at::Tensor)-&gt;at:Tensor

Invoked with: proxy(conv), tensor([2.0171]), tensor([-1.9950])
</code></pre>
<p>How to deal it, please help me , thx.</p>