Add Conformer RNN-T model prototype #2322
Conversation
hwangjeff
force-pushed
the
conformer_rnnt_prototype
branch
from
2245f8c to
444a87d
Compare
hwangjeff
force-pushed
the
conformer_rnnt_prototype
branch
2 times, most recently
from
eb2c046 to
c255d6b
Compare
…d _ConformerTranscriber
hwangjeff
force-pushed
the
conformer_rnnt_prototype
branch
from
c255d6b to
4e31ff2
Compare
hwangjeff
force-pushed
the
conformer_rnnt_prototype
branch
from
5619f61 to
490a981
Compare
hwangjeff
force-pushed
the
conformer_rnnt_prototype
branch
2 times, most recently
from
ef91d94 to
0514f97
Compare
hwangjeff
force-pushed
the
conformer_rnnt_prototype
branch
from
0514f97 to
84e6107
Compare
hwangjeff
requested review from
nateanl,
mthrok,
xiaohui-zhang and
carolineechen
|
otherwise the PR looks good |
torchaudio/prototype/models/rnnt.py
Outdated
| from torchaudio.models.rnnt import _Joiner, _Predictor, _TimeReduction, _Transcriber | ||
|
|
||
|
|
||
| class _ConformerTranscriber(torch.nn.Module, _Transcriber): |
There was a problem hiding this comment.
encoder might be a more standard name than transcriber?
There was a problem hiding this comment.
i can rename the class and associated variables in another pr, as there may be bc-breaking implications
There was a problem hiding this comment.
actually, renamed the concrete classes here — will handle the abstract class separately
hwangjeff
force-pushed
the
conformer_rnnt_prototype
branch
from
2bdbe42 to
1c44c67
Compare
|
@hwangjeff has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
1 similar comment
|
@hwangjeff has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
Hey @hwangjeff. |
| lstm_dropout: int, | ||
| joiner_activation: str, | ||
| ) -> RNNT: | ||
| r"""Builds Conformer-based recurrent neural network transducer (RNN-T) model. |
There was a problem hiding this comment.
Need to overwrite the return signature otherwise the documentation shows wrong module path.
|
|
||
|
|
||
| class ConformerRNNTTestImpl(TestBaseMixin): | ||
| def _get_input_config(self): |
There was a problem hiding this comment.
nit: I'd move these helper functions to module-level, so that the the code around the test class is more focused on test logic.
Summary: Adds Conformer RNN-T model as prototype feature, by way of factory functions `conformer_rnnt_model` and `conformer_rnnt_base`, which instantiates a baseline version of the model. Also includes the following: - Modifies `Conformer` to accept arguments `use_group_norm` and `convolution_first` to pass to each of its `ConformerLayer` instances. - Makes `_Predictor` an abstract class and introduces `_EmformerEncoder` and `_ConformerEncoder`. - Introduces tests for `conformer_rnnt_model`. - Adds docs. Pull Request resolved: pytorch#2322 Reviewed By: xiaohui-zhang Differential Revision: D35565987 Pulled By: hwangjeff fbshipit-source-id: cb37bb0477ae3d5fcf0b7124f334f4cbb89b5789
Adds Conformer RNN-T model as prototype feature, by way of factory functions
conformer_rnnt_modelandconformer_rnnt_base, which instantiates a baseline version of the model. Also includes the following:Conformerto accept argumentsuse_group_normandconvolution_firstto pass to each of itsConformerLayerinstances._Predictoran abstract class and introduces_EmformerEncoderand_ConformerEncoder.conformer_rnnt_model.