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

jit
1

Hi,

I’m a freshman to pytorch, recently I got some troubles when trying to convert a python models to c++. HELPPPPPP！！

enviroment:
pytorch1.1

When i trying to trace some custom layers which implemented by c++, i got an error like this:

could not export python function call <python_value>. 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
               ~~~~~~~~~~ <--- HERE

the code is here:

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

the toppool code is:

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
2

The error message says that you have a function that you did not provide script annotation. In your case you will need to script TopPoolFunction and TopPool to make the model to be exported.

3

Thanks for the reply.
Yes, you are right, but when i apply script on TopPoolFunction, i got an error:

attribute lookup is not defined on python value of type 'FunctionMeta':
@torch.jit.script_method
def forward(self, x):
    result = TopPoolFunction.apply(x)
             ~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    return result

Here is how i script TopPool:

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

For TopPoolFunction, i actually have no idea how to script…

4

seems like TopPoolFunction is a autograd function, we don’t support script autograd Functions, so it is not scriptable.

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++),

1 Like