Count PackedSequence cell steps in the flops profiler RNN hook - #8507
Count PackedSequence cell steps in the flops profiler RNN hook#8507ebarkhordar wants to merge 1 commit into
Conversation
_rnn_forward_hook read inp.shape[0] and inp.shape[1] off the first positional argument, but nn.RNN, nn.GRU and nn.LSTM all accept a PackedSequence there and it has no .shape, so profiling any packed-input RNN raised AttributeError. A PackedSequence runs one cell step per row of .data, so .data.shape[0] is already the batch_size * seq_length product the count multiplies by. Reading .data.shape[1] as the sequence length would scale the result by input_size. Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3dc3a8c7d6
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| inp = input[0] | ||
| batch_size = inp.shape[0] | ||
| seq_length = inp.shape[1] | ||
| if isinstance(inp, nn.utils.rnn.PackedSequence): |
There was a problem hiding this comment.
Add the required sign-off trailer
This is a non-merge commit, but its message has no Signed-off-by: trailer, so it does not satisfy the repository's commit and CI requirements. Recreate the commit with git commit --signoff using the configured author name and email.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
_rnn_forward_hookreadsinp.shape[0]andinp.shape[1], and nn.RNN, nn.GRU and nn.LSTM all accept a PackedSequence there, so profiling a packed-input RNN raises the AttributeError you hit.Slight correction to the patch in the issue:
inp.datais(sum(lengths), input_size), sodata.shape[1]is the input size and not a sequence length. Using it multiplies the count by that factor. On a 2-layerLSTM(8, 16)with 4 sequences of length 5 it gives a multiplier of 160 where 20 is right, so the crash turns into a silent 8x overcount.What the hook wants is the number of cell steps, and a PackedSequence runs one per row of
.data. Sobatch_size = inp.data.shape[0]andseq_length = 1.Added a test over one model with three inputs: the dense batch, the same batch packed at full length, and lengths
[5, 3, 2, 1]. Full-length packing comes out equal to the dense count (149760), and the ragged one at 11 of its 20 steps (82368). The two packed cases fail on master. The dense case is there as the reference the other two are checked against, and the RNN hook had no test before this.Only ran it on CPU.
Fixes #4333