Cant save models when having some custom layers using torch.jit.trace

username: patrickpjiang
cooked: <p>Hi,</p>
<p>I’m a freshman to pytorch, recently I got some troubles when trying to convert a python models to c++. HELPPPPPP！！</p>
<p>enviroment:<br>
pytorch1.1</p>
<p>When i trying to trace some custom layers which implemented by c++, i got an error like this:</p>
<pre><code class="lang-auto">could not export python function call &lt;python_value&gt;. Remove calls to Python functions before export. Did you forget add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__.:
@torch.jit.script_method
def forward(self, x):
    p1_conv1 = self.p1_conv1(x)
    pool1    = self.pool1(p1_conv1) # Troubles is here
               ~~~~~~~~~~ &lt;--- HERE


</code></pre>
<p>the code is here:</p>
<pre><code class="lang-auto">class corner_pool(torch.jit.ScriptModule):
    def _init_layers(self, dim, pool1, pool2):
        self.p1_conv1 = torch.jit.trace(convolution(3, dim, 128), torch.rand(1, 256, 64, 64))
        self.pool1 = TopPool()  #custom layers implemented in cpp

    @torch.jit.script_method
    def forward(self, x):
        p1_conv1 = self.p1_conv1(x)
        pool1    = self.pool1(p1_conv1) # Troubles is here
</code></pre>
<p>the toppool code is:</p>
<pre><code class="lang-auto">class TopPoolFunction(Function):
    @staticmethod
    def forward(ctx, input):
        output = TopPool.forward(input)[0]
        ctx.save_for_backward(input)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        input  = ctx.saved_variables[0]
        output = TopPool.backward(input, grad_output)[0]
        return output

class TopPool(nn.Module):
    def forward(self, x):
        result = TopPoolFunction.apply(x)
        return result
</code></pre>

username: wanchaol
cooked: <aside class="quote no-group quote-modified" data-username="patrickpjiang" data-post="1" data-topic="49072">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/p/9dc877/48.png" class="avatar"> patrickpjiang:</div>
<blockquote>
<p>could not export python function call &lt;python_value&gt;. Remove calls to Python functions before export. Did you forget add <span class="mention">@script</span> or <span class="mention">@script_method</span> annotation? If this is a nn.ModuleList, add it to <strong>constants</strong>.:</p>
</blockquote>
</aside>
<p>The error message says that you have a function that you did not provide script annotation. In your case you will need to script <code>TopPoolFunction</code> and <code>TopPool</code> to make the model to be exported.</p>

username: patrickpjiang
cooked: <p>Thanks for the reply.<br>
Yes, you are right, but when i apply script on <code>TopPoolFunction</code>, i got an error:</p>
<pre><code class="lang-auto">attribute lookup is not defined on python value of type 'FunctionMeta':
@torch.jit.script_method
def forward(self, x):
    result = TopPoolFunction.apply(x)
             ~~~~~~~~~~~~~~~~~~~~~ &lt;--- HERE
    return result


</code></pre>
<p>Here is how i script TopPool:</p>
<pre><code class="lang-auto">class TopPool(torch.jit.ScriptModule):
     # here is what i added, is there any problem?
    @torch.jit.script_method        
    def forward(self, x):
        result = TopPoolFunction.apply(x)
        return result
</code></pre>
<p>For <code>TopPoolFunction</code>, i actually have no idea how to script…</p>

username: wanchaol
cooked: <p>seems like <code>TopPoolFunction</code> is a autograd function, we don’t support script autograd Functions, so it is not scriptable.</p>
<p>If you want to script it, you will need to write your TopPool function in python or TorchScript in order for TorchScript to compile it(TorchScript can only compile python, not C++),</p>