Skip to content

Conversation

@BowenBao
Copy link
Collaborator

Following this tutorial torch script custom ops, users have created their custom ops implementation and registered in PyTorch.

class FooModel(torch.nn.Module):
    def __init__(self, attr1, attr2):
        super(FooModule, self).__init__()
        self.attr1 = attr1
        self.attr2 = attr2

    def forward(self, input1, input2):
        # Calling custom op
        return torch.ops.custom_ops.foo_forward(input1, input2, self.attr1, self.attr2)

model = FooModel(attr1, attr2)

With this PR, users will be able to provide their own symbolic function for their custom ops.

# Create custom symbolic function
from torch.onnx.symbolic import parse_args
@parse_args('v', 'v', 'f', 'i')
def symbolic_foo_forward(g, input1, input2, attr1, attr2):
    return g.op("Foo", input1, input2, attr1_f=attr1, attr2_i=attr2)

# Register custom symbolic function
from torch.onnx import register_custom_op_symbolic
register_custom_op_symbolic('custom_ops::foo_forward', symbolic_foo_forward)

And proceed to export to ONNX as usual.

torch.onnx.export(model, (dummy_input1, dummy_input2), 'model.onnx')

@BowenBao BowenBao force-pushed the onnx_custom_ops_sync branch from 3eedac3 to 7bcab3a Compare April 1, 2019 20:44
Copy link
Member

@houseroad houseroad left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a great move. Thanks for adding the support for custom op. It will benefits more users.

  1. please address the inline comments
  2. add a test case, probably in test/onnx/test_operators.py
  3. add appropriate docs

return getattr(self, sel)(k)

def register_custom_op_symbolic(symbolic_name, symbolic_fn):
ns, op_name = symbolic_name.split('::')
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add assertion to check the format of the symbolic_name, should be domain::op.

Also, we need to check the domain doesn't conflict with existing domain

@BowenBao
Copy link
Collaborator Author

Thanks for the comments @houseroad! Closing this PR now that it is covered in #19294

@BowenBao BowenBao closed this Apr 26, 2019
facebook-github-bot pushed a commit that referenced this pull request May 11, 2019
Summary:
Support exporting multiple ONNX opsets (more specifically opset 10 for now), following the proposal in https://gist.github.com/spandantiwari/99700e60919c43bd167838038d20f353.
And add support for custom ops (merge with #18297).

This PR will be followed by another PR containing the changes related to testing the ops for different opsets.
Pull Request resolved: #19294

Reviewed By: zrphercule

Differential Revision: D15043951

Pulled By: houseroad

fbshipit-source-id: d336fc35b8827145639137bc348ae07e3c14bb1c
zdevito pushed a commit to zdevito/ATen that referenced this pull request May 11, 2019
Summary:
Support exporting multiple ONNX opsets (more specifically opset 10 for now), following the proposal in https://gist.github.com/spandantiwari/99700e60919c43bd167838038d20f353.
And add support for custom ops (merge with pytorch/pytorch#18297).

This PR will be followed by another PR containing the changes related to testing the ops for different opsets.
Pull Request resolved: pytorch/pytorch#19294

Reviewed By: zrphercule

Differential Revision: D15043951

Pulled By: houseroad

fbshipit-source-id: d336fc35b8827145639137bc348ae07e3c14bb1c
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants