Add RNN Transducer Loss for CPU by vincentqb · Pull Request #1137 · pytorch/audio

vincentqb · 2020-12-30T18:18:12Z

This pull request introduces rnnt_loss and RNNTLoss as a prototype in torchaudio.prototype.transducer using HawkAaron's warp-transducer.

The python interface remains currently the same as the original.
This has been tested as integrated within ESPNet here.

Follow-up work detailed in #1240.

cc @astaff, internal, #1099

cpuhrsch · 2020-12-30T19:33:23Z

+def rnnt_loss(acts, labels, act_lens, label_lens, blank=0, reduction="mean"):
+    """RNN Transducer Loss
+
+    Args:


I think the documentation could be improved a bit. It could also be useful to reference the paper.

mthrok · 2020-12-30T20:45:49Z

        super().build_extension(ext)
+
+
+_TRANSDUCER_NAME = '_warp_transducer'


This will get installed in global namespace, outside of torchaudio package directory.
Please put it in torchaudio package.

mthrok · 2020-12-30T21:01:56Z

+    MESSAGE(STATUS "Building static library with GPU support")
+
+    CUDA_ADD_LIBRARY(warprnnt STATIC submodule/src/rnnt_entrypoint.cu)
+    IF (!Torch_FOUND)


If torch is not found, shouldn't it be failing?

cpuhrsch · 2020-12-30T21:10:02Z

+        self.reduction = reduction
+        self.loss = _RNNT.apply
+
+    def forward(self, acts, labels, act_lens, label_lens):


If you don't want to copy-paste the docs from the functional you could reference it here within the documentation.

vincentqb · 2020-12-31T23:34:15Z

Some TODOs:

Use TORCH_CHECK to raise an error instead of writing to standard output, here, e.g. used here.
Investigate the memory allocation for workspace, here.

Some follow-ups:

Move libsox to a third_party subfolder as suggested above.
Investigate using AT_DISPATCH_FLOATING_TYPES.
Migrate the checks to C++.
Add GPU implementation and compilation.
Patch the submodule to remove the pytorch deprecation warnings.
Refactor, see internal.

cpuhrsch · 2021-01-04T17:51:41Z

+    # Test if example provided in README runs
+    # https://github.com/HawkAaron/warp-transducer
+
+    acts = torch.FloatTensor(


nit: use the factory function torch.tensor([xyz], dtype=torch.float) instead of the type constructor. Same applies to IntTensor.

cpuhrsch · 2021-01-04T22:19:54Z

+                U = data["tgt_lengths"][b]
+                for t in range(gradients.shape[1]):
+                    for u in range(gradients.shape[2]):
+                        np.testing.assert_allclose(


self.assertEqual should be preferred

vincentqb · 2021-01-05T05:41:08Z

Some more TODOs:

Remove numpy from tests
Guard prototype csrc and extension

Some more follow-ups:

Pass along the DEBUG flag to cmake
Guard prototype python files by omitting them from torchaudio, see also comment
Guard building third party transducer even if not added as an extension
Enable building transducer in nightlies only, not release.
Build within same folders as libsox, comment and comment
Remove hardcoded O2/O3 optimization, see comment

Error below also happens on master:

conda_build.exceptions.DependencyNeedsBuildingError: Unsatisfiable dependencies for platform osx-64: {"python[version='>=2.7,<2.8.0a0|>=3.8,<3.9.0a0|>=3.5,<3.6.0a0|>=3.5']", "python[version='>=3.6,<3.7.0a0|>=3.7,<3.8.0a0|>=3.9,<3.10.0a0']", "python_abi[version='3.6.*|3.7.*',build='*_cp36m|*_cp37m']", "python[version='>=2.7,<2.8.0a0|>=3.6,<3.7.0a0|>=3.7,<3.8.0a0|>=3.9,<3.10.0a0|>=3.8,<3.9.0a0|>=3.5,<3.6.0a0']", 'python', 'python_abi=3.9[build=*_cp39]'}

cc comment above

cpuhrsch · 2021-01-05T17:46:14Z

+        loss = rnnt_loss(acts, labels, act_length, label_length)
+        loss.backward()
+
+    def _test_costs_and_gradients(


This could be inlined since it only has one call-site and is pretty small (but that's not the reason to remove an abstraction necessarily).

mthrok · 2021-01-06T16:02:55Z

+}
+
+TORCH_LIBRARY_IMPL(torchaudio, CPU, m) {
+    m.impl("rnnt_loss", &cpu_rnnt_loss);


@vincentqb Can you define a proper namespace? torchaudio::<something>::rnnt_loss

I am not sure how you want to move on, but if you have a plan to add different type of rnnt, then more descriptive name would work better later, like warprnnt

Adding anonymous namespace in #1159 for the time being.

mthrok · 2021-01-07T16:37:12Z

@vincentqb I update the followup description for things addressed in #1159 and #1161. Please stamp these PRs when you have time.

For Enable building transducer in nightlies only, disable for release. I am thinking to add master environment variable that has higher precedence than BUILD_TRANSDUCER, like DISABLE_PROTOTYPE and propagating it from CCI configuration.

mthrok · 2021-01-09T02:23:38Z

For C++ ABI issue ssee #880

* fdsa * Tutorial runs * clarify one scaler per convergence run * adjust sizes, dont run illustrative sections * satisfying ocd * MORE * fdsa * details * rephrase * fix formatting * move script to recipes * hopefully moved to recipes * fdsa * add amp_tutorial to toctree * amp_tutorial -> amp_recipe * looks like backtick highlights dont render in card_description * correct path for amp_recipe.html * arch notes and saving/restoring * formatting * fdsa * Clarify autograd-autocast interaction for custom ops * touchups Co-authored-by: Brian Johnson <brianjo@fb.com>

facebook-github-bot added the CLA Signed label Dec 30, 2020

vincentqb force-pushed the transducer-cpu branch 2 times, most recently from 9d6589a to e2e6562 Compare December 30, 2020 18:57

cpuhrsch changed the title ~~Add RNN Transducer Loss~~ Add RNN Transducer Loss for CPU Dec 30, 2020

cpuhrsch reviewed Dec 30, 2020

View reviewed changes

Comment thread build_tools/setup_helpers/extension.py Outdated

cpuhrsch reviewed Dec 30, 2020

View reviewed changes

Comment thread third_party/warp_transducer/CMakeLists.txt Outdated

cpuhrsch reviewed Dec 30, 2020

View reviewed changes

Comment thread third_party/warp_transducer/CMakeLists.txt Outdated

cpuhrsch reviewed Dec 30, 2020

View reviewed changes

Comment thread torchaudio/prototype/transducer.py Outdated

cpuhrsch reviewed Dec 30, 2020

View reviewed changes

Comment thread third_party/warp_transducer/CMakeLists.txt Outdated