Skip to content

Conversation

@BowenBao
Copy link
Collaborator

@BowenBao BowenBao commented Oct 2, 2019

This PR enables exporting aten::unbind created by the tracer. The traced version IR will always have this pattern aten::unbind -> prim::ListUnpack.
Another PR supporting scripted aten::unbind will be submitted separately later.

// Unbind is being converted to ONNX as Split + Squeeze.
// Example IR
// graph(%0 : Float(3, 4, 5)):
//   %7 : Long() = prim::Constant[value={0}]()
//   %3 : Tensor[] = aten::unbind(%0, %7)
//   %4 : Float(4, 5), %5 : Float(4, 5), %6 : Float(4, 5) = prim::ListUnpack(%3)
//   return (%4, %5, %6)
//
// Translates to ONNX:
// graph(%0 : Float(3, 4, 5)):
//   %1 : Tensor, %2 : Tensor, %3 : Tensor = onnx::Split[axis=0](%0)
//   %4 : Float(4, 5) = onnx::Squeeze[axes=[0]](%3)
//   %5 : Float(4, 5) = onnx::Squeeze[axes=[0]](%2)
//   %6 : Float(4, 5) = onnx::Squeeze[axes=[0]](%1)
//   return (%6, %5, %4)

@BowenBao BowenBao requested a review from apaszke as a code owner October 2, 2019 22:31
@pytorchbot pytorchbot added oncall: jit Add this issue/PR to JIT oncall triage queue module: onnx Related to torch.onnx labels Oct 2, 2019
@BowenBao
Copy link
Collaborator Author

BowenBao commented Oct 2, 2019

cc @lara-hdr @neginraoof @spandantiwari for review

Copy link
Contributor

@lara-hdr lara-hdr left a comment

Choose a reason for hiding this comment

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

LGTM

@neginraoof
Copy link
Contributor

Don't we need operator tests for expected model?

Copy link
Contributor

@facebook-github-bot facebook-github-bot left a comment

Choose a reason for hiding this comment

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

@houseroad has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

@houseroad
Copy link
Member

I think end to end test is good enough for this case.

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.

Looks good.

@houseroad
Copy link
Member

@pytorchbot rebase this please

Copy link
Contributor

@facebook-github-bot facebook-github-bot left a comment

Choose a reason for hiding this comment

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

@houseroad has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

@facebook-github-bot
Copy link
Contributor

@houseroad merged this pull request in ba79233.

thiagocrepaldi pushed a commit to thiagocrepaldi/pytorch that referenced this pull request Feb 4, 2020
Summary:
This PR enables exporting aten::unbind created by the tracer. The traced version IR will always have this pattern ```aten::unbind -> prim::ListUnpack```.
Another PR supporting scripted aten::unbind will be submitted separately later.
```
// Unbind is being converted to ONNX as Split + Squeeze.
// Example IR
// graph(%0 : Float(3, 4, 5)):
//   %7 : Long() = prim::Constant[value={0}]()
//   %3 : Tensor[] = aten::unbind(%0, %7)
//   %4 : Float(4, 5), %5 : Float(4, 5), %6 : Float(4, 5) = prim::ListUnpack(%3)
//   return (%4, %5, %6)
//
// Translates to ONNX:
// graph(%0 : Float(3, 4, 5)):
//   %1 : Tensor, %2 : Tensor, %3 : Tensor = onnx::Split[axis=0](%0)
//   %4 : Float(4, 5) = onnx::Squeeze[axes=[0]](%3)
//   %5 : Float(4, 5) = onnx::Squeeze[axes=[0]](%2)
//   %6 : Float(4, 5) = onnx::Squeeze[axes=[0]](%1)
//   return (%6, %5, %4)
```
Pull Request resolved: pytorch#27247

Reviewed By: hl475

Differential Revision: D17791095

Pulled By: houseroad

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

Labels

Merged module: onnx Related to torch.onnx oncall: jit Add this issue/PR to JIT oncall triage queue open source

Projects

None yet

Development

Successfully merging this pull request may close these issues.

8 participants