How to transform custom ops to onnx

I try to transform the custom ops to onnx as the nms set. I define the static symbolic method. when the output is single, I got the custom onnx node. however, when the number of output is more than one, the custom node would be optimized. Here is my simple code.

def func1(a, b):
    c = torch.add(a, b)
    d = torch.sub(c, a)
    return [c,d]

class NMSop(torch.autograd.Function):

    @staticmethod
    def forward(ctx, a, thr):
        return func1(a, thr)

    @staticmethod
    def symbolic(g, a, thr):
        x =  g.op('Aten::NMS', a, thr)
        return  x

class CustomNet(nn.Module):
    def __init__(self):
        super(CustomNet, self).__init__()
        self.net = NMSop()

    def forward(self, a, b):
        #import pdb;pdb.set_trace()
        x = NMSop.apply(a, b)
        #x = a.repeat_interleave(torch.tensor([2]), dim=0)
        return x

net = CustomNet()
t = torch.randn(3,4)
s = torch.randn(3,4)
#result = net(t)
#print(result)
torch.onnx.export(net, (t,s), 'test.onnx', verbose=True)

if the func1 return c, I can get the NMS node, if return [c,d], can not. So there are some methods to process multiple output?

[zgplvyou]

Sorry, I do not know how to add shape inference for the symbolic in PyTorch. Maybe you can add a custom schema with this.

ok, I trace the graph, the shape information is missing in the jit graph, I will continue to find the solution next week. Anyway, happy mid-autumn festival, today is going into garbage time.

[grimoire]

Sorry, I do not know how to add shape inference for the symbolic in PyTorch. Maybe you can add a custom schema with this.