add more ctc-loss related summaries #457
Conversation
axlearn/audio/asr_decoder.py
Outdated
| jnp.mean(source_lengths), batch_size | ||
| ), | ||
| "input_stats/frame_packing_effiency": WeightedScalar( | ||
| jnp.sum(source_lengths) / input_batch["paddings"].size, input_batch["paddings"].size |
There was a problem hiding this comment.
Guard against division by 0 here and below?
There was a problem hiding this comment.
I think this before: the denominator input_batch['paddings'].size is directly pulled from the input. If it is zero, many other places will become problematic well before we hit here (e.g., various normalization layer). Nevertheless, I have added safeguard for them.
There was a problem hiding this comment.
Thanks! I don't think we should make assumptions about the other components.
| ), | ||
| } | ||
| # pytype: enable=attribute-error | ||
| return ret_dict |
There was a problem hiding this comment.
Is there a specific reason to prefer returning the summaries instead of just adding them here?
There was a problem hiding this comment.
By returning a dictionary, we can possibly reuse the summaries value in the base class for a more specific summaries in the subclasss. Also it allows us to override some summaries in the subclass.
axlearn/audio/asr_decoder.py
Outdated
|
|
||
| def _input_stats_summary( |
There was a problem hiding this comment.
| def _input_stats_summary( | |
| def _input_stats_summaries( |
or def _add_input_stats_summaries if we decide to inline the add, which may be more similar to other callsites in the repo.
There was a problem hiding this comment.
Have the changes been pushed?
There was a problem hiding this comment.
sorry. was confused and pushed to the other repo. Just pushed the change.
axlearn/audio/asr_decoder.py
Outdated
| per_frame_loss = total_ctc_loss / num_valid_frames | ||
| per_label_loss = total_ctc_loss / num_valid_labels | ||
| batch_size = per_example_weight.shape[0] |
| per_frame_loss = total_ctc_loss / num_valid_frames | ||
| per_label_loss = total_ctc_loss / num_valid_labels |
axlearn/audio/asr_decoder.py
Outdated
| valid_label_mask = (1.0 - target_paddings) * per_example_weight[:, None] | ||
| num_valid_frames = jnp.sum(valid_frame_mask) | ||
| num_valid_labels = jnp.sum(valid_label_mask) | ||
| num_valid_examples = jnp.maximum(per_example_weight.sum(), 1.0) |
There was a problem hiding this comment.
A couple nits -- since we sum over weights, 1.0 may not always be appropriate. We might also consider renaming num_valid_examples to total_example_weight.
There was a problem hiding this comment.
Thanks for the suggestion.
yqwangustc
requested review from
zhiyun and
markblee
and removed request for
zhiyun
Add more detailed loss to easy debugging