[Frontend][ONNX] LSTM Support #4825
Conversation
|
@masahi, @soiferj, @mbrookhart can you take a look at this PR? |
|
Awesome!! I was looking at implementing this myself yesterday. I’ll take a look as soon as possible. Thanks for sending the PR! |
|
looks great! Will wait for @soiferj's review. |
| def _impl_v7(cls, inputs, attr, params): | ||
| # Unpack inputs, note that if optional and not provided then value will be None. | ||
| X = inputs[0] | ||
| W = inputs[1] |
There was a problem hiding this comment.
Is there any case when the weights won’t be constant? If they’re constant, we can remove some operations from the graph and compute them here (like squeeze).
By constant, I mean we can call infer_value on it.
There was a problem hiding this comment.
I think in almost all cases it'd be safe to assume weights are constant. However, the fold constant pass in relay will eliminate all operations on the weights anyway. Since treating the weights as a non-constant is slightly more flexible I prefer it.
There was a problem hiding this comment.
Sounds good. Thanks a lot for the updates.
| tvm.testing.assert_allclose(o_out, t_out, rtol=5e-3, atol=5e-3) | ||
|
|
||
|
|
||
| def test_lstm(): |
There was a problem hiding this comment.
Can you also add a test where initial c and h states are set to something other than 0?
There was a problem hiding this comment.
New tests for initial states and peephole weights are added. Glad you pointed this out since both those cases had some small bugs.
| @@ -32,6 +32,55 @@ | |||
| __all__ = ['from_onnx'] | |||
|
|
|||
|
|
|||
| class onnx_input(): | |||
There was a problem hiding this comment.
I really like this design - very sleek.
| def _impl_v7(cls, inputs, attr, params): | ||
| # Unpack inputs, note that if optional and not provided then value will be None. | ||
| X = inputs[0] | ||
| W = inputs[1] |
There was a problem hiding this comment.
Sounds good. Thanks a lot for the updates.
|
Thanks @jwfromm @soiferj @mbrookhart |
* Initial version working and passing tests. * WIP on supporting other activations. * add support for multiple activation functions in lstm * All tests working and code cleaned up. * Undo import swap to avoid conflict with masahi. * Added new tests and related bug fixes. Co-authored-by: Matthew Brookhart <mbrookhart@octoml.ai>
* Initial version working and passing tests. * WIP on supporting other activations. * add support for multiple activation functions in lstm * All tests working and code cleaned up. * Undo import swap to avoid conflict with masahi. * Added new tests and related bug fixes. Co-authored-by: Matthew Brookhart <mbrookhart@octoml.ai>
* Initial version working and passing tests. * WIP on supporting other activations. * add support for multiple activation functions in lstm * All tests working and code cleaned up. * Undo import swap to avoid conflict with masahi. * Added new tests and related bug fixes. Co-authored-by: Matthew Brookhart <mbrookhart@octoml.ai>
* Initial version working and passing tests. * WIP on supporting other activations. * add support for multiple activation functions in lstm * All tests working and code cleaned up. * Undo import swap to avoid conflict with masahi. * Added new tests and related bug fixes. Co-authored-by: Matthew Brookhart <mbrookhart@octoml.ai>
This PR adds LSTM support to the relay Onnx frontend.
Besides adding the LSTM parser itself, we encountered an issue where for some Onnx operations (like LSTMs) arguments are optional. The current method for passing arguments to converters is just to pack them into a list however as some arguments are optional the position of each input is inconsistent. Instead, we should be using a dictionary mapping input names to their value. However, changing all inputs to a dictionary would require changing all the current operators and present problems with direct Onnx to relay conversions. Our workaround here is to add the
onnx_inputclass that can be accessed as a list as we previously did or with input name dictionary style lookup.