diff --git a/.flake8 b/.flake8 new file mode 100644 index 0000000000..090e97971b --- /dev/null +++ b/.flake8 @@ -0,0 +1,8 @@ +[flake8] +show-source=true +statistics=true +max-line-length = 80 + +exclude = + .git, + **/data/** diff --git a/.github/workflows/style_check.yml b/.github/workflows/style_check.yml new file mode 100644 index 0000000000..20c3363b41 --- /dev/null +++ b/.github/workflows/style_check.yml @@ -0,0 +1,62 @@ +# Copyright 2021 Fangjun Kuang (csukuangfj@gmail.com) + +# See ../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name: style_check + +on: + push: + branches: + - master + pull_request: + branches: + - master + +jobs: + style_check: + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ubuntu-18.04, macos-10.15] + python-version: [3.7, 3.9] + fail-fast: false + + steps: + - uses: actions/checkout@v2 + with: + fetch-depth: 0 + + - name: Setup Python ${{ matrix.python-version }} + uses: actions/setup-python@v1 + with: + python-version: ${{ matrix.python-version }} + + - name: Install Python dependencies + run: | + python3 -m pip install --upgrade pip black flake8 + + - name: Run flake8 + shell: bash + working-directory: ${{github.workspace}} + run: | + # stop the build if there are Python syntax errors or undefined names + flake8 . --count --show-source --statistics + flake8 . + + - name: Run black + shell: bash + working-directory: ${{github.workspace}} + run: | + black --check --diff . diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000000..9a298877a5 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,77 @@ +# Copyright 2021 Fangjun Kuang (csukuangfj@gmail.com) + +# See ../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name: test + +on: + push: + branches: + - master + pull_request: + branches: + - master + +jobs: + test: + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ubuntu-18.04, macos-10.15] + python-version: [3.6, 3.7, 3.8, 3.9] + torch: ["1.8.1"] + k2-version: ["1.2.dev20210724"] + fail-fast: false + + steps: + - uses: actions/checkout@v2 + with: + fetch-depth: 0 + + - name: Setup Python ${{ matrix.python-version }} + uses: actions/setup-python@v1 + with: + python-version: ${{ matrix.python-version }} + + - name: Install Python dependencies + run: | + python3 -m pip install --upgrade pip pytest kaldialign + pip install k2==${{ matrix.k2-version }}+cpu.torch${{ matrix.torch }} -f https://k2-fsa.org/nightly/ + + # Don't use: pip install lhotse + # since it installs a version of PyTorch that is not predictable + git clone --depth 1 https://github.com/lhotse-speech/lhotse + cd lhotse + sed -i.bak "/torch/d" requirements.txt + pip install -r ./requirements.txt + + + - name: Run tests + if: startsWith(matrix.os, 'ubuntu') + run: | + ls -lh + export PYTHONPATH=$PWD:$PWD/lhotse:$PYTHONPATH + echo $PYTHONPATH + pytest ./test + + - name: Run tests + if: startsWith(matrix.os, 'macos') + run: | + ls -lh + export PYTHONPATH=$PWD:$PWD/lhotse:$PYTHONPATH + lib_path=$(python -c "from distutils.sysconfig import get_python_lib; print(get_python_lib())") + echo "lib_path: $lib_path" + export DYLD_LIBRARY_PATH=$lib_path:$DYLD_LIBRARY_PATH + pytest ./test diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000..6cb9f22997 --- /dev/null +++ b/.gitignore @@ -0,0 +1,6 @@ +data +__pycache__ +path.sh +exp +exp*/ +*.pt diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000000..b59784dbf4 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,26 @@ +repos: + - repo: https://github.com/psf/black + rev: 21.6b0 + hooks: + - id: black + args: [--line-length=80] + + - repo: https://github.com/PyCQA/flake8 + rev: 3.9.2 + hooks: + - id: flake8 + args: [--max-line-length=80] + + - repo: https://github.com/pycqa/isort + rev: 5.9.2 + hooks: + - id: isort + args: [--profile=black, --line-length=80] + + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.0.1 + hooks: + - id: check-executables-have-shebangs + - id: end-of-file-fixer + - id: mixed-line-ending + - id: trailing-whitespace diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000..ee06cfc772 --- /dev/null +++ b/LICENSE @@ -0,0 +1,211 @@ + + Legal Notices + + NOTE (this is not from the Apache License): The copyright model is that + authors (or their employers, if noted in individual files) own their + individual contributions. The authors' contributions can be discerned + from the git history. + + ------------------------------------------------------------------------- + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/contributing.md b/contributing.md new file mode 100644 index 0000000000..c8f06fdaef --- /dev/null +++ b/contributing.md @@ -0,0 +1,39 @@ + +## Pre-commit hooks + +We use [git][git] [pre-commit][pre-commit] [hooks][hooks] to check that files +going to be committed: + + - contain no trailing spaces + - are formatted with [black][black] + - are compatible to [PEP8][PEP8] (checked by [flake8][flake8]) + - end in a newline and only a newline + - contain sorted `imports` (checked by [isort][isort]) + +These hooks are disabled by default. Please use the following commands to enable them: + +```bash +pip install pre-commit # run it only once +pre-commit install # run it only once, it will install all hooks + +# modify some files +git add +git commit # It runs all hooks automatically. + +# If all hooks run successfully, you can write the commit message now. Done! +# +# If any hook failed, your commit was not successful. +# Please read the error messages and make changes accordingly. +# And rerun + +git add +git commit +``` + +[git]: https://git-scm.com/book/en/v2/Customizing-Git-Git-Hooks +[flake8]: https://github.com/PyCQA/flake8 +[PEP8]: https://www.python.org/dev/peps/pep-0008/ +[black]: https://github.com/psf/black +[hooks]: https://github.com/pre-commit/pre-commit-hooks +[pre-commit]: https://github.com/pre-commit/pre-commit +[isort]: https://github.com/PyCQA/isort diff --git a/egs/librispeech/ASR/README.md b/egs/librispeech/ASR/README.md new file mode 100644 index 0000000000..45c9ef4dec --- /dev/null +++ b/egs/librispeech/ASR/README.md @@ -0,0 +1,121 @@ + +Run `./prepare.sh` to prepare the data. + +Run `./xxx_train.py` (to be added) to train a model. + +## Conformer-CTC +Results of the pre-trained model from +`` +are given below + +### HLG - no LM rescoring + +(output beam size is 8) + +#### 1-best decoding + +``` +[test-clean-no_rescore] %WER 3.15% [1656 / 52576, 127 ins, 377 del, 1152 sub ] +[test-other-no_rescore] %WER 7.03% [3682 / 52343, 220 ins, 1024 del, 2438 sub ] +``` + +#### n-best decoding + +For n=100, + +``` +[test-clean-no_rescore-100] %WER 3.15% [1656 / 52576, 127 ins, 377 del, 1152 sub ] +[test-other-no_rescore-100] %WER 7.14% [3737 / 52343, 275 ins, 1020 del, 2442 sub ] +``` + +For n=200, + +``` +[test-clean-no_rescore-200] %WER 3.16% [1660 / 52576, 125 ins, 378 del, 1157 sub ] +[test-other-no_rescore-200] %WER 7.04% [3684 / 52343, 228 ins, 1012 del, 2444 sub ] +``` + +### HLG - with LM rescoring + +#### Whole lattice rescoring + +``` +[test-clean-lm_scale_0.8] %WER 2.77% [1456 / 52576, 150 ins, 210 del, 1096 sub ] +[test-other-lm_scale_0.8] %WER 6.23% [3262 / 52343, 246 ins, 635 del, 2381 sub ] +``` + +WERs of different LM scales are: + +``` +For test-clean, WER of different settings are: +lm_scale_0.8 2.77 best for test-clean +lm_scale_0.9 2.87 +lm_scale_1.0 3.06 +lm_scale_1.1 3.34 +lm_scale_1.2 3.71 +lm_scale_1.3 4.18 +lm_scale_1.4 4.8 +lm_scale_1.5 5.48 +lm_scale_1.6 6.08 +lm_scale_1.7 6.79 +lm_scale_1.8 7.49 +lm_scale_1.9 8.14 +lm_scale_2.0 8.82 + +For test-other, WER of different settings are: +lm_scale_0.8 6.23 best for test-other +lm_scale_0.9 6.37 +lm_scale_1.0 6.62 +lm_scale_1.1 6.99 +lm_scale_1.2 7.46 +lm_scale_1.3 8.13 +lm_scale_1.4 8.84 +lm_scale_1.5 9.61 +lm_scale_1.6 10.32 +lm_scale_1.7 11.17 +lm_scale_1.8 12.12 +lm_scale_1.9 12.93 +lm_scale_2.0 13.77 +``` + +#### n-best LM rescoring + +n = 100 + +``` +[test-clean-lm_scale_0.8] %WER 2.79% [1469 / 52576, 149 ins, 212 del, 1108 sub ] +[test-other-lm_scale_0.8] %WER 6.36% [3329 / 52343, 259 ins, 666 del, 2404 sub ] +``` + +WERs of different LM scales are: + +``` +For test-clean, WER of different settings are: +lm_scale_0.8 2.79 best for test-clean +lm_scale_0.9 2.89 +lm_scale_1.0 3.03 +lm_scale_1.1 3.28 +lm_scale_1.2 3.52 +lm_scale_1.3 3.78 +lm_scale_1.4 4.04 +lm_scale_1.5 4.24 +lm_scale_1.6 4.45 +lm_scale_1.7 4.58 +lm_scale_1.8 4.7 +lm_scale_1.9 4.8 +lm_scale_2.0 4.92 +For test-other, WER of different settings are: +lm_scale_0.8 6.36 best for test-other +lm_scale_0.9 6.45 +lm_scale_1.0 6.64 +lm_scale_1.1 6.92 +lm_scale_1.2 7.25 +lm_scale_1.3 7.59 +lm_scale_1.4 7.88 +lm_scale_1.5 8.13 +lm_scale_1.6 8.36 +lm_scale_1.7 8.54 +lm_scale_1.8 8.71 +lm_scale_1.9 8.88 +lm_scale_2.0 9.02 +``` diff --git a/egs/librispeech/ASR/conformer_ctc/__init__.py b/egs/librispeech/ASR/conformer_ctc/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/egs/librispeech/ASR/conformer_ctc/conformer.py b/egs/librispeech/ASR/conformer_ctc/conformer.py new file mode 100644 index 0000000000..1e82eff2fa --- /dev/null +++ b/egs/librispeech/ASR/conformer_ctc/conformer.py @@ -0,0 +1,914 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2021 University of Chinese Academy of Sciences (author: Han Zhu) +# Apache 2.0 + +import math +import warnings +from typing import Optional, Tuple + +import torch +from torch import Tensor, nn +from transformer import Supervisions, Transformer, encoder_padding_mask + + +class Conformer(Transformer): + """ + Args: + num_features (int): Number of input features + num_classes (int): Number of output classes + subsampling_factor (int): subsampling factor of encoder (the convolution layers before transformers) + d_model (int): attention dimension + nhead (int): number of head + dim_feedforward (int): feedforward dimention + num_encoder_layers (int): number of encoder layers + num_decoder_layers (int): number of decoder layers + dropout (float): dropout rate + cnn_module_kernel (int): Kernel size of convolution module + normalize_before (bool): whether to use layer_norm before the first block. + vgg_frontend (bool): whether to use vgg frontend. + """ + + def __init__( + self, + num_features: int, + num_classes: int, + subsampling_factor: int = 4, + d_model: int = 256, + nhead: int = 4, + dim_feedforward: int = 2048, + num_encoder_layers: int = 12, + num_decoder_layers: int = 6, + dropout: float = 0.1, + cnn_module_kernel: int = 31, + normalize_before: bool = True, + vgg_frontend: bool = False, + is_espnet_structure: bool = False, + mmi_loss: bool = True, + use_feat_batchnorm: bool = False, + ) -> None: + super(Conformer, self).__init__( + num_features=num_features, + num_classes=num_classes, + subsampling_factor=subsampling_factor, + d_model=d_model, + nhead=nhead, + dim_feedforward=dim_feedforward, + num_encoder_layers=num_encoder_layers, + num_decoder_layers=num_decoder_layers, + dropout=dropout, + normalize_before=normalize_before, + vgg_frontend=vgg_frontend, + mmi_loss=mmi_loss, + use_feat_batchnorm=use_feat_batchnorm, + ) + + self.encoder_pos = RelPositionalEncoding(d_model, dropout) + + encoder_layer = ConformerEncoderLayer( + d_model, + nhead, + dim_feedforward, + dropout, + cnn_module_kernel, + normalize_before, + is_espnet_structure, + ) + self.encoder = ConformerEncoder(encoder_layer, num_encoder_layers) + self.normalize_before = normalize_before + self.is_espnet_structure = is_espnet_structure + if self.normalize_before and self.is_espnet_structure: + self.after_norm = nn.LayerNorm(d_model) + else: + # Note: TorchScript detects that self.after_norm could be used inside forward() + # and throws an error without this change. + self.after_norm = identity + + def encode( + self, x: Tensor, supervisions: Optional[Supervisions] = None + ) -> Tuple[Tensor, Optional[Tensor]]: + """ + Args: + x: Tensor of dimension (batch_size, num_features, input_length). + supervisions : Supervison in lhotse format, i.e., batch['supervisions'] + + Returns: + Tensor: Predictor tensor of dimension (input_length, batch_size, d_model). + Tensor: Mask tensor of dimension (batch_size, input_length) + """ + x = x.permute(0, 2, 1) # (B, F, T) -> (B, T, F) + + x = self.encoder_embed(x) + x, pos_emb = self.encoder_pos(x) + x = x.permute(1, 0, 2) # (B, T, F) -> (T, B, F) + mask = encoder_padding_mask(x.size(0), supervisions) + if mask is not None: + mask = mask.to(x.device) + x = self.encoder(x, pos_emb, src_key_padding_mask=mask) # (T, B, F) + + if self.normalize_before and self.is_espnet_structure: + x = self.after_norm(x) + + return x, mask + + +class ConformerEncoderLayer(nn.Module): + """ + ConformerEncoderLayer is made up of self-attn, feedforward and convolution networks. + See: "Conformer: Convolution-augmented Transformer for Speech Recognition" + + Args: + d_model: the number of expected features in the input (required). + nhead: the number of heads in the multiheadattention models (required). + dim_feedforward: the dimension of the feedforward network model (default=2048). + dropout: the dropout value (default=0.1). + cnn_module_kernel (int): Kernel size of convolution module. + normalize_before: whether to use layer_norm before the first block. + + Examples:: + >>> encoder_layer = ConformerEncoderLayer(d_model=512, nhead=8) + >>> src = torch.rand(10, 32, 512) + >>> pos_emb = torch.rand(32, 19, 512) + >>> out = encoder_layer(src, pos_emb) + """ + + def __init__( + self, + d_model: int, + nhead: int, + dim_feedforward: int = 2048, + dropout: float = 0.1, + cnn_module_kernel: int = 31, + normalize_before: bool = True, + is_espnet_structure: bool = False, + ) -> None: + super(ConformerEncoderLayer, self).__init__() + self.self_attn = RelPositionMultiheadAttention( + d_model, nhead, dropout=0.0, is_espnet_structure=is_espnet_structure + ) + + self.feed_forward = nn.Sequential( + nn.Linear(d_model, dim_feedforward), + Swish(), + nn.Dropout(dropout), + nn.Linear(dim_feedforward, d_model), + ) + + self.feed_forward_macaron = nn.Sequential( + nn.Linear(d_model, dim_feedforward), + Swish(), + nn.Dropout(dropout), + nn.Linear(dim_feedforward, d_model), + ) + + self.conv_module = ConvolutionModule(d_model, cnn_module_kernel) + + self.norm_ff_macaron = nn.LayerNorm( + d_model + ) # for the macaron style FNN module + self.norm_ff = nn.LayerNorm(d_model) # for the FNN module + self.norm_mha = nn.LayerNorm(d_model) # for the MHA module + + self.ff_scale = 0.5 + + self.norm_conv = nn.LayerNorm(d_model) # for the CNN module + self.norm_final = nn.LayerNorm( + d_model + ) # for the final output of the block + + self.dropout = nn.Dropout(dropout) + + self.normalize_before = normalize_before + + def forward( + self, + src: Tensor, + pos_emb: Tensor, + src_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + ) -> Tensor: + """ + Pass the input through the encoder layer. + + Args: + src: the sequence to the encoder layer (required). + pos_emb: Positional embedding tensor (required). + src_mask: the mask for the src sequence (optional). + src_key_padding_mask: the mask for the src keys per batch (optional). + + Shape: + src: (S, N, E). + pos_emb: (N, 2*S-1, E) + src_mask: (S, S). + src_key_padding_mask: (N, S). + S is the source sequence length, N is the batch size, E is the feature number + """ + + # macaron style feed forward module + residual = src + if self.normalize_before: + src = self.norm_ff_macaron(src) + src = residual + self.ff_scale * self.dropout( + self.feed_forward_macaron(src) + ) + if not self.normalize_before: + src = self.norm_ff_macaron(src) + + # multi-headed self-attention module + residual = src + if self.normalize_before: + src = self.norm_mha(src) + src_att = self.self_attn( + src, + src, + src, + pos_emb=pos_emb, + attn_mask=src_mask, + key_padding_mask=src_key_padding_mask, + )[0] + src = residual + self.dropout(src_att) + if not self.normalize_before: + src = self.norm_mha(src) + + # convolution module + residual = src + if self.normalize_before: + src = self.norm_conv(src) + src = residual + self.dropout(self.conv_module(src)) + if not self.normalize_before: + src = self.norm_conv(src) + + # feed forward module + residual = src + if self.normalize_before: + src = self.norm_ff(src) + src = residual + self.ff_scale * self.dropout(self.feed_forward(src)) + if not self.normalize_before: + src = self.norm_ff(src) + + if self.normalize_before: + src = self.norm_final(src) + + return src + + +class ConformerEncoder(nn.TransformerEncoder): + r"""ConformerEncoder is a stack of N encoder layers + + Args: + encoder_layer: an instance of the ConformerEncoderLayer() class (required). + num_layers: the number of sub-encoder-layers in the encoder (required). + norm: the layer normalization component (optional). + + Examples:: + >>> encoder_layer = ConformerEncoderLayer(d_model=512, nhead=8) + >>> conformer_encoder = ConformerEncoder(encoder_layer, num_layers=6) + >>> src = torch.rand(10, 32, 512) + >>> pos_emb = torch.rand(32, 19, 512) + >>> out = conformer_encoder(src, pos_emb) + """ + + def __init__( + self, encoder_layer: nn.Module, num_layers: int, norm: nn.Module = None + ) -> None: + super(ConformerEncoder, self).__init__( + encoder_layer=encoder_layer, num_layers=num_layers, norm=norm + ) + + def forward( + self, + src: Tensor, + pos_emb: Tensor, + mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + ) -> Tensor: + r"""Pass the input through the encoder layers in turn. + + Args: + src: the sequence to the encoder (required). + pos_emb: Positional embedding tensor (required). + mask: the mask for the src sequence (optional). + src_key_padding_mask: the mask for the src keys per batch (optional). + + Shape: + src: (S, N, E). + pos_emb: (N, 2*S-1, E) + mask: (S, S). + src_key_padding_mask: (N, S). + S is the source sequence length, T is the target sequence length, N is the batch size, E is the feature number + + """ + output = src + + for mod in self.layers: + output = mod( + output, + pos_emb, + src_mask=mask, + src_key_padding_mask=src_key_padding_mask, + ) + + if self.norm is not None: + output = self.norm(output) + + return output + + +class RelPositionalEncoding(torch.nn.Module): + """Relative positional encoding module. + + See : Appendix B in "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context" + Modified from https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/transformer/embedding.py + + Args: + d_model: Embedding dimension. + dropout_rate: Dropout rate. + max_len: Maximum input length. + + """ + + def __init__( + self, d_model: int, dropout_rate: float, max_len: int = 5000 + ) -> None: + """Construct an PositionalEncoding object.""" + super(RelPositionalEncoding, self).__init__() + self.d_model = d_model + self.xscale = math.sqrt(self.d_model) + self.dropout = torch.nn.Dropout(p=dropout_rate) + self.pe = None + self.extend_pe(torch.tensor(0.0).expand(1, max_len)) + + def extend_pe(self, x: Tensor) -> None: + """Reset the positional encodings.""" + if self.pe is not None: + # self.pe contains both positive and negative parts + # the length of self.pe is 2 * input_len - 1 + if self.pe.size(1) >= x.size(1) * 2 - 1: + # Note: TorchScript doesn't implement operator== for torch.Device + if self.pe.dtype != x.dtype or str(self.pe.device) != str( + x.device + ): + self.pe = self.pe.to(dtype=x.dtype, device=x.device) + return + # Suppose `i` means to the position of query vecotr and `j` means the + # position of key vector. We use position relative positions when keys + # are to the left (i>j) and negative relative positions otherwise (i Tuple[Tensor, Tensor]: + """Add positional encoding. + + Args: + x (torch.Tensor): Input tensor (batch, time, `*`). + + Returns: + torch.Tensor: Encoded tensor (batch, time, `*`). + torch.Tensor: Encoded tensor (batch, 2*time-1, `*`). + + """ + self.extend_pe(x) + x = x * self.xscale + pos_emb = self.pe[ + :, + self.pe.size(1) // 2 + - x.size(1) + + 1 : self.pe.size(1) // 2 + + x.size(1), + ] + return self.dropout(x), self.dropout(pos_emb) + + +class RelPositionMultiheadAttention(nn.Module): + r"""Multi-Head Attention layer with relative position encoding + + See reference: "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context" + + Args: + embed_dim: total dimension of the model. + num_heads: parallel attention heads. + dropout: a Dropout layer on attn_output_weights. Default: 0.0. + + Examples:: + + >>> rel_pos_multihead_attn = RelPositionMultiheadAttention(embed_dim, num_heads) + >>> attn_output, attn_output_weights = multihead_attn(query, key, value, pos_emb) + """ + + def __init__( + self, + embed_dim: int, + num_heads: int, + dropout: float = 0.0, + is_espnet_structure: bool = False, + ) -> None: + super(RelPositionMultiheadAttention, self).__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.dropout = dropout + self.head_dim = embed_dim // num_heads + assert ( + self.head_dim * num_heads == self.embed_dim + ), "embed_dim must be divisible by num_heads" + + self.in_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=True) + self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True) + + # linear transformation for positional encoding. + self.linear_pos = nn.Linear(embed_dim, embed_dim, bias=False) + # these two learnable bias are used in matrix c and matrix d + # as described in "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context" Section 3.3 + self.pos_bias_u = nn.Parameter(torch.Tensor(num_heads, self.head_dim)) + self.pos_bias_v = nn.Parameter(torch.Tensor(num_heads, self.head_dim)) + + self._reset_parameters() + + self.is_espnet_structure = is_espnet_structure + + def _reset_parameters(self) -> None: + nn.init.xavier_uniform_(self.in_proj.weight) + nn.init.constant_(self.in_proj.bias, 0.0) + nn.init.constant_(self.out_proj.bias, 0.0) + + nn.init.xavier_uniform_(self.pos_bias_u) + nn.init.xavier_uniform_(self.pos_bias_v) + + def forward( + self, + query: Tensor, + key: Tensor, + value: Tensor, + pos_emb: Tensor, + key_padding_mask: Optional[Tensor] = None, + need_weights: bool = True, + attn_mask: Optional[Tensor] = None, + ) -> Tuple[Tensor, Optional[Tensor]]: + r""" + Args: + query, key, value: map a query and a set of key-value pairs to an output. + pos_emb: Positional embedding tensor + key_padding_mask: if provided, specified padding elements in the key will + be ignored by the attention. When given a binary mask and a value is True, + the corresponding value on the attention layer will be ignored. When given + a byte mask and a value is non-zero, the corresponding value on the attention + layer will be ignored + need_weights: output attn_output_weights. + attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all + the batches while a 3D mask allows to specify a different mask for the entries of each batch. + + Shape: + - Inputs: + - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is + the embedding dimension. + - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is + the embedding dimension. + - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is + the embedding dimension. + - pos_emb: :math:`(N, 2*L-1, E)` where L is the target sequence length, N is the batch size, E is + the embedding dimension. + - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length. + If a ByteTensor is provided, the non-zero positions will be ignored while the position + with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the + value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged. + - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length. + 3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length, + S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked + positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend + while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True`` + is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor + is provided, it will be added to the attention weight. + + - Outputs: + - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, + E is the embedding dimension. + - attn_output_weights: :math:`(N, L, S)` where N is the batch size, + L is the target sequence length, S is the source sequence length. + """ + return self.multi_head_attention_forward( + query, + key, + value, + pos_emb, + self.embed_dim, + self.num_heads, + self.in_proj.weight, + self.in_proj.bias, + self.dropout, + self.out_proj.weight, + self.out_proj.bias, + training=self.training, + key_padding_mask=key_padding_mask, + need_weights=need_weights, + attn_mask=attn_mask, + ) + + def rel_shift(self, x: Tensor) -> Tensor: + """Compute relative positional encoding. + + Args: + x: Input tensor (batch, head, time1, 2*time1-1). + time1 means the length of query vector. + + Returns: + Tensor: tensor of shape (batch, head, time1, time2) + (note: time2 has the same value as time1, but it is for + the key, while time1 is for the query). + """ + (batch_size, num_heads, time1, n) = x.shape + assert n == 2 * time1 - 1 + # Note: TorchScript requires explicit arg for stride() + batch_stride = x.stride(0) + head_stride = x.stride(1) + time1_stride = x.stride(2) + n_stride = x.stride(3) + return x.as_strided( + (batch_size, num_heads, time1, time1), + (batch_stride, head_stride, time1_stride - n_stride, n_stride), + storage_offset=n_stride * (time1 - 1), + ) + + def multi_head_attention_forward( + self, + query: Tensor, + key: Tensor, + value: Tensor, + pos_emb: Tensor, + embed_dim_to_check: int, + num_heads: int, + in_proj_weight: Tensor, + in_proj_bias: Tensor, + dropout_p: float, + out_proj_weight: Tensor, + out_proj_bias: Tensor, + training: bool = True, + key_padding_mask: Optional[Tensor] = None, + need_weights: bool = True, + attn_mask: Optional[Tensor] = None, + ) -> Tuple[Tensor, Optional[Tensor]]: + r""" + Args: + query, key, value: map a query and a set of key-value pairs to an output. + pos_emb: Positional embedding tensor + embed_dim_to_check: total dimension of the model. + num_heads: parallel attention heads. + in_proj_weight, in_proj_bias: input projection weight and bias. + dropout_p: probability of an element to be zeroed. + out_proj_weight, out_proj_bias: the output projection weight and bias. + training: apply dropout if is ``True``. + key_padding_mask: if provided, specified padding elements in the key will + be ignored by the attention. This is an binary mask. When the value is True, + the corresponding value on the attention layer will be filled with -inf. + need_weights: output attn_output_weights. + attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all + the batches while a 3D mask allows to specify a different mask for the entries of each batch. + + Shape: + Inputs: + - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is + the embedding dimension. + - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is + the embedding dimension. + - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is + the embedding dimension. + - pos_emb: :math:`(N, 2*L-1, E)` or :math:`(1, 2*L-1, E)` where L is the target sequence + length, N is the batch size, E is the embedding dimension. + - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length. + If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions + will be unchanged. If a BoolTensor is provided, the positions with the + value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged. + - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length. + 3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length, + S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked + positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend + while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True`` + are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor + is provided, it will be added to the attention weight. + + Outputs: + - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, + E is the embedding dimension. + - attn_output_weights: :math:`(N, L, S)` where N is the batch size, + L is the target sequence length, S is the source sequence length. + """ + + tgt_len, bsz, embed_dim = query.size() + assert embed_dim == embed_dim_to_check + assert key.size(0) == value.size(0) and key.size(1) == value.size(1) + + head_dim = embed_dim // num_heads + assert ( + head_dim * num_heads == embed_dim + ), "embed_dim must be divisible by num_heads" + scaling = float(head_dim) ** -0.5 + + if torch.equal(query, key) and torch.equal(key, value): + # self-attention + q, k, v = nn.functional.linear( + query, in_proj_weight, in_proj_bias + ).chunk(3, dim=-1) + + elif torch.equal(key, value): + # encoder-decoder attention + # This is inline in_proj function with in_proj_weight and in_proj_bias + _b = in_proj_bias + _start = 0 + _end = embed_dim + _w = in_proj_weight[_start:_end, :] + if _b is not None: + _b = _b[_start:_end] + q = nn.functional.linear(query, _w, _b) + # This is inline in_proj function with in_proj_weight and in_proj_bias + _b = in_proj_bias + _start = embed_dim + _end = None + _w = in_proj_weight[_start:, :] + if _b is not None: + _b = _b[_start:] + k, v = nn.functional.linear(key, _w, _b).chunk(2, dim=-1) + + else: + # This is inline in_proj function with in_proj_weight and in_proj_bias + _b = in_proj_bias + _start = 0 + _end = embed_dim + _w = in_proj_weight[_start:_end, :] + if _b is not None: + _b = _b[_start:_end] + q = nn.functional.linear(query, _w, _b) + + # This is inline in_proj function with in_proj_weight and in_proj_bias + _b = in_proj_bias + _start = embed_dim + _end = embed_dim * 2 + _w = in_proj_weight[_start:_end, :] + if _b is not None: + _b = _b[_start:_end] + k = nn.functional.linear(key, _w, _b) + + # This is inline in_proj function with in_proj_weight and in_proj_bias + _b = in_proj_bias + _start = embed_dim * 2 + _end = None + _w = in_proj_weight[_start:, :] + if _b is not None: + _b = _b[_start:] + v = nn.functional.linear(value, _w, _b) + + if not self.is_espnet_structure: + q = q * scaling + + if attn_mask is not None: + assert ( + attn_mask.dtype == torch.float32 + or attn_mask.dtype == torch.float64 + or attn_mask.dtype == torch.float16 + or attn_mask.dtype == torch.uint8 + or attn_mask.dtype == torch.bool + ), "Only float, byte, and bool types are supported for attn_mask, not {}".format( + attn_mask.dtype + ) + if attn_mask.dtype == torch.uint8: + warnings.warn( + "Byte tensor for attn_mask is deprecated. Use bool tensor instead." + ) + attn_mask = attn_mask.to(torch.bool) + + if attn_mask.dim() == 2: + attn_mask = attn_mask.unsqueeze(0) + if list(attn_mask.size()) != [1, query.size(0), key.size(0)]: + raise RuntimeError( + "The size of the 2D attn_mask is not correct." + ) + elif attn_mask.dim() == 3: + if list(attn_mask.size()) != [ + bsz * num_heads, + query.size(0), + key.size(0), + ]: + raise RuntimeError( + "The size of the 3D attn_mask is not correct." + ) + else: + raise RuntimeError( + "attn_mask's dimension {} is not supported".format( + attn_mask.dim() + ) + ) + # attn_mask's dim is 3 now. + + # convert ByteTensor key_padding_mask to bool + if ( + key_padding_mask is not None + and key_padding_mask.dtype == torch.uint8 + ): + warnings.warn( + "Byte tensor for key_padding_mask is deprecated. Use bool tensor instead." + ) + key_padding_mask = key_padding_mask.to(torch.bool) + + q = q.contiguous().view(tgt_len, bsz, num_heads, head_dim) + k = k.contiguous().view(-1, bsz, num_heads, head_dim) + v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1) + + src_len = k.size(0) + + if key_padding_mask is not None: + assert key_padding_mask.size(0) == bsz, "{} == {}".format( + key_padding_mask.size(0), bsz + ) + assert key_padding_mask.size(1) == src_len, "{} == {}".format( + key_padding_mask.size(1), src_len + ) + + q = q.transpose(0, 1) # (batch, time1, head, d_k) + + pos_emb_bsz = pos_emb.size(0) + assert pos_emb_bsz in (1, bsz) # actually it is 1 + p = self.linear_pos(pos_emb).view(pos_emb_bsz, -1, num_heads, head_dim) + p = p.transpose(1, 2) # (batch, head, 2*time1-1, d_k) + + q_with_bias_u = (q + self.pos_bias_u).transpose( + 1, 2 + ) # (batch, head, time1, d_k) + + q_with_bias_v = (q + self.pos_bias_v).transpose( + 1, 2 + ) # (batch, head, time1, d_k) + + # compute attention score + # first compute matrix a and matrix c + # as described in "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context" Section 3.3 + k = k.permute(1, 2, 3, 0) # (batch, head, d_k, time2) + matrix_ac = torch.matmul( + q_with_bias_u, k + ) # (batch, head, time1, time2) + + # compute matrix b and matrix d + matrix_bd = torch.matmul( + q_with_bias_v, p.transpose(-2, -1) + ) # (batch, head, time1, 2*time1-1) + matrix_bd = self.rel_shift(matrix_bd) + + if not self.is_espnet_structure: + attn_output_weights = ( + matrix_ac + matrix_bd + ) # (batch, head, time1, time2) + else: + attn_output_weights = ( + matrix_ac + matrix_bd + ) * scaling # (batch, head, time1, time2) + + attn_output_weights = attn_output_weights.view( + bsz * num_heads, tgt_len, -1 + ) + + assert list(attn_output_weights.size()) == [ + bsz * num_heads, + tgt_len, + src_len, + ] + + if attn_mask is not None: + if attn_mask.dtype == torch.bool: + attn_output_weights.masked_fill_(attn_mask, float("-inf")) + else: + attn_output_weights += attn_mask + + if key_padding_mask is not None: + attn_output_weights = attn_output_weights.view( + bsz, num_heads, tgt_len, src_len + ) + attn_output_weights = attn_output_weights.masked_fill( + key_padding_mask.unsqueeze(1).unsqueeze(2), + float("-inf"), + ) + attn_output_weights = attn_output_weights.view( + bsz * num_heads, tgt_len, src_len + ) + + attn_output_weights = nn.functional.softmax(attn_output_weights, dim=-1) + attn_output_weights = nn.functional.dropout( + attn_output_weights, p=dropout_p, training=training + ) + + attn_output = torch.bmm(attn_output_weights, v) + assert list(attn_output.size()) == [bsz * num_heads, tgt_len, head_dim] + attn_output = ( + attn_output.transpose(0, 1) + .contiguous() + .view(tgt_len, bsz, embed_dim) + ) + attn_output = nn.functional.linear( + attn_output, out_proj_weight, out_proj_bias + ) + + if need_weights: + # average attention weights over heads + attn_output_weights = attn_output_weights.view( + bsz, num_heads, tgt_len, src_len + ) + return attn_output, attn_output_weights.sum(dim=1) / num_heads + else: + return attn_output, None + + +class ConvolutionModule(nn.Module): + """ConvolutionModule in Conformer model. + Modified from https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/conformer/convolution.py + + Args: + channels (int): The number of channels of conv layers. + kernel_size (int): Kernerl size of conv layers. + bias (bool): Whether to use bias in conv layers (default=True). + + """ + + def __init__( + self, channels: int, kernel_size: int, bias: bool = True + ) -> None: + """Construct an ConvolutionModule object.""" + super(ConvolutionModule, self).__init__() + # kernerl_size should be a odd number for 'SAME' padding + assert (kernel_size - 1) % 2 == 0 + + self.pointwise_conv1 = nn.Conv1d( + channels, + 2 * channels, + kernel_size=1, + stride=1, + padding=0, + bias=bias, + ) + self.depthwise_conv = nn.Conv1d( + channels, + channels, + kernel_size, + stride=1, + padding=(kernel_size - 1) // 2, + groups=channels, + bias=bias, + ) + self.norm = nn.BatchNorm1d(channels) + self.pointwise_conv2 = nn.Conv1d( + channels, + channels, + kernel_size=1, + stride=1, + padding=0, + bias=bias, + ) + self.activation = Swish() + + def forward(self, x: Tensor) -> Tensor: + """Compute convolution module. + + Args: + x: Input tensor (#time, batch, channels). + + Returns: + Tensor: Output tensor (#time, batch, channels). + + """ + # exchange the temporal dimension and the feature dimension + x = x.permute(1, 2, 0) # (#batch, channels, time). + + # GLU mechanism + x = self.pointwise_conv1(x) # (batch, 2*channels, time) + x = nn.functional.glu(x, dim=1) # (batch, channels, time) + + # 1D Depthwise Conv + x = self.depthwise_conv(x) + x = self.activation(self.norm(x)) + + x = self.pointwise_conv2(x) # (batch, channel, time) + + return x.permute(2, 0, 1) + + +class Swish(torch.nn.Module): + """Construct an Swish object.""" + + def forward(self, x: Tensor) -> Tensor: + """Return Swich activation function.""" + return x * torch.sigmoid(x) + + +def identity(x): + return x diff --git a/egs/librispeech/ASR/conformer_ctc/decode.py b/egs/librispeech/ASR/conformer_ctc/decode.py new file mode 100755 index 0000000000..d1cbc14de9 --- /dev/null +++ b/egs/librispeech/ASR/conformer_ctc/decode.py @@ -0,0 +1,474 @@ +#!/usr/bin/env python3 + +# Copyright 2021 Xiaomi Corporation (Author: Liyong Guo, Fangjun Kuang) + +# (still working in progress) + +import argparse +import logging +from collections import defaultdict +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +import k2 +import torch +import torch.nn as nn +from conformer import Conformer + +from icefall.checkpoint import average_checkpoints, load_checkpoint +from icefall.dataset.librispeech import LibriSpeechAsrDataModule +from icefall.decode import ( + get_lattice, + nbest_decoding, + one_best_decoding, + rescore_with_attention_decoder, + rescore_with_n_best_list, + rescore_with_whole_lattice, +) +from icefall.lexicon import Lexicon +from icefall.utils import ( + AttributeDict, + get_texts, + setup_logger, + store_transcripts, + write_error_stats, +) + + +def get_parser(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + + parser.add_argument( + "--epoch", + type=int, + default=9, + help="It specifies the checkpoint to use for decoding." + "Note: Epoch counts from 0.", + ) + parser.add_argument( + "--avg", + type=int, + default=1, + help="Number of checkpoints to average. Automatically select " + "consecutive checkpoints before the checkpoint specified by " + "'--epoch'. ", + ) + return parser + + +def get_params() -> AttributeDict: + params = AttributeDict( + { + "exp_dir": Path("conformer_ctc/exp"), + "lang_dir": Path("data/lang/bpe"), + "lm_dir": Path("data/lm"), + "feature_dim": 80, + "nhead": 8, + "attention_dim": 512, + "subsampling_factor": 4, + "num_decoder_layers": 6, + "vgg_frontend": False, + "is_espnet_structure": True, + "mmi_loss": False, + "use_feat_batchnorm": True, + "search_beam": 20, + "output_beam": 8, + "min_active_states": 30, + "max_active_states": 10000, + "use_double_scores": True, + # Possible values for method: + # - 1best + # - nbest + # - nbest-rescoring + # - whole-lattice-rescoring + # - attention-decoder + # "method": "whole-lattice-rescoring", + "method": "1best", + # num_paths is used when method is "nbest", "nbest-rescoring", + # and attention-decoder + "num_paths": 100, + } + ) + return params + + +def decode_one_batch( + params: AttributeDict, + model: nn.Module, + HLG: k2.Fsa, + batch: dict, + lexicon: Lexicon, + G: Optional[k2.Fsa] = None, +) -> Dict[str, List[List[int]]]: + """Decode one batch and return the result in a dict. The dict has the + following format: + + - key: It indicates the setting used for decoding. For example, + if no rescoring is used, the key is the string `no_rescore`. + If LM rescoring is used, the key is the string `lm_scale_xxx`, + where `xxx` is the value of `lm_scale`. An example key is + `lm_scale_0.7` + - value: It contains the decoding result. `len(value)` equals to + batch size. `value[i]` is the decoding result for the i-th + utterance in the given batch. + Args: + params: + It's the return value of :func:`get_params`. + + - params.method is "1best", it uses 1best decoding without LM rescoring. + - params.method is "nbest", it uses nbest decoding without LM rescoring. + - params.method is "nbest-rescoring", it uses nbest LM rescoring. + - params.method is "whole-lattice-rescoring", it uses whole lattice LM + rescoring. + + model: + The neural model. + HLG: + The decoding graph. + batch: + It is the return value from iterating + `lhotse.dataset.K2SpeechRecognitionDataset`. See its documentation + for the format of the `batch`. + lexicon: + It contains word symbol table. + G: + An LM. It is not None when params.method is "nbest-rescoring" + or "whole-lattice-rescoring". In general, the G in HLG + is a 3-gram LM, while this G is a 4-gram LM. + Returns: + Return the decoding result. See above description for the format of + the returned dict. + """ + device = HLG.device + feature = batch["inputs"] + assert feature.ndim == 3 + feature = feature.to(device) + # at entry, feature is [N, T, C] + + feature = feature.permute(0, 2, 1) # now feature is [N, C, T] + + supervisions = batch["supervisions"] + + nnet_output, memory, memory_key_padding_mask = model(feature, supervisions) + # nnet_output is [N, C, T] + + nnet_output = nnet_output.permute(0, 2, 1) + # now nnet_output is [N, T, C] + + supervision_segments = torch.stack( + ( + supervisions["sequence_idx"], + supervisions["start_frame"] // params.subsampling_factor, + supervisions["num_frames"] // params.subsampling_factor, + ), + 1, + ).to(torch.int32) + + lattice = get_lattice( + nnet_output=nnet_output, + HLG=HLG, + supervision_segments=supervision_segments, + search_beam=params.search_beam, + output_beam=params.output_beam, + min_active_states=params.min_active_states, + max_active_states=params.max_active_states, + subsampling_factor=params.subsampling_factor, + ) + + if params.method in ["1best", "nbest"]: + if params.method == "1best": + best_path = one_best_decoding( + lattice=lattice, use_double_scores=params.use_double_scores + ) + key = "no_rescore" + else: + best_path = nbest_decoding( + lattice=lattice, + num_paths=params.num_paths, + use_double_scores=params.use_double_scores, + ) + key = f"no_rescore-{params.num_paths}" + + hyps = get_texts(best_path) + hyps = [[lexicon.word_table[i] for i in ids] for ids in hyps] + return {key: hyps} + + assert params.method in [ + "nbest-rescoring", + "whole-lattice-rescoring", + "attention-decoder", + ] + + lm_scale_list = [0.8, 0.9, 1.0, 1.1, 1.2, 1.3] + lm_scale_list += [1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2.0] + + if params.method == "nbest-rescoring": + best_path_dict = rescore_with_n_best_list( + lattice=lattice, + G=G, + num_paths=params.num_paths, + lm_scale_list=lm_scale_list, + ) + elif params.method == "whole-lattice-rescoring": + best_path_dict = rescore_with_whole_lattice( + lattice=lattice, G_with_epsilon_loops=G, lm_scale_list=lm_scale_list + ) + elif params.method == "attention-decoder": + # lattice uses a 3-gram Lm. We rescore it with a 4-gram LM. + rescored_lattice = rescore_with_whole_lattice( + lattice=lattice, G_with_epsilon_loops=G, lm_scale_list=None + ) + + best_path_dict = rescore_with_attention_decoder( + lattice=rescored_lattice, + num_paths=params.num_paths, + model=model, + memory=memory, + memory_key_padding_mask=memory_key_padding_mask, + ) + else: + assert False, f"Unsupported decoding method: {params.method}" + + ans = dict() + for lm_scale_str, best_path in best_path_dict.items(): + hyps = get_texts(best_path) + hyps = [[lexicon.word_table[i] for i in ids] for ids in hyps] + ans[lm_scale_str] = hyps + return ans + + +def decode_dataset( + dl: torch.utils.data.DataLoader, + params: AttributeDict, + model: nn.Module, + HLG: k2.Fsa, + lexicon: Lexicon, + G: Optional[k2.Fsa] = None, +) -> Dict[str, List[Tuple[List[int], List[int]]]]: + """Decode dataset. + + Args: + dl: + PyTorch's dataloader containing the dataset to decode. + params: + It is returned by :func:`get_params`. + model: + The neural model. + HLG: + The decoding graph. + lexicon: + It contains word symbol table. + G: + An LM. It is not None when params.method is "nbest-rescoring" + or "whole-lattice-rescoring". In general, the G in HLG + is a 3-gram LM, while this G is a 4-gram LM. + Returns: + Return a dict, whose key may be "no-rescore" if no LM rescoring + is used, or it may be "lm_scale_0.7" if LM rescoring is used. + Its value is a list of tuples. Each tuple contains two elements: + The first is the reference transcript, and the second is the + predicted result. + """ + results = [] + + num_cuts = 0 + tot_num_cuts = len(dl.dataset.cuts) + + results = defaultdict(list) + for batch_idx, batch in enumerate(dl): + texts = batch["supervisions"]["text"] + + hyps_dict = decode_one_batch( + params=params, + model=model, + HLG=HLG, + batch=batch, + lexicon=lexicon, + G=G, + ) + + for lm_scale, hyps in hyps_dict.items(): + this_batch = [] + assert len(hyps) == len(texts) + for hyp_words, ref_text in zip(hyps, texts): + ref_words = ref_text.split() + this_batch.append((ref_words, hyp_words)) + + results[lm_scale].extend(this_batch) + + num_cuts += len(batch["supervisions"]["text"]) + + if batch_idx % 100 == 0: + logging.info( + f"batch {batch_idx}, cuts processed until now is " + f"{num_cuts}/{tot_num_cuts} " + f"({float(num_cuts)/tot_num_cuts*100:.6f}%)" + ) + return results + + +def save_results( + params: AttributeDict, + test_set_name: str, + results_dict: Dict[str, List[Tuple[List[int], List[int]]]], +): + test_set_wers = dict() + for key, results in results_dict.items(): + recog_path = params.exp_dir / f"recogs-{test_set_name}-{key}.txt" + store_transcripts(filename=recog_path, texts=results) + logging.info(f"The transcripts are stored in {recog_path}") + + # The following prints out WERs, per-word error statistics and aligned + # ref/hyp pairs. + errs_filename = params.exp_dir / f"errs-{test_set_name}-{key}.txt" + with open(errs_filename, "w") as f: + wer = write_error_stats(f, f"{test_set_name}-{key}", results) + test_set_wers[key] = wer + + logging.info("Wrote detailed error stats to {}".format(errs_filename)) + + test_set_wers = sorted(test_set_wers.items(), key=lambda x: x[1]) + errs_info = params.exp_dir / f"wer-summary-{test_set_name}.txt" + with open(errs_info, "w") as f: + print("settings\tWER", file=f) + for key, val in test_set_wers: + print("{}\t{}".format(key, val), file=f) + + s = "\nFor {}, WER of different settings are:\n".format(test_set_name) + note = "\tbest for {}".format(test_set_name) + for key, val in test_set_wers: + s += "{}\t{}{}\n".format(key, val, note) + note = "" + logging.info(s) + + +@torch.no_grad() +def main(): + parser = get_parser() + LibriSpeechAsrDataModule.add_arguments(parser) + args = parser.parse_args() + + params = get_params() + params.update(vars(args)) + + setup_logger(f"{params.exp_dir}/log/log-decode") + logging.info("Decoding started") + logging.info(params) + + lexicon = Lexicon(params.lang_dir) + max_token_id = max(lexicon.tokens) + num_classes = max_token_id + 1 # +1 for the blank + + device = torch.device("cpu") + if torch.cuda.is_available(): + device = torch.device("cuda", 0) + + logging.info(f"device: {device}") + + HLG = k2.Fsa.from_dict(torch.load(f"{params.lm_dir}/HLG_bpe.pt")) + HLG = HLG.to(device) + assert HLG.requires_grad is False + + if not hasattr(HLG, "lm_scores"): + HLG.lm_scores = HLG.scores.clone() + + # HLG = k2.ctc_topo(4999).to(device) + + if params.method in ( + "nbest-rescoring", + "whole-lattice-rescoring", + "attention-decoder", + ): + if not (params.lm_dir / "G_4_gram.pt").is_file(): + logging.info("Loading G_4_gram.fst.txt") + logging.warning("It may take 8 minutes.") + with open(params.lm_dir / "G_4_gram.fst.txt") as f: + first_word_disambig_id = lexicon.word_table["#0"] + + G = k2.Fsa.from_openfst(f.read(), acceptor=False) + # G.aux_labels is not needed in later computations, so + # remove it here. + del G.aux_labels + # CAUTION: The following line is crucial. + # Arcs entering the back-off state have label equal to #0. + # We have to change it to 0 here. + G.labels[G.labels >= first_word_disambig_id] = 0 + G = k2.Fsa.from_fsas([G]).to(device) + G = k2.arc_sort(G) + torch.save(G.as_dict(), params.lm_dir / "G_4_gram.pt") + else: + logging.info("Loading pre-compiled G_4_gram.pt") + d = torch.load(params.lm_dir / "G_4_gram.pt") + G = k2.Fsa.from_dict(d).to(device) + + if params.method in ["whole-lattice-rescoring", "attention-decoder"]: + # Add epsilon self-loops to G as we will compose + # it with the whole lattice later + G = k2.add_epsilon_self_loops(G) + G = k2.arc_sort(G) + G = G.to(device) + + # G.lm_scores is used to replace HLG.lm_scores during + # LM rescoring. + G.lm_scores = G.scores.clone() + else: + G = None + + model = Conformer( + num_features=params.feature_dim, + nhead=params.nhead, + d_model=params.attention_dim, + num_classes=num_classes, + subsampling_factor=params.subsampling_factor, + num_decoder_layers=params.num_decoder_layers, + vgg_frontend=params.vgg_frontend, + is_espnet_structure=params.is_espnet_structure, + mmi_loss=params.mmi_loss, + use_feat_batchnorm=params.use_feat_batchnorm, + ) + + if params.avg == 1: + load_checkpoint(f"{params.exp_dir}/epoch-{params.epoch}.pt", model) + else: + start = params.epoch - params.avg + 1 + filenames = [] + for i in range(start, params.epoch + 1): + if start >= 0: + filenames.append(f"{params.exp_dir}/epoch-{i}.pt") + logging.info(f"averaging {filenames}") + model.load_state_dict(average_checkpoints(filenames)) + + model.to(device) + model.eval() + num_param = sum([p.numel() for p in model.parameters()]) + logging.info(f"Number of model parameters: {num_param}") + + librispeech = LibriSpeechAsrDataModule(args) + # CAUTION: `test_sets` is for displaying only. + # If you want to skip test-clean, you have to skip + # it inside the for loop. That is, use + # + # if test_set == 'test-clean': continue + # + test_sets = ["test-clean", "test-other"] + for test_set, test_dl in zip(test_sets, librispeech.test_dataloaders()): + results_dict = decode_dataset( + dl=test_dl, + params=params, + model=model, + HLG=HLG, + lexicon=lexicon, + G=G, + ) + + save_results( + params=params, test_set_name=test_set, results_dict=results_dict + ) + + logging.info("Done!") + + +if __name__ == "__main__": + main() diff --git a/egs/librispeech/ASR/conformer_ctc/train.py b/egs/librispeech/ASR/conformer_ctc/train.py new file mode 100755 index 0000000000..40d3cf7fbb --- /dev/null +++ b/egs/librispeech/ASR/conformer_ctc/train.py @@ -0,0 +1,679 @@ +#!/usr/bin/env python3 + +# This is just at the very beginning ... + +import argparse +import logging +from pathlib import Path +from shutil import copyfile +from typing import Optional + +import k2 +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn as nn +from conformer import Conformer +from lhotse.utils import fix_random_seed +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.utils.tensorboard import SummaryWriter +from transformer import Noam + +from icefall.bpe_graph_compiler import BpeCtcTrainingGraphCompiler +from icefall.checkpoint import load_checkpoint +from icefall.checkpoint import save_checkpoint as save_checkpoint_impl +from icefall.dataset.librispeech import LibriSpeechAsrDataModule +from icefall.dist import cleanup_dist, setup_dist +from icefall.lexicon import Lexicon +from icefall.utils import ( + AttributeDict, + encode_supervisions, + setup_logger, + str2bool, +) + + +def get_parser(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + + parser.add_argument( + "--world-size", + type=int, + default=1, + help="Number of GPUs for DDP training.", + ) + + parser.add_argument( + "--master-port", + type=int, + default=12354, + help="Master port to use for DDP training.", + ) + + parser.add_argument( + "--tensorboard", + type=str2bool, + default=True, + help="Should various information be logged in tensorboard.", + ) + + # TODO: add extra arguments and support DDP training. + # Currently, only single GPU training is implemented. Will add + # DDP training once single GPU training is finished. + return parser + + +def get_params() -> AttributeDict: + """Return a dict containing training parameters. + + All training related parameters that are not passed from the commandline + is saved in the variable `params`. + + Commandline options are merged into `params` after they are parsed, so + you can also access them via `params`. + + Explanation of options saved in `params`: + + - exp_dir: It specifies the directory where all training related + files, e.g., checkpoints, log, etc, are saved + + - lang_dir: It contains language related input files such as + "lexicon.txt" + + - lr: It specifies the initial learning rate + + - feature_dim: The model input dim. It has to match the one used + in computing features. + + - weight_decay: The weight_decay for the optimizer. + + - subsampling_factor: The subsampling factor for the model. + + - start_epoch: If it is not zero, load checkpoint `start_epoch-1` + and continue training from that checkpoint. + + - num_epochs: Number of epochs to train. + + - best_train_loss: Best training loss so far. It is used to select + the model that has the lowest training loss. It is + updated during the training. + + - best_valid_loss: Best validation loss so far. It is used to select + the model that has the lowest validation loss. It is + updated during the training. + + - best_train_epoch: It is the epoch that has the best training loss. + + - best_valid_epoch: It is the epoch that has the best validation loss. + + - batch_idx_train: Used to writing statistics to tensorboard. It + contains number of batches trained so far across + epochs. + + - log_interval: Print training loss if batch_idx % log_interval` is 0 + + - valid_interval: Run validation if batch_idx % valid_interval` is 0 + + - beam_size: It is used in k2.ctc_loss + + - reduction: It is used in k2.ctc_loss + + - use_double_scores: It is used in k2.ctc_loss + """ + params = AttributeDict( + { + "exp_dir": Path("conformer_ctc/exp"), + "lang_dir": Path("data/lang/bpe"), + "feature_dim": 80, + "weight_decay": 0.0, + "subsampling_factor": 4, + "start_epoch": 0, + "num_epochs": 50, + "best_train_loss": float("inf"), + "best_valid_loss": float("inf"), + "best_train_epoch": -1, + "best_valid_epoch": -1, + "batch_idx_train": 0, + "log_interval": 10, + "valid_interval": 3000, + "beam_size": 10, + "reduction": "sum", + "use_double_scores": True, + # + "accum_grad": 1, + "att_rate": 0.7, + "attention_dim": 512, + "nhead": 8, + "num_decoder_layers": 6, + "is_espnet_structure": True, + "mmi_loss": False, + "use_feat_batchnorm": True, + "lr_factor": 5.0, + "warm_step": 80000, + } + ) + + return params + + +def load_checkpoint_if_available( + params: AttributeDict, + model: nn.Module, + optimizer: Optional[torch.optim.Optimizer] = None, + scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None, +) -> None: + """Load checkpoint from file. + + If params.start_epoch is positive, it will load the checkpoint from + `params.start_epoch - 1`. Otherwise, this function does nothing. + + Apart from loading state dict for `model`, `optimizer` and `scheduler`, + it also updates `best_train_epoch`, `best_train_loss`, `best_valid_epoch`, + and `best_valid_loss` in `params`. + + Args: + params: + The return value of :func:`get_params`. + model: + The training model. + optimizer: + The optimizer that we are using. + scheduler: + The learning rate scheduler we are using. + Returns: + Return None. + """ + if params.start_epoch <= 0: + return + + filename = params.exp_dir / f"epoch-{params.start_epoch-1}.pt" + saved_params = load_checkpoint( + filename, + model=model, + optimizer=optimizer, + scheduler=scheduler, + ) + + keys = [ + "best_train_epoch", + "best_valid_epoch", + "batch_idx_train", + "best_train_loss", + "best_valid_loss", + ] + for k in keys: + params[k] = saved_params[k] + + return saved_params + + +def save_checkpoint( + params: AttributeDict, + model: nn.Module, + optimizer: Optional[torch.optim.Optimizer] = None, + scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None, + rank: int = 0, +) -> None: + """Save model, optimizer, scheduler and training stats to file. + + Args: + params: + It is returned by :func:`get_params`. + model: + The training model. + """ + if rank != 0: + return + filename = params.exp_dir / f"epoch-{params.cur_epoch}.pt" + save_checkpoint_impl( + filename=filename, + model=model, + params=params, + optimizer=optimizer, + scheduler=scheduler, + rank=rank, + ) + + if params.best_train_epoch == params.cur_epoch: + best_train_filename = params.exp_dir / "best-train-loss.pt" + copyfile(src=filename, dst=best_train_filename) + + if params.best_valid_epoch == params.cur_epoch: + best_valid_filename = params.exp_dir / "best-valid-loss.pt" + copyfile(src=filename, dst=best_valid_filename) + + +def compute_loss( + params: AttributeDict, + model: nn.Module, + batch: dict, + graph_compiler: BpeCtcTrainingGraphCompiler, + is_training: bool, +): + """ + Compute CTC loss given the model and its inputs. + + Args: + params: + Parameters for training. See :func:`get_params`. + model: + The model for training. It is an instance of Conformer in our case. + batch: + A batch of data. See `lhotse.dataset.K2SpeechRecognitionDataset()` + for the content in it. + graph_compiler: + It is used to build a decoding graph from a ctc topo and training + transcript. The training transcript is contained in the given `batch`, + while the ctc topo is built when this compiler is instantiated. + is_training: + True for training. False for validation. When it is True, this + function enables autograd during computation; when it is False, it + disables autograd. + """ + device = graph_compiler.device + feature = batch["inputs"] + # at entry, feature is [N, T, C] + feature = feature.permute(0, 2, 1) # now feature is [N, C, T] + assert feature.ndim == 3 + feature = feature.to(device) + + supervisions = batch["supervisions"] + with torch.set_grad_enabled(is_training): + nnet_output, encoder_memory, memory_mask = model(feature, supervisions) + # nnet_output is [N, C, T] + nnet_output = nnet_output.permute(0, 2, 1) # [N, C, T] -> [N, T, C] + + # NOTE: We need `encode_supervisions` to sort sequences with + # different duration in decreasing order, required by + # `k2.intersect_dense` called in `k2.ctc_loss` + supervision_segments, texts = encode_supervisions( + supervisions, subsampling_factor=params.subsampling_factor + ) + + token_ids = graph_compiler.texts_to_ids(texts) + + decoding_graph = graph_compiler.compile(token_ids) + + dense_fsa_vec = k2.DenseFsaVec( + nnet_output, + supervision_segments, + allow_truncate=params.subsampling_factor - 1, + ) + + ctc_loss = k2.ctc_loss( + decoding_graph=decoding_graph, + dense_fsa_vec=dense_fsa_vec, + output_beam=params.beam_size, + reduction=params.reduction, + use_double_scores=params.use_double_scores, + ) + + if params.att_rate != 0.0: + with torch.set_grad_enabled(is_training): + if hasattr(model, "module"): + att_loss = model.module.decoder_forward( + encoder_memory, + memory_mask, + token_ids=token_ids, + sos_id=graph_compiler.sos_id, + eos_id=graph_compiler.eos_id, + ) + else: + att_loss = model.decoder_forward( + encoder_memory, + memory_mask, + token_ids=token_ids, + sos_id=graph_compiler.sos_id, + eos_id=graph_compiler.eos_id, + ) + loss = (1.0 - params.att_rate) * ctc_loss + params.att_rate * att_loss + else: + loss = ctc_loss + att_loss = torch.tensor([0]) + + # train_frames and valid_frames are used for printing. + if is_training: + params.train_frames = supervision_segments[:, 2].sum().item() + else: + params.valid_frames = supervision_segments[:, 2].sum().item() + + assert loss.requires_grad == is_training + + return loss, ctc_loss.detach(), att_loss.detach() + + +def compute_validation_loss( + params: AttributeDict, + model: nn.Module, + graph_compiler: BpeCtcTrainingGraphCompiler, + valid_dl: torch.utils.data.DataLoader, + world_size: int = 1, +) -> None: + """Run the validation process. The validation loss + is saved in `params.valid_loss`. + """ + model.eval() + + tot_loss = 0.0 + tot_ctc_loss = 0.0 + tot_att_loss = 0.0 + tot_frames = 0.0 + for batch_idx, batch in enumerate(valid_dl): + loss, ctc_loss, att_loss = compute_loss( + params=params, + model=model, + batch=batch, + graph_compiler=graph_compiler, + is_training=False, + ) + assert loss.requires_grad is False + assert ctc_loss.requires_grad is False + assert att_loss.requires_grad is False + + loss_cpu = loss.detach().cpu().item() + tot_loss += loss_cpu + + tot_ctc_loss += ctc_loss.detach().cpu().item() + tot_att_loss += att_loss.detach().cpu().item() + + tot_frames += params.valid_frames + + if world_size > 1: + s = torch.tensor( + [tot_loss, tot_ctc_loss, tot_att_loss, tot_frames], + device=loss.device, + ) + dist.all_reduce(s, op=dist.ReduceOp.SUM) + s = s.cpu().tolist() + tot_loss = s[0] + tot_ctc_loss = s[1] + tot_att_loss = s[2] + tot_frames = s[3] + + params.valid_loss = tot_loss / tot_frames + params.valid_ctc_loss = tot_ctc_loss / tot_frames + params.valid_att_loss = tot_att_loss / tot_frames + + if params.valid_loss < params.best_valid_loss: + params.best_valid_epoch = params.cur_epoch + params.best_valid_loss = params.valid_loss + + +def train_one_epoch( + params: AttributeDict, + model: nn.Module, + optimizer: torch.optim.Optimizer, + graph_compiler: BpeCtcTrainingGraphCompiler, + train_dl: torch.utils.data.DataLoader, + valid_dl: torch.utils.data.DataLoader, + tb_writer: Optional[SummaryWriter] = None, + world_size: int = 1, +) -> None: + """Train the model for one epoch. + + The training loss from the mean of all frames is saved in + `params.train_loss`. It runs the validation process every + `params.valid_interval` batches. + + Args: + params: + It is returned by :func:`get_params`. + model: + The model for training. + optimizer: + The optimizer we are using. + graph_compiler: + It is used to convert transcripts to FSAs. + train_dl: + Dataloader for the training dataset. + valid_dl: + Dataloader for the validation dataset. + tb_writer: + Writer to write log messages to tensorboard. + world_size: + Number of nodes in DDP training. If it is 1, DDP is disabled. + """ + model.train() + + tot_loss = 0.0 # sum of losses over all batches + tot_ctc_loss = 0.0 + tot_att_loss = 0.0 + + tot_frames = 0.0 # sum of frames over all batches + for batch_idx, batch in enumerate(train_dl): + params.batch_idx_train += 1 + batch_size = len(batch["supervisions"]["text"]) + + loss, ctc_loss, att_loss = compute_loss( + params=params, + model=model, + batch=batch, + graph_compiler=graph_compiler, + is_training=True, + ) + + # NOTE: We use reduction==sum and loss is computed over utterances + # in the batch and there is no normalization to it so far. + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + loss_cpu = loss.detach().cpu().item() + ctc_loss_cpu = ctc_loss.detach().cpu().item() + att_loss_cpu = att_loss.detach().cpu().item() + + tot_frames += params.train_frames + tot_loss += loss_cpu + tot_ctc_loss += ctc_loss_cpu + tot_att_loss += att_loss_cpu + + tot_avg_loss = tot_loss / tot_frames + tot_avg_ctc_loss = tot_ctc_loss / tot_frames + tot_avg_att_loss = tot_att_loss / tot_frames + + if batch_idx % params.log_interval == 0: + logging.info( + f"Epoch {params.cur_epoch}, batch {batch_idx}, " + f"batch avg ctc loss {ctc_loss_cpu/params.train_frames:.4f}, " + f"batch avg att loss {att_loss_cpu/params.train_frames:.4f}, " + f"batch avg loss {loss_cpu/params.train_frames:.4f}, " + f"total avg ctc loss: {tot_avg_ctc_loss:.4f}, " + f"total avg att loss: {tot_avg_att_loss:.4f}, " + f"total avg loss: {tot_avg_loss:.4f}, " + f"batch size: {batch_size}" + ) + + if tb_writer is not None: + tb_writer.add_scalar( + "train/current_ctc_loss", + ctc_loss_cpu / params.train_frames, + params.batch_idx_train, + ) + tb_writer.add_scalar( + "train/current_att_loss", + att_loss_cpu / params.train_frames, + params.batch_idx_train, + ) + tb_writer.add_scalar( + "train/current_loss", + loss_cpu / params.train_frames, + params.batch_idx_train, + ) + tb_writer.add_scalar( + "train/tot_avg_ctc_loss", + tot_avg_ctc_loss, + params.batch_idx_train, + ) + + tb_writer.add_scalar( + "train/tot_avg_att_loss", + tot_avg_att_loss, + params.batch_idx_train, + ) + tb_writer.add_scalar( + "train/tot_avg_loss", + tot_avg_loss, + params.batch_idx_train, + ) + + if batch_idx > 0 and batch_idx % params.valid_interval == 0: + compute_validation_loss( + params=params, + model=model, + graph_compiler=graph_compiler, + valid_dl=valid_dl, + world_size=world_size, + ) + model.train() + logging.info( + f"Epoch {params.cur_epoch}, " + f"valid ctc loss {params.valid_ctc_loss:.4f}," + f"valid att loss {params.valid_att_loss:.4f}," + f"valid loss {params.valid_loss:.4f}," + f" best valid loss: {params.best_valid_loss:.4f} " + f"best valid epoch: {params.best_valid_epoch}" + ) + + params.train_loss = tot_loss / tot_frames + + if params.train_loss < params.best_train_loss: + params.best_train_epoch = params.cur_epoch + params.best_train_loss = params.train_loss + + +def run(rank, world_size, args): + """ + Args: + rank: + It is a value between 0 and `world_size-1`, which is + passed automatically by `mp.spawn()` in :func:`main`. + The node with rank 0 is responsible for saving checkpoint. + world_size: + Number of GPUs for DDP training. + args: + The return value of get_parser().parse_args() + """ + params = get_params() + params.update(vars(args)) + + fix_random_seed(42) + if world_size > 1: + setup_dist(rank, world_size, params.master_port) + + setup_logger(f"{params.exp_dir}/log/log-train") + logging.info("Training started") + logging.info(params) + + if args.tensorboard and rank == 0: + tb_writer = SummaryWriter(log_dir=f"{params.exp_dir}/tensorboard") + else: + tb_writer = None + + lexicon = Lexicon(params.lang_dir) + max_token_id = max(lexicon.tokens) + num_classes = max_token_id + 1 # +1 for the blank + + device = torch.device("cpu") + if torch.cuda.is_available(): + device = torch.device("cuda", rank) + + graph_compiler = BpeCtcTrainingGraphCompiler( + params.lang_dir, + device=device, + sos_token="", + eos_token="", + ) + + logging.info("About to create model") + model = Conformer( + num_features=params.feature_dim, + nhead=params.nhead, + d_model=params.attention_dim, + num_classes=num_classes, + subsampling_factor=params.subsampling_factor, + num_decoder_layers=params.num_decoder_layers, + vgg_frontend=False, + is_espnet_structure=params.is_espnet_structure, + mmi_loss=params.mmi_loss, + use_feat_batchnorm=params.use_feat_batchnorm, + ) + + checkpoints = load_checkpoint_if_available(params=params, model=model) + + model.to(device) + if world_size > 1: + model = DDP(model, device_ids=[rank]) + + optimizer = Noam( + model.parameters(), + model_size=params.attention_dim, + factor=params.lr_factor, + warm_step=params.warm_step, + weight_decay=params.weight_decay, + ) + + if checkpoints: + optimizer.load_state_dict(checkpoints["optimizer"]) + + librispeech = LibriSpeechAsrDataModule(args) + train_dl = librispeech.train_dataloaders() + valid_dl = librispeech.valid_dataloaders() + + for epoch in range(params.start_epoch, params.num_epochs): + train_dl.sampler.set_epoch(epoch) + + cur_lr = optimizer._rate + if tb_writer is not None: + tb_writer.add_scalar( + "train/learning_rate", cur_lr, params.batch_idx_train + ) + tb_writer.add_scalar("train/epoch", epoch, params.batch_idx_train) + + if rank == 0: + logging.info("epoch {}, learning rate {}".format(epoch, cur_lr)) + + params.cur_epoch = epoch + + train_one_epoch( + params=params, + model=model, + optimizer=optimizer, + graph_compiler=graph_compiler, + train_dl=train_dl, + valid_dl=valid_dl, + tb_writer=tb_writer, + world_size=world_size, + ) + + save_checkpoint( + params=params, + model=model, + optimizer=optimizer, + rank=rank, + ) + + logging.info("Done!") + + if world_size > 1: + torch.distributed.barrier() + cleanup_dist() + + +def main(): + parser = get_parser() + LibriSpeechAsrDataModule.add_arguments(parser) + args = parser.parse_args() + + world_size = args.world_size + assert world_size >= 1 + if world_size > 1: + mp.spawn(run, args=(world_size, args), nprocs=world_size, join=True) + else: + run(rank=0, world_size=1, args=args) + + +if __name__ == "__main__": + main() diff --git a/egs/librispeech/ASR/conformer_ctc/transformer.py b/egs/librispeech/ASR/conformer_ctc/transformer.py new file mode 100644 index 0000000000..1df16e3467 --- /dev/null +++ b/egs/librispeech/ASR/conformer_ctc/transformer.py @@ -0,0 +1,1095 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2021 University of Chinese Academy of Sciences (author: Han Zhu) +# Apache 2.0 + +import math +from typing import Dict, List, Optional, Tuple + +import k2 +import torch +from torch import Tensor, nn + +from icefall.utils import get_texts + +# Note: TorchScript requires Dict/List/etc. to be fully typed. +Supervisions = Dict[str, Tensor] + + +class Transformer(nn.Module): + """ + Args: + num_features (int): Number of input features + num_classes (int): Number of output classes + subsampling_factor (int): subsampling factor of encoder (the convolution layers before transformers) + d_model (int): attention dimension + nhead (int): number of head + dim_feedforward (int): feedforward dimention + num_encoder_layers (int): number of encoder layers + num_decoder_layers (int): number of decoder layers + dropout (float): dropout rate + normalize_before (bool): whether to use layer_norm before the first block. + vgg_frontend (bool): whether to use vgg frontend. + """ + + def __init__( + self, + num_features: int, + num_classes: int, + subsampling_factor: int = 4, + d_model: int = 256, + nhead: int = 4, + dim_feedforward: int = 2048, + num_encoder_layers: int = 12, + num_decoder_layers: int = 6, + dropout: float = 0.1, + normalize_before: bool = True, + vgg_frontend: bool = False, + mmi_loss: bool = True, + use_feat_batchnorm: bool = False, + ) -> None: + super().__init__() + self.use_feat_batchnorm = use_feat_batchnorm + if use_feat_batchnorm: + self.feat_batchnorm = nn.BatchNorm1d(num_features) + + self.num_features = num_features + self.num_classes = num_classes + self.subsampling_factor = subsampling_factor + if subsampling_factor != 4: + raise NotImplementedError("Support only 'subsampling_factor=4'.") + + self.encoder_embed = ( + VggSubsampling(num_features, d_model) + if vgg_frontend + else Conv2dSubsampling(num_features, d_model) + ) + self.encoder_pos = PositionalEncoding(d_model, dropout) + + encoder_layer = TransformerEncoderLayer( + d_model, + nhead, + dim_feedforward, + dropout, + normalize_before=normalize_before, + ) + + if normalize_before: + encoder_norm = nn.LayerNorm(d_model) + else: + encoder_norm = None + + self.encoder = nn.TransformerEncoder( + encoder_layer, num_encoder_layers, encoder_norm + ) + + self.encoder_output_layer = nn.Sequential( + nn.Dropout(p=dropout), nn.Linear(d_model, num_classes) + ) + + if num_decoder_layers > 0: + if mmi_loss: + self.decoder_num_class = ( + self.num_classes + 1 + ) # +1 for the sos/eos symbol + else: + self.decoder_num_class = ( + self.num_classes + ) # bpe model already has sos/eos symbol + + self.decoder_embed = nn.Embedding(self.decoder_num_class, d_model) + self.decoder_pos = PositionalEncoding(d_model, dropout) + + decoder_layer = TransformerDecoderLayer( + d_model, + nhead, + dim_feedforward, + dropout, + normalize_before=normalize_before, + ) + + if normalize_before: + decoder_norm = nn.LayerNorm(d_model) + else: + decoder_norm = None + + self.decoder = nn.TransformerDecoder( + decoder_layer, num_decoder_layers, decoder_norm + ) + + self.decoder_output_layer = torch.nn.Linear( + d_model, self.decoder_num_class + ) + + self.decoder_criterion = LabelSmoothingLoss(self.decoder_num_class) + else: + self.decoder_criterion = None + + def forward( + self, x: Tensor, supervision: Optional[Supervisions] = None + ) -> Tuple[Tensor, Tensor, Optional[Tensor]]: + """ + Args: + x: Tensor of dimension (batch_size, num_features, input_length). + supervision: Supervison in lhotse format, get from batch['supervisions'] + + Returns: + Tensor: After log-softmax tensor of dimension (batch_size, number_of_classes, input_length). + Tensor: Before linear layer tensor of dimension (input_length, batch_size, d_model). + Optional[Tensor]: Mask tensor of dimension (batch_size, input_length) or None. + + """ + if self.use_feat_batchnorm: + x = self.feat_batchnorm(x) + encoder_memory, memory_mask = self.encode(x, supervision) + x = self.encoder_output(encoder_memory) + return x, encoder_memory, memory_mask + + def encode( + self, x: Tensor, supervisions: Optional[Supervisions] = None + ) -> Tuple[Tensor, Optional[Tensor]]: + """ + Args: + x: Tensor of dimension (batch_size, num_features, input_length). + supervisions : Supervison in lhotse format, i.e., batch['supervisions'] + + Returns: + Tensor: Predictor tensor of dimension (input_length, batch_size, d_model). + Optional[Tensor]: Mask tensor of dimension (batch_size, input_length) or None. + """ + x = x.permute(0, 2, 1) # (B, F, T) -> (B, T, F) + + x = self.encoder_embed(x) + x = self.encoder_pos(x) + x = x.permute(1, 0, 2) # (B, T, F) -> (T, B, F) + mask = encoder_padding_mask(x.size(0), supervisions) + mask = mask.to(x.device) if mask != None else None + x = self.encoder(x, src_key_padding_mask=mask) # (T, B, F) + + return x, mask + + def encoder_output(self, x: Tensor) -> Tensor: + """ + Args: + x: Tensor of dimension (input_length, batch_size, d_model). + + Returns: + Tensor: After log-softmax tensor of dimension (batch_size, number_of_classes, input_length). + """ + x = self.encoder_output_layer(x).permute( + 1, 2, 0 + ) # (T, B, F) ->(B, F, T) + x = nn.functional.log_softmax(x, dim=1) # (B, F, T) + return x + + def decoder_forward( + self, + x: Tensor, + encoder_mask: Tensor, + supervision: Supervisions = None, + graph_compiler: object = None, + token_ids: List[int] = None, + sos_id: Optional[int] = None, + eos_id: Optional[int] = None, + ) -> Tensor: + """ + Args: + x: Tensor of dimension (input_length, batch_size, d_model). + encoder_mask: Mask tensor of dimension (batch_size, input_length) + supervision: Supervison in lhotse format, get from batch['supervisions'] + graph_compiler: use graph_compiler.L_inv (Its labels are words, while its aux_labels are phones) + , graph_compiler.words and graph_compiler.oov + sos_id: sos token id + eos_id: eos token id + + Returns: + Tensor: Decoder loss. + """ + if supervision is not None and graph_compiler is not None: + batch_text = get_normal_transcripts( + supervision, graph_compiler.lexicon.words, graph_compiler.oov + ) + ys_in_pad, ys_out_pad = add_sos_eos( + batch_text, graph_compiler.L_inv, sos_id, eos_id, + ) + elif token_ids is not None: + _sos = torch.tensor([sos_id]) + _eos = torch.tensor([eos_id]) + ys_in = [ + torch.cat([_sos, torch.tensor(y)], dim=0) for y in token_ids + ] + ys_out = [ + torch.cat([torch.tensor(y), _eos], dim=0) for y in token_ids + ] + ys_in_pad = pad_list(ys_in, eos_id) + ys_out_pad = pad_list(ys_out, -1) + + else: + raise ValueError("Invalid input for decoder self attetion") + + ys_in_pad = ys_in_pad.to(x.device) + ys_out_pad = ys_out_pad.to(x.device) + + tgt_mask = generate_square_subsequent_mask(ys_in_pad.shape[-1]).to( + x.device + ) + + tgt_key_padding_mask = decoder_padding_mask(ys_in_pad) + + tgt = self.decoder_embed(ys_in_pad) # (B, T) -> (B, T, F) + tgt = self.decoder_pos(tgt) + tgt = tgt.permute(1, 0, 2) # (B, T, F) -> (T, B, F) + pred_pad = self.decoder( + tgt=tgt, + memory=x, + tgt_mask=tgt_mask, + tgt_key_padding_mask=tgt_key_padding_mask, + memory_key_padding_mask=encoder_mask, + ) # (T, B, F) + pred_pad = pred_pad.permute(1, 0, 2) # (T, B, F) -> (B, T, F) + pred_pad = self.decoder_output_layer(pred_pad) # (B, T, F) + + decoder_loss = self.decoder_criterion(pred_pad, ys_out_pad) + + return decoder_loss + + def decoder_nll( + self, + x: Tensor, + encoder_mask: Tensor, + token_ids: List[List[int]], + sos_id: int, + eos_id: int, + ) -> Tensor: + """ + Args: + x: encoder-output, Tensor of dimension (input_length, batch_size, d_model). + encoder_mask: Mask tensor of dimension (batch_size, input_length) + token_ids: n-best list extracted from lattice before rescore + + Returns: + Tensor: negative log-likelihood. + """ + # The common part between this fuction and decoder_forward could be + # extracted as a seperated function. + if token_ids is not None: + _sos = torch.tensor([sos_id]) + _eos = torch.tensor([eos_id]) + ys_in = [ + torch.cat([_sos, torch.tensor(y)], dim=0) for y in token_ids + ] + ys_out = [ + torch.cat([torch.tensor(y), _eos], dim=0) for y in token_ids + ] + ys_in_pad = pad_list(ys_in, eos_id) + ys_out_pad = pad_list(ys_out, -1) + else: + raise ValueError("Invalid input for decoder self attetion") + + ys_in_pad = ys_in_pad.to(x.device, dtype=torch.int64) + ys_out_pad = ys_out_pad.to(x.device, dtype=torch.int64) + + tgt_mask = generate_square_subsequent_mask(ys_in_pad.shape[-1]).to( + x.device + ) + + tgt_key_padding_mask = decoder_padding_mask(ys_in_pad) + + tgt = self.decoder_embed(ys_in_pad) # (B, T) -> (B, T, F) + tgt = self.decoder_pos(tgt) + tgt = tgt.permute(1, 0, 2) # (B, T, F) -> (T, B, F) + pred_pad = self.decoder( + tgt=tgt, + memory=x, + tgt_mask=tgt_mask, + tgt_key_padding_mask=tgt_key_padding_mask, + memory_key_padding_mask=encoder_mask, + ) # (T, B, F) + pred_pad = pred_pad.permute(1, 0, 2) # (T, B, F) -> (B, T, F) + pred_pad = self.decoder_output_layer(pred_pad) # (B, T, F) + # nll: negative log-likelihood + nll = torch.nn.functional.cross_entropy( + pred_pad.view(-1, self.decoder_num_class), + ys_out_pad.view(-1), + ignore_index=-1, + reduction="none", + ) + + nll = nll.view(pred_pad.shape[0], -1) + + return nll + + +class TransformerEncoderLayer(nn.Module): + """ + Modified from torch.nn.TransformerEncoderLayer. Add support of normalize_before, + i.e., use layer_norm before the first block. + + Args: + d_model: the number of expected features in the input (required). + nhead: the number of heads in the multiheadattention models (required). + dim_feedforward: the dimension of the feedforward network model (default=2048). + dropout: the dropout value (default=0.1). + activation: the activation function of intermediate layer, relu or gelu (default=relu). + normalize_before: whether to use layer_norm before the first block. + + Examples:: + >>> encoder_layer = TransformerEncoderLayer(d_model=512, nhead=8) + >>> src = torch.rand(10, 32, 512) + >>> out = encoder_layer(src) + """ + + def __init__( + self, + d_model: int, + nhead: int, + dim_feedforward: int = 2048, + dropout: float = 0.1, + activation: str = "relu", + normalize_before: bool = True, + ) -> None: + super(TransformerEncoderLayer, self).__init__() + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=0.0) + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + + self.normalize_before = normalize_before + + def __setstate__(self, state): + if "activation" not in state: + state["activation"] = nn.functional.relu + super(TransformerEncoderLayer, self).__setstate__(state) + + def forward( + self, + src: Tensor, + src_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + ) -> Tensor: + """ + Pass the input through the encoder layer. + + Args: + src: the sequence to the encoder layer (required). + src_mask: the mask for the src sequence (optional). + src_key_padding_mask: the mask for the src keys per batch (optional). + + Shape: + src: (S, N, E). + src_mask: (S, S). + src_key_padding_mask: (N, S). + S is the source sequence length, T is the target sequence length, N is the batch size, E is the feature number + """ + residual = src + if self.normalize_before: + src = self.norm1(src) + src2 = self.self_attn( + src, + src, + src, + attn_mask=src_mask, + key_padding_mask=src_key_padding_mask, + )[0] + src = residual + self.dropout1(src2) + if not self.normalize_before: + src = self.norm1(src) + + residual = src + if self.normalize_before: + src = self.norm2(src) + src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) + src = residual + self.dropout2(src2) + if not self.normalize_before: + src = self.norm2(src) + return src + + +class TransformerDecoderLayer(nn.Module): + """ + Modified from torch.nn.TransformerDecoderLayer. Add support of normalize_before, + i.e., use layer_norm before the first block. + + Args: + d_model: the number of expected features in the input (required). + nhead: the number of heads in the multiheadattention models (required). + dim_feedforward: the dimension of the feedforward network model (default=2048). + dropout: the dropout value (default=0.1). + activation: the activation function of intermediate layer, relu or gelu (default=relu). + + Examples:: + >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8) + >>> memory = torch.rand(10, 32, 512) + >>> tgt = torch.rand(20, 32, 512) + >>> out = decoder_layer(tgt, memory) + """ + + def __init__( + self, + d_model: int, + nhead: int, + dim_feedforward: int = 2048, + dropout: float = 0.1, + activation: str = "relu", + normalize_before: bool = True, + ) -> None: + super(TransformerDecoderLayer, self).__init__() + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=0.0) + self.src_attn = nn.MultiheadAttention(d_model, nhead, dropout=0.0) + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.norm3 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + self.dropout3 = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + + self.normalize_before = normalize_before + + def __setstate__(self, state): + if "activation" not in state: + state["activation"] = nn.functional.relu + super(TransformerDecoderLayer, self).__setstate__(state) + + def forward( + self, + tgt: Tensor, + memory: Tensor, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + ) -> Tensor: + """Pass the inputs (and mask) through the decoder layer. + + Args: + tgt: the sequence to the decoder layer (required). + memory: the sequence from the last layer of the encoder (required). + tgt_mask: the mask for the tgt sequence (optional). + memory_mask: the mask for the memory sequence (optional). + tgt_key_padding_mask: the mask for the tgt keys per batch (optional). + memory_key_padding_mask: the mask for the memory keys per batch (optional). + + Shape: + tgt: (T, N, E). + memory: (S, N, E). + tgt_mask: (T, T). + memory_mask: (T, S). + tgt_key_padding_mask: (N, T). + memory_key_padding_mask: (N, S). + S is the source sequence length, T is the target sequence length, N is the batch size, E is the feature number + """ + residual = tgt + if self.normalize_before: + tgt = self.norm1(tgt) + tgt2 = self.self_attn( + tgt, + tgt, + tgt, + attn_mask=tgt_mask, + key_padding_mask=tgt_key_padding_mask, + )[0] + tgt = residual + self.dropout1(tgt2) + if not self.normalize_before: + tgt = self.norm1(tgt) + + residual = tgt + if self.normalize_before: + tgt = self.norm2(tgt) + tgt2 = self.src_attn( + tgt, + memory, + memory, + attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask, + )[0] + tgt = residual + self.dropout2(tgt2) + if not self.normalize_before: + tgt = self.norm2(tgt) + + residual = tgt + if self.normalize_before: + tgt = self.norm3(tgt) + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) + tgt = residual + self.dropout3(tgt2) + if not self.normalize_before: + tgt = self.norm3(tgt) + return tgt + + +def _get_activation_fn(activation: str): + if activation == "relu": + return nn.functional.relu + elif activation == "gelu": + return nn.functional.gelu + + raise RuntimeError( + "activation should be relu/gelu, not {}".format(activation) + ) + + +class Conv2dSubsampling(nn.Module): + """Convolutional 2D subsampling (to 1/4 length). + Modified from https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/transformer/subsampling.py + + Args: + idim: Input dimension. + odim: Output dimension. + + """ + + def __init__(self, idim: int, odim: int) -> None: + """Construct a Conv2dSubsampling object.""" + super(Conv2dSubsampling, self).__init__() + self.conv = nn.Sequential( + nn.Conv2d( + in_channels=1, out_channels=odim, kernel_size=3, stride=2 + ), + nn.ReLU(), + nn.Conv2d( + in_channels=odim, out_channels=odim, kernel_size=3, stride=2 + ), + nn.ReLU(), + ) + self.out = nn.Linear(odim * (((idim - 1) // 2 - 1) // 2), odim) + + def forward(self, x: Tensor) -> Tensor: + """Subsample x. + + Args: + x: Input tensor of dimension (batch_size, input_length, num_features). (#batch, time, idim). + + Returns: + torch.Tensor: Subsampled tensor of dimension (batch_size, input_length, d_model). + where time' = time // 4. + + """ + x = x.unsqueeze(1) # (b, c, t, f) + x = self.conv(x) + b, c, t, f = x.size() + x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f)) + return x + + +class VggSubsampling(nn.Module): + """Trying to follow the setup described here https://arxiv.org/pdf/1910.09799.pdf + This paper is not 100% explicit so I am guessing to some extent, + and trying to compare with other VGG implementations. + + Args: + idim: Input dimension. + odim: Output dimension. + + """ + + def __init__(self, idim: int, odim: int) -> None: + """Construct a VggSubsampling object. This uses 2 VGG blocks with 2 + Conv2d layers each, subsampling its input by a factor of 4 in the + time dimensions. + + Args: + idim: Number of features at input, e.g. 40 or 80 for MFCC + (will be treated as the image height). + odim: Output dimension (number of features), e.g. 256 + """ + super(VggSubsampling, self).__init__() + + cur_channels = 1 + layers = [] + block_dims = [32, 64] + + # The decision to use padding=1 for the 1st convolution, then padding=0 + # for the 2nd and for the max-pooling, and ceil_mode=True, was driven by + # a back-compatibility concern so that the number of frames at the + # output would be equal to: + # (((T-1)//2)-1)//2. + # We can consider changing this by using padding=1 on the 2nd convolution, + # so the num-frames at the output would be T//4. + for block_dim in block_dims: + layers.append( + torch.nn.Conv2d( + in_channels=cur_channels, + out_channels=block_dim, + kernel_size=3, + padding=1, + stride=1, + ) + ) + layers.append(torch.nn.ReLU()) + layers.append( + torch.nn.Conv2d( + in_channels=block_dim, + out_channels=block_dim, + kernel_size=3, + padding=0, + stride=1, + ) + ) + layers.append( + torch.nn.MaxPool2d( + kernel_size=2, stride=2, padding=0, ceil_mode=True + ) + ) + cur_channels = block_dim + + self.layers = nn.Sequential(*layers) + + self.out = nn.Linear( + block_dims[-1] * (((idim - 1) // 2 - 1) // 2), odim + ) + + def forward(self, x: Tensor) -> Tensor: + """Subsample x. + + Args: + x: Input tensor of dimension (batch_size, input_length, num_features). (#batch, time, idim). + + Returns: + torch.Tensor: Subsampled tensor of dimension (batch_size, input_length', d_model). + where input_length' == (((input_length - 1) // 2) - 1) // 2 + + """ + x = x.unsqueeze(1) # (b, c, t, f) + x = self.layers(x) + b, c, t, f = x.size() + x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f)) + return x + + +class PositionalEncoding(nn.Module): + """ + Positional encoding. + + Args: + d_model: Embedding dimension. + dropout: Dropout rate. + max_len: Maximum input length. + + """ + + def __init__( + self, d_model: int, dropout: float = 0.1, max_len: int = 5000 + ) -> None: + """Construct an PositionalEncoding object.""" + super(PositionalEncoding, self).__init__() + self.d_model = d_model + self.xscale = math.sqrt(self.d_model) + self.dropout = nn.Dropout(p=dropout) + self.pe = None + self.extend_pe(torch.tensor(0.0).expand(1, max_len)) + + def extend_pe(self, x: Tensor) -> None: + """Reset the positional encodings.""" + if self.pe is not None: + if self.pe.size(1) >= x.size(1): + if self.pe.dtype != x.dtype or self.pe.device != x.device: + self.pe = self.pe.to(dtype=x.dtype, device=x.device) + return + pe = torch.zeros(x.size(1), self.d_model) + position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, self.d_model, 2, dtype=torch.float32) + * -(math.log(10000.0) / self.d_model) + ) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + pe = pe.unsqueeze(0) + self.pe = pe.to(device=x.device, dtype=x.dtype) + + def forward(self, x: Tensor) -> Tensor: + """ + Add positional encoding. + + Args: + x: Input tensor of dimention (batch_size, input_length, d_model). + + Returns: + torch.Tensor: Encoded tensor of dimention (batch_size, input_length, d_model). + + """ + self.extend_pe(x) + x = x * self.xscale + self.pe[:, : x.size(1)] + return self.dropout(x) + + +class Noam(object): + """ + Implements Noam optimizer. Proposed in "Attention Is All You Need", https://arxiv.org/pdf/1706.03762.pdf + Modified from https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/transformer/optimizer.py + + Args: + params (iterable): iterable of parameters to optimize or dicts defining parameter groups + model_size: attention dimension of the transformer model + factor: learning rate factor + warm_step: warmup steps + """ + + def __init__( + self, + params, + model_size: int = 256, + factor: float = 10.0, + warm_step: int = 25000, + weight_decay=0, + ) -> None: + """Construct an Noam object.""" + self.optimizer = torch.optim.Adam( + params, lr=0, betas=(0.9, 0.98), eps=1e-9, weight_decay=weight_decay + ) + self._step = 0 + self.warmup = warm_step + self.factor = factor + self.model_size = model_size + self._rate = 0 + + @property + def param_groups(self): + """Return param_groups.""" + return self.optimizer.param_groups + + def step(self): + """Update parameters and rate.""" + self._step += 1 + rate = self.rate() + for p in self.optimizer.param_groups: + p["lr"] = rate + self._rate = rate + self.optimizer.step() + + def rate(self, step=None): + """Implement `lrate` above.""" + if step is None: + step = self._step + return ( + self.factor + * self.model_size ** (-0.5) + * min(step ** (-0.5), step * self.warmup ** (-1.5)) + ) + + def zero_grad(self): + """Reset gradient.""" + self.optimizer.zero_grad() + + def state_dict(self): + """Return state_dict.""" + return { + "_step": self._step, + "warmup": self.warmup, + "factor": self.factor, + "model_size": self.model_size, + "_rate": self._rate, + "optimizer": self.optimizer.state_dict(), + } + + def load_state_dict(self, state_dict): + """Load state_dict.""" + for key, value in state_dict.items(): + if key == "optimizer": + self.optimizer.load_state_dict(state_dict["optimizer"]) + else: + setattr(self, key, value) + + +class LabelSmoothingLoss(nn.Module): + """ + Label-smoothing loss. KL-divergence between q_{smoothed ground truth prob.}(w) + and p_{prob. computed by model}(w) is minimized. + Modified from https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/transformer/label_smoothing_loss.py + + Args: + size: the number of class + padding_idx: padding_idx: ignored class id + smoothing: smoothing rate (0.0 means the conventional CE) + normalize_length: normalize loss by sequence length if True + criterion: loss function to be smoothed + """ + + def __init__( + self, + size: int, + padding_idx: int = -1, + smoothing: float = 0.1, + normalize_length: bool = False, + criterion: nn.Module = nn.KLDivLoss(reduction="none"), + ) -> None: + """Construct an LabelSmoothingLoss object.""" + super(LabelSmoothingLoss, self).__init__() + self.criterion = criterion + self.padding_idx = padding_idx + assert 0.0 < smoothing <= 1.0 + self.confidence = 1.0 - smoothing + self.smoothing = smoothing + self.size = size + self.true_dist = None + self.normalize_length = normalize_length + + def forward(self, x: Tensor, target: Tensor) -> Tensor: + """ + Compute loss between x and target. + + Args: + x: prediction of dimention (batch_size, input_length, number_of_classes). + target: target masked with self.padding_id of dimention (batch_size, input_length). + + Returns: + torch.Tensor: scalar float value + """ + assert x.size(2) == self.size + batch_size = x.size(0) + x = x.view(-1, self.size) + target = target.view(-1) + with torch.no_grad(): + true_dist = x.clone() + true_dist.fill_(self.smoothing / (self.size - 1)) + ignore = target == self.padding_idx # (B,) + total = len(target) - ignore.sum().item() + target = target.masked_fill(ignore, 0) # avoid -1 index + true_dist.scatter_(1, target.unsqueeze(1), self.confidence) + kl = self.criterion(torch.log_softmax(x, dim=1), true_dist) + # denom = total if self.normalize_length else batch_size + denom = total if self.normalize_length else 1 + return kl.masked_fill(ignore.unsqueeze(1), 0).sum() / denom + + +def encoder_padding_mask( + max_len: int, supervisions: Optional[Supervisions] = None +) -> Optional[Tensor]: + """Make mask tensor containing indices of padded part. + + Args: + max_len: maximum length of input features + supervisions : Supervison in lhotse format, i.e., batch['supervisions'] + + Returns: + Tensor: Mask tensor of dimension (batch_size, input_length), True denote the masked indices. + """ + if supervisions is None: + return None + + supervision_segments = torch.stack( + ( + supervisions["sequence_idx"], + supervisions["start_frame"], + supervisions["num_frames"], + ), + 1, + ).to(torch.int32) + + lengths = [ + 0 for _ in range(int(supervision_segments[:, 0].max().item()) + 1) + ] + for idx in range(supervision_segments.size(0)): + # Note: TorchScript doesn't allow to unpack tensors as tuples + sequence_idx = supervision_segments[idx, 0].item() + start_frame = supervision_segments[idx, 1].item() + num_frames = supervision_segments[idx, 2].item() + lengths[sequence_idx] = start_frame + num_frames + + lengths = [((i - 1) // 2 - 1) // 2 for i in lengths] + bs = int(len(lengths)) + seq_range = torch.arange(0, max_len, dtype=torch.int64) + seq_range_expand = seq_range.unsqueeze(0).expand(bs, max_len) + # Note: TorchScript doesn't implement Tensor.new() + seq_length_expand = torch.tensor( + lengths, device=seq_range_expand.device, dtype=seq_range_expand.dtype + ).unsqueeze(-1) + mask = seq_range_expand >= seq_length_expand + + return mask + + +def decoder_padding_mask(ys_pad: Tensor, ignore_id: int = -1) -> Tensor: + """Generate a length mask for input. The masked position are filled with bool(True), + Unmasked positions are filled with bool(False). + + Args: + ys_pad: padded tensor of dimension (batch_size, input_length). + ignore_id: the ignored number (the padding number) in ys_pad + + Returns: + Tensor: a mask tensor of dimension (batch_size, input_length). + """ + ys_mask = ys_pad == ignore_id + return ys_mask + + +def get_normal_transcripts( + supervision: Supervisions, words: k2.SymbolTable, oov: str = "" +) -> List[List[int]]: + """Get normal transcripts (1 input recording has 1 transcript) from lhotse cut format. + Achieved by concatenate the transcripts corresponding to the same recording. + + Args: + supervision : Supervison in lhotse format, i.e., batch['supervisions'] + words: The word symbol table. + oov: Out of vocabulary word. + + Returns: + List[List[int]]: List of concatenated transcripts, length is batch_size + """ + + texts = [ + [token if token in words else oov for token in text.split(" ")] + for text in supervision["text"] + ] + texts_ids = [[words[token] for token in text] for text in texts] + + batch_text = [ + [] for _ in range(int(supervision["sequence_idx"].max().item()) + 1) + ] + for sequence_idx, text in zip(supervision["sequence_idx"], texts_ids): + batch_text[sequence_idx] = batch_text[sequence_idx] + text + return batch_text + + +def generate_square_subsequent_mask(sz: int) -> Tensor: + """Generate a square mask for the sequence. The masked positions are filled with float('-inf'). + Unmasked positions are filled with float(0.0). + + Args: + sz: mask size + + Returns: + Tensor: a square mask of dimension (sz, sz) + """ + mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1) + mask = ( + mask.float() + .masked_fill(mask == 0, float("-inf")) + .masked_fill(mask == 1, float(0.0)) + ) + return mask + + +def add_sos_eos( + ys: List[List[int]], + lexicon: k2.Fsa, + sos_id: int, + eos_id: int, + ignore_id: int = -1, +) -> Tuple[Tensor, Tensor]: + """Add and labels. + + Args: + ys: batch of unpadded target sequences + lexicon: Its labels are words, while its aux_labels are phones. + sos_id: index of + eos_id: index of + ignore_id: index of padding + + Returns: + Tensor: Input of transformer decoder. Padded tensor of dimention (batch_size, max_length). + Tensor: Output of transformer decoder. padded tensor of dimention (batch_size, max_length). + """ + + _sos = torch.tensor([sos_id]) + _eos = torch.tensor([eos_id]) + ys = get_hierarchical_targets(ys, lexicon) + ys_in = [torch.cat([_sos, y], dim=0) for y in ys] + ys_out = [torch.cat([y, _eos], dim=0) for y in ys] + return pad_list(ys_in, eos), pad_list(ys_out, ignore_id) + + +def pad_list(ys: List[Tensor], pad_value: float) -> Tensor: + """Perform padding for the list of tensors. + + Args: + ys: List of tensors. len(ys) = batch_size. + pad_value: Value for padding. + + Returns: + Tensor: Padded tensor (batch_size, max_length, `*`). + + Examples: + >>> x = [torch.ones(4), torch.ones(2), torch.ones(1)] + >>> x + [tensor([1., 1., 1., 1.]), tensor([1., 1.]), tensor([1.])] + >>> pad_list(x, 0) + tensor([[1., 1., 1., 1.], + [1., 1., 0., 0.], + [1., 0., 0., 0.]]) + + """ + n_batch = len(ys) + max_len = max(x.size(0) for x in ys) + pad = ys[0].new_full((n_batch, max_len, *ys[0].size()[1:]), pad_value) + + for i in range(n_batch): + pad[i, : ys[i].size(0)] = ys[i] + + return pad + + +def get_hierarchical_targets( + ys: List[List[int]], lexicon: k2.Fsa +) -> List[Tensor]: + """Get hierarchical transcripts (i.e., phone level transcripts) from transcripts (i.e., word level transcripts). + + Args: + ys: Word level transcripts. + lexicon: Its labels are words, while its aux_labels are phones. + + Returns: + List[Tensor]: Phone level transcripts. + + """ + + if lexicon is None: + return ys + else: + L_inv = lexicon + + n_batch = len(ys) + device = L_inv.device + + transcripts = k2.create_fsa_vec( + [k2.linear_fsa(x, device=device) for x in ys] + ) + transcripts_with_self_loops = k2.add_epsilon_self_loops(transcripts) + + transcripts_lexicon = k2.intersect( + L_inv, transcripts_with_self_loops, treat_epsilons_specially=False + ) + # Don't call invert_() above because we want to return phone IDs, + # which is the `aux_labels` of transcripts_lexicon + transcripts_lexicon = k2.remove_epsilon(transcripts_lexicon) + transcripts_lexicon = k2.top_sort(transcripts_lexicon) + + transcripts_lexicon = k2.shortest_path( + transcripts_lexicon, use_double_scores=True + ) + + ys = get_texts(transcripts_lexicon) + ys = [torch.tensor(y) for y in ys] + + return ys + + +def test_transformer(): + t = Transformer(40, 1281) + T = 200 + f = torch.rand(31, 40, T) + g, _, _ = t(f) + assert g.shape == (31, 1281, (((T - 1) // 2) - 1) // 2) + + +def main(): + test_transformer() + + +if __name__ == "__main__": + main() diff --git a/egs/librispeech/ASR/local/__init__.py b/egs/librispeech/ASR/local/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/egs/librispeech/ASR/local/compile_hlg.py b/egs/librispeech/ASR/local/compile_hlg.py new file mode 100755 index 0000000000..605d72daed --- /dev/null +++ b/egs/librispeech/ASR/local/compile_hlg.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python3 + +""" +This script compiles HLG from + + - H, the ctc topology, built from tokens contained in lexicon.txt + - L, the lexicon, built from L_disambig.pt + + Caution: We use a lexicon that contains disambiguation symbols + + - G, the LM, built from data/lm/G_3_gram.fst.txt + +The generated HLG is saved in data/lm/HLG.pt (phone based) +or data/lm/HLG_bpe.pt (BPE based) +""" +import logging +from pathlib import Path + +import k2 +import torch + +from icefall.lexicon import Lexicon + + +def compile_HLG(lang_dir: str) -> k2.Fsa: + """ + Args: + lang_dir: + The language directory, e.g., data/lang or data/lang/bpe. + + Return: + An FSA representing HLG. + """ + lexicon = Lexicon(lang_dir) + max_token_id = max(lexicon.tokens) + logging.info(f"Building ctc_topo. max_token_id: {max_token_id}") + H = k2.ctc_topo(max_token_id) + L = k2.Fsa.from_dict(torch.load(f"{lang_dir}/L_disambig.pt")) + + if Path("data/lm/G_3_gram.pt").is_file(): + logging.info("Loading pre-compiled G_3_gram") + d = torch.load("data/lm/G_3_gram.pt") + G = k2.Fsa.from_dict(d) + else: + logging.info("Loading G_3_gram.fst.txt") + with open("data/lm/G_3_gram.fst.txt") as f: + G = k2.Fsa.from_openfst(f.read(), acceptor=False) + torch.save(G.as_dict(), "G_3_gram.pt") + + first_token_disambig_id = lexicon.token_table["#0"] + first_word_disambig_id = lexicon.word_table["#0"] + + L = k2.arc_sort(L) + G = k2.arc_sort(G) + + logging.info("Intersecting L and G") + LG = k2.compose(L, G) + logging.info(f"LG shape: {LG.shape}") + + logging.info("Connecting LG") + LG = k2.connect(LG) + logging.info(f"LG shape after k2.connect: {LG.shape}") + + logging.info(type(LG.aux_labels)) + logging.info("Determinizing LG") + + LG = k2.determinize(LG) + logging.info(type(LG.aux_labels)) + + logging.info("Connecting LG after k2.determinize") + LG = k2.connect(LG) + + logging.info("Removing disambiguation symbols on LG") + + LG.labels[LG.labels >= first_token_disambig_id] = 0 + + assert isinstance(LG.aux_labels, k2.RaggedInt) + LG.aux_labels.values()[LG.aux_labels.values() >= first_word_disambig_id] = 0 + + LG = k2.remove_epsilon(LG) + logging.info(f"LG shape after k2.remove_epsilon: {LG.shape}") + + LG = k2.connect(LG) + LG.aux_labels = k2.ragged.remove_values_eq(LG.aux_labels, 0) + + logging.info("Arc sorting LG") + LG = k2.arc_sort(LG) + + logging.info("Composing H and LG") + # CAUTION: The name of the inner_labels is fixed + # to `tokens`. If you want to change it, please + # also change other places in icefall that are using + # it. + HLG = k2.compose(H, LG, inner_labels="tokens") + + logging.info("Connecting LG") + HLG = k2.connect(HLG) + + logging.info("Arc sorting LG") + HLG = k2.arc_sort(HLG) + logging.info(f"HLG.shape: {HLG.shape}") + + return HLG + + +def phone_based_HLG(): + if Path("data/lm/HLG.pt").is_file(): + return + + logging.info("Compiling phone based HLG") + HLG = compile_HLG("data/lang") + + logging.info("Saving HLG.pt to data/lm") + torch.save(HLG.as_dict(), "data/lm/HLG.pt") + + +def bpe_based_HLG(): + if Path("data/lm/HLG_bpe.pt").is_file(): + return + + logging.info("Compiling BPE based HLG") + HLG = compile_HLG("data/lang/bpe") + logging.info("Saving HLG_bpe.pt to data/lm") + torch.save(HLG.as_dict(), "data/lm/HLG_bpe.pt") + + +def main(): + phone_based_HLG() + bpe_based_HLG() + + +if __name__ == "__main__": + formatter = ( + "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s" + ) + + logging.basicConfig(format=formatter, level=logging.INFO) + + main() diff --git a/egs/librispeech/ASR/local/compute_fbank_librispeech.py b/egs/librispeech/ASR/local/compute_fbank_librispeech.py new file mode 100755 index 0000000000..947d9f8d9d --- /dev/null +++ b/egs/librispeech/ASR/local/compute_fbank_librispeech.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 + +""" +This file computes fbank features of the librispeech dataset. +Its looks for manifests in the directory data/manifests +and generated fbank features are saved in data/fbank. +""" + +import os +from pathlib import Path + +from lhotse import CutSet, Fbank, FbankConfig, LilcomHdf5Writer +from lhotse.recipes.utils import read_manifests_if_cached + +from icefall.utils import get_executor + + +def compute_fbank_librispeech(): + src_dir = Path("data/manifests") + output_dir = Path("data/fbank") + num_jobs = min(15, os.cpu_count()) + num_mel_bins = 80 + + dataset_parts = ( + "dev-clean", + "dev-other", + "test-clean", + "test-other", + "train-clean-100", + "train-clean-360", + "train-other-500", + ) + manifests = read_manifests_if_cached( + dataset_parts=dataset_parts, output_dir=src_dir + ) + assert manifests is not None + + extractor = Fbank(FbankConfig(num_mel_bins=num_mel_bins)) + + with get_executor() as ex: # Initialize the executor only once. + for partition, m in manifests.items(): + if (output_dir / f"cuts_{partition}.json.gz").is_file(): + print(f"{partition} already exists - skipping.") + continue + print("Processing", partition) + cut_set = CutSet.from_manifests( + recordings=m["recordings"], + supervisions=m["supervisions"], + ) + if "train" in partition: + cut_set = ( + cut_set + + cut_set.perturb_speed(0.9) + + cut_set.perturb_speed(1.1) + ) + cut_set = cut_set.compute_and_store_features( + extractor=extractor, + storage_path=f"{output_dir}/feats_{partition}", + # when an executor is specified, make more partitions + num_jobs=num_jobs if ex is None else 80, + executor=ex, + storage_type=LilcomHdf5Writer, + ) + cut_set.to_json(output_dir / f"cuts_{partition}.json.gz") + + +if __name__ == "__main__": + compute_fbank_librispeech() diff --git a/egs/librispeech/ASR/local/compute_fbank_musan.py b/egs/librispeech/ASR/local/compute_fbank_musan.py new file mode 100755 index 0000000000..d63131da89 --- /dev/null +++ b/egs/librispeech/ASR/local/compute_fbank_musan.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 + +""" +This file computes fbank features of the musan dataset. +Its looks for manifests in the directory data/manifests +and generated fbank features are saved in data/fbank. +""" + +import os +from pathlib import Path + +from lhotse import CutSet, Fbank, FbankConfig, LilcomHdf5Writer, combine +from lhotse.recipes.utils import read_manifests_if_cached + +from icefall.utils import get_executor + + +def compute_fbank_musan(): + src_dir = Path("data/manifests") + output_dir = Path("data/fbank") + num_jobs = min(15, os.cpu_count()) + num_mel_bins = 80 + + dataset_parts = ( + "music", + "speech", + "noise", + ) + manifests = read_manifests_if_cached( + dataset_parts=dataset_parts, output_dir=src_dir + ) + assert manifests is not None + + musan_cuts_path = output_dir / "cuts_musan.json.gz" + + if musan_cuts_path.is_file(): + print(f"{musan_cuts_path} already exists - skipping") + return + + print("Extracting features for Musan") + + extractor = Fbank(FbankConfig(num_mel_bins=num_mel_bins)) + + with get_executor() as ex: # Initialize the executor only once. + # create chunks of Musan with duration 5 - 10 seconds + musan_cuts = ( + CutSet.from_manifests( + recordings=combine( + part["recordings"] for part in manifests.values() + ) + ) + .cut_into_windows(10.0) + .filter(lambda c: c.duration > 5) + .compute_and_store_features( + extractor=extractor, + storage_path=f"{output_dir}/feats_musan", + num_jobs=num_jobs if ex is None else 80, + executor=ex, + storage_type=LilcomHdf5Writer, + ) + ) + musan_cuts.to_json(musan_cuts_path) + + +if __name__ == "__main__": + compute_fbank_musan() diff --git a/egs/librispeech/ASR/local/download_lm.py b/egs/librispeech/ASR/local/download_lm.py new file mode 100755 index 0000000000..0bdc2935ba --- /dev/null +++ b/egs/librispeech/ASR/local/download_lm.py @@ -0,0 +1,50 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2021 Xiaomi Corporation (authors: Fangjun Kuang) +""" +This file downloads librispeech LM files to data/lm +""" + +import gzip +import os +import shutil +from pathlib import Path + +from lhotse.utils import urlretrieve_progress +from tqdm.auto import tqdm + + +def download_lm(): + url = "http://www.openslr.org/resources/11" + target_dir = Path("data/lm") + + files_to_download = ( + "3-gram.pruned.1e-7.arpa.gz", + "4-gram.arpa.gz", + "librispeech-vocab.txt", + "librispeech-lexicon.txt", + ) + + for f in tqdm(files_to_download, desc="Downloading LibriSpeech LM files"): + filename = target_dir / f + if filename.is_file() is False: + urlretrieve_progress( + f"{url}/{f}", + filename=filename, + desc=f"Downloading {filename}", + ) + else: + print(f"{filename} already exists - skipping") + + if ".gz" in str(filename): + unzip_file = Path(os.path.splitext(filename)[0]) + if unzip_file.is_file() is False: + with gzip.open(filename, "rb") as f_in: + with open(unzip_file, "wb") as f_out: + shutil.copyfileobj(f_in, f_out) + else: + print(f"{unzip_file} already exist - skipping") + + +if __name__ == "__main__": + download_lm() diff --git a/egs/librispeech/ASR/local/parse_options.sh b/egs/librispeech/ASR/local/parse_options.sh new file mode 100755 index 0000000000..71fb9e5ea1 --- /dev/null +++ b/egs/librispeech/ASR/local/parse_options.sh @@ -0,0 +1,97 @@ +#!/usr/bin/env bash + +# Copyright 2012 Johns Hopkins University (Author: Daniel Povey); +# Arnab Ghoshal, Karel Vesely + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED +# WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE, +# MERCHANTABLITY OR NON-INFRINGEMENT. +# See the Apache 2 License for the specific language governing permissions and +# limitations under the License. + + +# Parse command-line options. +# To be sourced by another script (as in ". parse_options.sh"). +# Option format is: --option-name arg +# and shell variable "option_name" gets set to value "arg." +# The exception is --help, which takes no arguments, but prints the +# $help_message variable (if defined). + + +### +### The --config file options have lower priority to command line +### options, so we need to import them first... +### + +# Now import all the configs specified by command-line, in left-to-right order +for ((argpos=1; argpos<$#; argpos++)); do + if [ "${!argpos}" == "--config" ]; then + argpos_plus1=$((argpos+1)) + config=${!argpos_plus1} + [ ! -r $config ] && echo "$0: missing config '$config'" && exit 1 + . $config # source the config file. + fi +done + + +### +### Now we process the command line options +### +while true; do + [ -z "${1:-}" ] && break; # break if there are no arguments + case "$1" in + # If the enclosing script is called with --help option, print the help + # message and exit. Scripts should put help messages in $help_message + --help|-h) if [ -z "$help_message" ]; then echo "No help found." 1>&2; + else printf "$help_message\n" 1>&2 ; fi; + exit 0 ;; + --*=*) echo "$0: options to scripts must be of the form --name value, got '$1'" + exit 1 ;; + # If the first command-line argument begins with "--" (e.g. --foo-bar), + # then work out the variable name as $name, which will equal "foo_bar". + --*) name=`echo "$1" | sed s/^--// | sed s/-/_/g`; + # Next we test whether the variable in question is undefned-- if so it's + # an invalid option and we die. Note: $0 evaluates to the name of the + # enclosing script. + # The test [ -z ${foo_bar+xxx} ] will return true if the variable foo_bar + # is undefined. We then have to wrap this test inside "eval" because + # foo_bar is itself inside a variable ($name). + eval '[ -z "${'$name'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1; + + oldval="`eval echo \\$$name`"; + # Work out whether we seem to be expecting a Boolean argument. + if [ "$oldval" == "true" ] || [ "$oldval" == "false" ]; then + was_bool=true; + else + was_bool=false; + fi + + # Set the variable to the right value-- the escaped quotes make it work if + # the option had spaces, like --cmd "queue.pl -sync y" + eval $name=\"$2\"; + + # Check that Boolean-valued arguments are really Boolean. + if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then + echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2 + exit 1; + fi + shift 2; + ;; + *) break; + esac +done + + +# Check for an empty argument to the --cmd option, which can easily occur as a +# result of scripting errors. +[ ! -z "${cmd+xxx}" ] && [ -z "$cmd" ] && echo "$0: empty argument to --cmd option" 1>&2 && exit 1; + + +true; # so this script returns exit code 0. diff --git a/egs/librispeech/ASR/local/prepare_lang.py b/egs/librispeech/ASR/local/prepare_lang.py new file mode 100755 index 0000000000..b9d13f5bb4 --- /dev/null +++ b/egs/librispeech/ASR/local/prepare_lang.py @@ -0,0 +1,367 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2021 Xiaomi Corporation (authors: Fangjun Kuang) + +""" +This script takes as input a lexicon file "data/lang/lexicon.txt" +consisting of words and tokens (i.e., phones) and does the following: + +1. Add disambiguation symbols to the lexicon and generate lexicon_disambig.txt + +2. Generate tokens.txt, the token table mapping a token to a unique integer. + +3. Generate words.txt, the word table mapping a word to a unique integer. + +4. Generate L.pt, in k2 format. It can be loaded by + + d = torch.load("L.pt") + lexicon = k2.Fsa.from_dict(d) + +5. Generate L_disambig.pt, in k2 format. +""" +import math +import re +import sys +from collections import defaultdict +from pathlib import Path +from typing import Any, Dict, List, Tuple + +import k2 +import torch + +from icefall.lexicon import read_lexicon, write_lexicon + +Lexicon = List[Tuple[str, List[str]]] + + +def write_mapping(filename: str, sym2id: Dict[str, int]) -> None: + """Write a symbol to ID mapping to a file. + + Note: + No need to implement `read_mapping` as it can be done + through :func:`k2.SymbolTable.from_file`. + + Args: + filename: + Filename to save the mapping. + sym2id: + A dict mapping symbols to IDs. + Returns: + Return None. + """ + with open(filename, "w", encoding="utf-8") as f: + for sym, i in sym2id.items(): + f.write(f"{sym} {i}\n") + + +def get_tokens(lexicon: Lexicon) -> List[str]: + """Get tokens from a lexicon. + + Args: + lexicon: + It is the return value of :func:`read_lexicon`. + Returns: + Return a list of unique tokens. + """ + ans = set() + for _, tokens in lexicon: + ans.update(tokens) + sorted_ans = sorted(list(ans)) + return sorted_ans + + +def get_words(lexicon: Lexicon) -> List[str]: + """Get words from a lexicon. + + Args: + lexicon: + It is the return value of :func:`read_lexicon`. + Returns: + Return a list of unique words. + """ + ans = set() + for word, _ in lexicon: + ans.add(word) + sorted_ans = sorted(list(ans)) + return sorted_ans + + +def add_disambig_symbols(lexicon: Lexicon) -> Tuple[Lexicon, int]: + """It adds pseudo-token disambiguation symbols #1, #2 and so on + at the ends of tokens to ensure that all pronunciations are different, + and that none is a prefix of another. + + See also add_lex_disambig.pl from kaldi. + + Args: + lexicon: + It is returned by :func:`read_lexicon`. + Returns: + Return a tuple with two elements: + + - The output lexicon with disambiguation symbols + - The ID of the max disambiguation symbol that appears + in the lexicon + """ + + # (1) Work out the count of each token-sequence in the + # lexicon. + count = defaultdict(int) + for _, tokens in lexicon: + count[" ".join(tokens)] += 1 + + # (2) For each left sub-sequence of each token-sequence, note down + # that it exists (for identifying prefixes of longer strings). + issubseq = defaultdict(int) + for _, tokens in lexicon: + tokens = tokens.copy() + tokens.pop() + while tokens: + issubseq[" ".join(tokens)] = 1 + tokens.pop() + + # (3) For each entry in the lexicon: + # if the token sequence is unique and is not a + # prefix of another word, no disambig symbol. + # Else output #1, or #2, #3, ... if the same token-seq + # has already been assigned a disambig symbol. + ans = [] + + # We start with #1 since #0 has its own purpose + first_allowed_disambig = 1 + max_disambig = first_allowed_disambig - 1 + last_used_disambig_symbol_of = defaultdict(int) + + for word, tokens in lexicon: + tokenseq = " ".join(tokens) + assert tokenseq != "" + if issubseq[tokenseq] == 0 and count[tokenseq] == 1: + ans.append((word, tokens)) + continue + + cur_disambig = last_used_disambig_symbol_of[tokenseq] + if cur_disambig == 0: + cur_disambig = first_allowed_disambig + else: + cur_disambig += 1 + + if cur_disambig > max_disambig: + max_disambig = cur_disambig + last_used_disambig_symbol_of[tokenseq] = cur_disambig + tokenseq += f" #{cur_disambig}" + ans.append((word, tokenseq.split())) + return ans, max_disambig + + +def generate_id_map(symbols: List[str]) -> Dict[str, int]: + """Generate ID maps, i.e., map a symbol to a unique ID. + + Args: + symbols: + A list of unique symbols. + Returns: + A dict containing the mapping between symbols and IDs. + """ + return {sym: i for i, sym in enumerate(symbols)} + + +def add_self_loops( + arcs: List[List[Any]], disambig_token: int, disambig_word: int +) -> List[List[Any]]: + """Adds self-loops to states of an FST to propagate disambiguation symbols + through it. They are added on each state with non-epsilon output symbols + on at least one arc out of the state. + + See also fstaddselfloops.pl from Kaldi. One difference is that + Kaldi uses OpenFst style FSTs and it has multiple final states. + This function uses k2 style FSTs and it does not need to add self-loops + to the final state. + + The input label of a self-loop is `disambig_token`, while the output + label is `disambig_word`. + + Args: + arcs: + A list-of-list. The sublist contains + `[src_state, dest_state, label, aux_label, score]` + disambig_token: + It is the token ID of the symbol `#0`. + disambig_word: + It is the word ID of the symbol `#0`. + + Return: + Return new `arcs` containing self-loops. + """ + states_needs_self_loops = set() + for arc in arcs: + src, dst, ilabel, olabel, score = arc + if olabel != 0: + states_needs_self_loops.add(src) + + ans = [] + for s in states_needs_self_loops: + ans.append([s, s, disambig_token, disambig_word, 0]) + + return arcs + ans + + +def lexicon_to_fst( + lexicon: Lexicon, + token2id: Dict[str, int], + word2id: Dict[str, int], + sil_token: str = "SIL", + sil_prob: float = 0.5, + need_self_loops: bool = False, +) -> k2.Fsa: + """Convert a lexicon to an FST (in k2 format) with optional silence at + the beginning and end of each word. + + Args: + lexicon: + The input lexicon. See also :func:`read_lexicon` + token2id: + A dict mapping tokens to IDs. + word2id: + A dict mapping words to IDs. + sil_token: + The silence token. + sil_prob: + The probability for adding a silence at the beginning and end + of the word. + need_self_loops: + If True, add self-loop to states with non-epsilon output symbols + on at least one arc out of the state. The input label for this + self loop is `token2id["#0"]` and the output label is `word2id["#0"]`. + Returns: + Return an instance of `k2.Fsa` representing the given lexicon. + """ + assert sil_prob > 0.0 and sil_prob < 1.0 + # CAUTION: we use score, i.e, negative cost. + sil_score = math.log(sil_prob) + no_sil_score = math.log(1.0 - sil_prob) + + start_state = 0 + loop_state = 1 # words enter and leave from here + sil_state = 2 # words terminate here when followed by silence; this state + # has a silence transition to loop_state. + next_state = 3 # the next un-allocated state, will be incremented as we go. + arcs = [] + + assert token2id[""] == 0 + assert word2id[""] == 0 + + eps = 0 + + sil_token = token2id[sil_token] + + arcs.append([start_state, loop_state, eps, eps, no_sil_score]) + arcs.append([start_state, sil_state, eps, eps, sil_score]) + arcs.append([sil_state, loop_state, sil_token, eps, 0]) + + for word, tokens in lexicon: + assert len(tokens) > 0, f"{word} has no pronunciations" + cur_state = loop_state + + word = word2id[word] + tokens = [token2id[i] for i in tokens] + + for i in range(len(tokens) - 1): + w = word if i == 0 else eps + arcs.append([cur_state, next_state, tokens[i], w, 0]) + + cur_state = next_state + next_state += 1 + + # now for the last token of this word + # It has two out-going arcs, one to the loop state, + # the other one to the sil_state. + i = len(tokens) - 1 + w = word if i == 0 else eps + arcs.append([cur_state, loop_state, tokens[i], w, no_sil_score]) + arcs.append([cur_state, sil_state, tokens[i], w, sil_score]) + + if need_self_loops: + disambig_token = token2id["#0"] + disambig_word = word2id["#0"] + arcs = add_self_loops( + arcs, disambig_token=disambig_token, disambig_word=disambig_word, + ) + + final_state = next_state + arcs.append([loop_state, final_state, -1, -1, 0]) + arcs.append([final_state]) + + arcs = sorted(arcs, key=lambda arc: arc[0]) + arcs = [[str(i) for i in arc] for arc in arcs] + arcs = [" ".join(arc) for arc in arcs] + arcs = "\n".join(arcs) + + fsa = k2.Fsa.from_str(arcs, acceptor=False) + return fsa + + +def main(): + out_dir = Path("data/lang") + lexicon_filename = out_dir / "lexicon.txt" + sil_token = "SIL" + sil_prob = 0.5 + + lexicon = read_lexicon(lexicon_filename) + tokens = get_tokens(lexicon) + words = get_words(lexicon) + + lexicon_disambig, max_disambig = add_disambig_symbols(lexicon) + + for i in range(max_disambig + 1): + disambig = f"#{i}" + assert disambig not in tokens + tokens.append(f"#{i}") + + assert "" not in tokens + tokens = [""] + tokens + + assert "" not in words + assert "#0" not in words + assert "" not in words + assert "" not in words + + words = [""] + words + ["#0", "", ""] + + token2id = generate_id_map(tokens) + word2id = generate_id_map(words) + + write_mapping(out_dir / "tokens.txt", token2id) + write_mapping(out_dir / "words.txt", word2id) + write_lexicon(out_dir / "lexicon_disambig.txt", lexicon_disambig) + + L = lexicon_to_fst( + lexicon, + token2id=token2id, + word2id=word2id, + sil_token=sil_token, + sil_prob=sil_prob, + ) + + L_disambig = lexicon_to_fst( + lexicon_disambig, + token2id=token2id, + word2id=word2id, + sil_token=sil_token, + sil_prob=sil_prob, + need_self_loops=True, + ) + torch.save(L.as_dict(), out_dir / "L.pt") + torch.save(L_disambig.as_dict(), out_dir / "L_disambig.pt") + + if False: + # Just for debugging, will remove it + L.labels_sym = k2.SymbolTable.from_file(out_dir / "tokens.txt") + L.aux_labels_sym = k2.SymbolTable.from_file(out_dir / "words.txt") + L_disambig.labels_sym = L.labels_sym + L_disambig.aux_labels_sym = L.aux_labels_sym + L.draw(out_dir / "L.png", title="L") + L_disambig.draw(out_dir / "L_disambig.png", title="L_disambig") + + +if __name__ == "__main__": + main() diff --git a/egs/librispeech/ASR/local/prepare_lang_bpe.py b/egs/librispeech/ASR/local/prepare_lang_bpe.py new file mode 100755 index 0000000000..0c3e9ede54 --- /dev/null +++ b/egs/librispeech/ASR/local/prepare_lang_bpe.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2021 Xiaomi Corporation (authors: Fangjun Kuang) + +""" +This script takes as inputs the following two files: + + - data/lang/bpe/bpe.model, + - data/lang/bpe/words.txt + +and generates the following files in the directory data/lang/bpe: + + - lexicon.txt + - lexicon_disambig.txt + - L.pt + - L_disambig.pt + - tokens.txt +""" + +from pathlib import Path +from typing import Dict, List, Tuple + +import k2 +import sentencepiece as spm +import torch +from prepare_lang import ( + Lexicon, + add_disambig_symbols, + add_self_loops, + write_lexicon, + write_mapping, +) + + +def lexicon_to_fst_no_sil( + lexicon: Lexicon, + token2id: Dict[str, int], + word2id: Dict[str, int], + need_self_loops: bool = False, +) -> k2.Fsa: + """Convert a lexicon to an FST (in k2 format). + + Args: + lexicon: + The input lexicon. See also :func:`read_lexicon` + token2id: + A dict mapping tokens to IDs. + word2id: + A dict mapping words to IDs. + need_self_loops: + If True, add self-loop to states with non-epsilon output symbols + on at least one arc out of the state. The input label for this + self loop is `token2id["#0"]` and the output label is `word2id["#0"]`. + Returns: + Return an instance of `k2.Fsa` representing the given lexicon. + """ + loop_state = 0 # words enter and leave from here + next_state = 1 # the next un-allocated state, will be incremented as we go + + arcs = [] + + # The blank symbol is defined in local/train_bpe_model.py + assert token2id[""] == 0 + assert word2id[""] == 0 + + eps = 0 + + for word, pieces in lexicon: + assert len(pieces) > 0, f"{word} has no pronunciations" + cur_state = loop_state + + word = word2id[word] + pieces = [token2id[i] for i in pieces] + + for i in range(len(pieces) - 1): + w = word if i == 0 else eps + arcs.append([cur_state, next_state, pieces[i], w, 0]) + + cur_state = next_state + next_state += 1 + + # now for the last piece of this word + i = len(pieces) - 1 + w = word if i == 0 else eps + arcs.append([cur_state, loop_state, pieces[i], w, 0]) + + if need_self_loops: + disambig_token = token2id["#0"] + disambig_word = word2id["#0"] + arcs = add_self_loops( + arcs, disambig_token=disambig_token, disambig_word=disambig_word, + ) + + final_state = next_state + arcs.append([loop_state, final_state, -1, -1, 0]) + arcs.append([final_state]) + + arcs = sorted(arcs, key=lambda arc: arc[0]) + arcs = [[str(i) for i in arc] for arc in arcs] + arcs = [" ".join(arc) for arc in arcs] + arcs = "\n".join(arcs) + + fsa = k2.Fsa.from_str(arcs, acceptor=False) + return fsa + + +def generate_lexicon( + model_file: str, words: List[str] +) -> Tuple[Lexicon, Dict[str, int]]: + """Generate a lexicon from a BPE model. + + Args: + model_file: + Path to a sentencepiece model. + words: + A list of strings representing words. + Returns: + Return a tuple with two elements: + - A dict whose keys are words and values are the corresponding + word pieces. + - A dict representing the token symbol, mapping from tokens to IDs. + """ + sp = spm.SentencePieceProcessor() + sp.load(str(model_file)) + + words_pieces: List[List[str]] = sp.encode(words, out_type=str) + + lexicon = [] + for word, pieces in zip(words, words_pieces): + lexicon.append((word, pieces)) + + # The OOV word is + lexicon.append(("", [sp.id_to_piece(sp.unk_id())])) + + token2id: Dict[str, int] = dict() + for i in range(sp.vocab_size()): + token2id[sp.id_to_piece(i)] = i + + return lexicon, token2id + + +def main(): + lang_dir = Path("data/lang/bpe") + model_file = lang_dir / "bpe.model" + + word_sym_table = k2.SymbolTable.from_file(lang_dir / "words.txt") + + words = word_sym_table.symbols + + excluded = ["", "!SIL", "", "", "#0", "", ""] + for w in excluded: + if w in words: + words.remove(w) + + lexicon, token_sym_table = generate_lexicon(model_file, words) + + lexicon_disambig, max_disambig = add_disambig_symbols(lexicon) + + next_token_id = max(token_sym_table.values()) + 1 + for i in range(max_disambig + 1): + disambig = f"#{i}" + assert disambig not in token_sym_table + token_sym_table[disambig] = next_token_id + next_token_id += 1 + + word_sym_table.add("#0") + word_sym_table.add("") + word_sym_table.add("") + + write_mapping(lang_dir / "tokens.txt", token_sym_table) + + write_lexicon(lang_dir / "lexicon.txt", lexicon) + write_lexicon(lang_dir / "lexicon_disambig.txt", lexicon_disambig) + + L = lexicon_to_fst_no_sil( + lexicon, token2id=token_sym_table, word2id=word_sym_table, + ) + + L_disambig = lexicon_to_fst_no_sil( + lexicon_disambig, + token2id=token_sym_table, + word2id=word_sym_table, + need_self_loops=True, + ) + torch.save(L.as_dict(), lang_dir / "L.pt") + torch.save(L_disambig.as_dict(), lang_dir / "L_disambig.pt") + + if False: + # Just for debugging, will remove it + L.labels_sym = k2.SymbolTable.from_file(lang_dir / "tokens.txt") + L.aux_labels_sym = k2.SymbolTable.from_file(lang_dir / "words.txt") + L_disambig.labels_sym = L.labels_sym + L_disambig.aux_labels_sym = L.aux_labels_sym + L.draw(lang_dir / "L.svg", title="L") + L_disambig.draw(lang_dir / "L_disambig.svg", title="L_disambig") + + +if __name__ == "__main__": + main() diff --git a/egs/librispeech/ASR/local/test_prepare_lang.py b/egs/librispeech/ASR/local/test_prepare_lang.py new file mode 100755 index 0000000000..23ab53c7d1 --- /dev/null +++ b/egs/librispeech/ASR/local/test_prepare_lang.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2021 Xiaomi Corporation (authors: Fangjun Kuang) + +import os +import tempfile + +import k2 +from prepare_lang import ( + add_disambig_symbols, + generate_id_map, + get_phones, + get_words, + lexicon_to_fst, + read_lexicon, + write_lexicon, + write_mapping, +) + + +def generate_lexicon_file() -> str: + fd, filename = tempfile.mkstemp() + os.close(fd) + s = """ + !SIL SIL + SPN + SPN + f f + a a + foo f o o + bar b a r + bark b a r k + food f o o d + food2 f o o d + fo f o + """.strip() + with open(filename, "w") as f: + f.write(s) + return filename + + +def test_read_lexicon(filename: str): + lexicon = read_lexicon(filename) + phones = get_phones(lexicon) + words = get_words(lexicon) + print(lexicon) + print(phones) + print(words) + lexicon_disambig, max_disambig = add_disambig_symbols(lexicon) + print(lexicon_disambig) + print("max disambig:", f"#{max_disambig}") + + phones = ["", "SIL", "SPN"] + phones + for i in range(max_disambig + 1): + phones.append(f"#{i}") + words = [""] + words + + phone2id = generate_id_map(phones) + word2id = generate_id_map(words) + + print(phone2id) + print(word2id) + + write_mapping("phones.txt", phone2id) + write_mapping("words.txt", word2id) + + write_lexicon("a.txt", lexicon) + write_lexicon("a_disambig.txt", lexicon_disambig) + + fsa = lexicon_to_fst(lexicon, phone2id=phone2id, word2id=word2id) + fsa.labels_sym = k2.SymbolTable.from_file("phones.txt") + fsa.aux_labels_sym = k2.SymbolTable.from_file("words.txt") + fsa.draw("L.pdf", title="L") + + fsa_disambig = lexicon_to_fst( + lexicon_disambig, phone2id=phone2id, word2id=word2id + ) + fsa_disambig.labels_sym = k2.SymbolTable.from_file("phones.txt") + fsa_disambig.aux_labels_sym = k2.SymbolTable.from_file("words.txt") + fsa_disambig.draw("L_disambig.pdf", title="L_disambig") + + +def main(): + filename = generate_lexicon_file() + test_read_lexicon(filename) + os.remove(filename) + + +if __name__ == "__main__": + main() diff --git a/egs/librispeech/ASR/local/train_bpe_model.py b/egs/librispeech/ASR/local/train_bpe_model.py new file mode 100755 index 0000000000..b5c6c7541a --- /dev/null +++ b/egs/librispeech/ASR/local/train_bpe_model.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 + +""" +This script takes as input "data/lang/bpe/train.txt" +and generates "data/lang/bpe/bep.model". +""" + +# You can install sentencepiece via: +# +# pip install sentencepiece +# +# Due to an issue reported in +# https://github.com/google/sentencepiece/pull/642#issuecomment-857972030 +# +# Please install a version >=0.1.96 + +from pathlib import Path + +import sentencepiece as spm + +import shutil + + +def main(): + model_type = "unigram" + vocab_size = 5000 + model_prefix = f"data/lang/bpe/{model_type}_{vocab_size}" + train_text = "data/lang/bpe/train.txt" + character_coverage = 1.0 + input_sentence_size = 100000000 + + user_defined_symbols = ["", ""] + unk_id = len(user_defined_symbols) + # Note: unk_id is fixed to 2. + # If you change it, you should also change other + # places that are using it. + + model_file = Path(model_prefix + ".model") + if not model_file.is_file(): + spm.SentencePieceTrainer.train( + input=train_text, + vocab_size=vocab_size, + model_type=model_type, + model_prefix=model_prefix, + input_sentence_size=input_sentence_size, + character_coverage=character_coverage, + user_defined_symbols=user_defined_symbols, + unk_id=unk_id, + bos_id=-1, + eos_id=-1, + ) + + sp = spm.SentencePieceProcessor(model_file=str(model_file)) + vocab_size = sp.vocab_size() + + shutil.copyfile(model_file, "data/lang/bpe/bpe.model") + + +if __name__ == "__main__": + main() diff --git a/egs/librispeech/ASR/prepare.sh b/egs/librispeech/ASR/prepare.sh new file mode 100755 index 0000000000..406527b713 --- /dev/null +++ b/egs/librispeech/ASR/prepare.sh @@ -0,0 +1,150 @@ +#!/usr/bin/env bash + +set -eou pipefail + +nj=15 +stage=-1 +stop_stage=100 + +. local/parse_options.sh || exit 1 + +mkdir -p data + +log() { + # This function is from espnet + local fname=${BASH_SOURCE[1]##*/} + echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*" +} + +if [ $stage -le -1 ] && [ $stop_stage -ge -1 ]; then + log "stage -1: Download LM" + mkdir -p data/lm + ./local/download_lm.py +fi + +if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then + log "stage 0: Download data" + + # If you have pre-downloaded it to /path/to/LibriSpeech, + # you can create a symlink + # + # ln -sfv /path/to/LibriSpeech data/ + # + # The script checks that if + # + # data/LibriSpeech/test-clean/.completed exists, + # + # it will not re-download it. + # + # The same goes for dev-clean, dev-other, test-other, train-clean-100 + # train-clean-360, and train-other-500 + + mkdir -p data/LibriSpeech + lhotse download librispeech --full data + + # If you have pre-downloaded it to /path/to/musan, + # you can create a symlink + # + # ln -sfv /path/to/musan data/ + # + # and create a file data/.musan_completed + # to avoid downloading it again + if [ ! -f data/.musan_completed ]; then + lhotse download musan data + fi +fi + +if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then + log "Stage 1: Prepare librispeech manifest" + # We assume that you have downloaded the librispeech corpus + # to data/LibriSpeech + mkdir -p data/manifests + lhotse prepare librispeech -j $nj data/LibriSpeech data/manifests +fi + +if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then + log "Stage 2: Prepare musan manifest" + # We assume that you have downloaded the musan corpus + # to data/musan + mkdir -p data/manifests + lhotse prepare musan data/musan data/manifests +fi + +if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then + log "Stage 3: Compute fbank for librispeech" + mkdir -p data/fbank + ./local/compute_fbank_librispeech.py +fi + +if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then + log "Stage 4: Compute fbank for musan" + mkdir -p data/fbank + ./local/compute_fbank_musan.py +fi + +if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then + log "Stage 5: Prepare phone based lang" + # TODO: add BPE based lang + mkdir -p data/lang + + (echo '!SIL SIL'; echo ' SPN'; echo ' SPN'; ) | + cat - data/lm/librispeech-lexicon.txt | + sort | uniq > data/lang/lexicon.txt + + if [ ! -f data/lang/L_disambig.pt ]; then + ./local/prepare_lang.py + fi +fi + +if [ $stage -le 6 ] && [ $stop_stage -ge 6 ]; then + log "State 6: Prepare BPE based lang" + mkdir -p data/lang/bpe + cp data/lang/words.txt data/lang/bpe/ + + if [ ! -f data/lang/bpe/train.txt ]; then + log "Generate data for BPE training" + files=$( + find "data/LibriSpeech/train-clean-100" -name "*.trans.txt" + find "data/LibriSpeech/train-clean-360" -name "*.trans.txt" + find "data/LibriSpeech/train-other-500" -name "*.trans.txt" + ) + for f in ${files[@]}; do + cat $f | cut -d " " -f 2- + done > data/lang/bpe/train.txt + fi + + python3 ./local/train_bpe_model.py + + if [ ! -f data/lang/bpe/L_disambig.pt ]; then + ./local/prepare_lang_bpe.py + fi +fi + +if [ $stage -le 7 ] && [ $stop_stage -ge 7 ]; then + log "Stage 7: Prepare G" + # We assume you have install kaldilm, if not, please install + # it using: pip install kaldilm + + if [ ! -f data/lm/G_3_gram.fst.txt ]; then + # It is used in building HLG + python3 -m kaldilm \ + --read-symbol-table="data/lang/words.txt" \ + --disambig-symbol='#0' \ + --max-order=3 \ + data/lm/3-gram.pruned.1e-7.arpa > data/lm/G_3_gram.fst.txt + fi + + if [ ! -f data/lm/G_4_gram.fst.txt ]; then + # It is used for LM rescoring + python3 -m kaldilm \ + --read-symbol-table="data/lang/words.txt" \ + --disambig-symbol='#0' \ + --max-order=4 \ + data/lm/4-gram.arpa > data/lm/G_4_gram.fst.txt + fi +fi + +if [ $stage -le 8 ] && [ $stop_stage -ge 8 ]; then + log "Stage 8: Compile HLG" + python3 ./local/compile_hlg.py +fi diff --git a/egs/librispeech/ASR/tdnn_lstm_ctc/README.md b/egs/librispeech/ASR/tdnn_lstm_ctc/README.md new file mode 100644 index 0000000000..401f3e319f --- /dev/null +++ b/egs/librispeech/ASR/tdnn_lstm_ctc/README.md @@ -0,0 +1,22 @@ +## (To be filled in) + +It will contain: + +- How to run +- WERs + +```bash +cd $PWD/.. + +./prepare.sh + +./tdnn_lstm_ctc/train.py +``` + +If you have 4 GPUs and want to use GPU 1 and GPU 3 for DDP training, +you can do the following: + +``` +export CUDA_VISIBLE_DEVICES="1,3" +./tdnn_lstm_ctc/train.py --world-size=2 +``` diff --git a/egs/librispeech/ASR/tdnn_lstm_ctc/__init__.py b/egs/librispeech/ASR/tdnn_lstm_ctc/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/egs/librispeech/ASR/tdnn_lstm_ctc/decode.py b/egs/librispeech/ASR/tdnn_lstm_ctc/decode.py new file mode 100755 index 0000000000..2c45b4e317 --- /dev/null +++ b/egs/librispeech/ASR/tdnn_lstm_ctc/decode.py @@ -0,0 +1,419 @@ +#!/usr/bin/env python3 + + +import argparse +import logging +from collections import defaultdict +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +import k2 +import torch +import torch.nn as nn +from model import TdnnLstm + +from icefall.checkpoint import average_checkpoints, load_checkpoint +from icefall.dataset.librispeech import LibriSpeechAsrDataModule +from icefall.decode import ( + get_lattice, + nbest_decoding, + one_best_decoding, + rescore_with_n_best_list, + rescore_with_whole_lattice, +) +from icefall.lexicon import Lexicon +from icefall.utils import ( + AttributeDict, + get_texts, + setup_logger, + store_transcripts, + write_error_stats, +) + + +def get_parser(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + + parser.add_argument( + "--epoch", + type=int, + default=9, + help="It specifies the checkpoint to use for decoding." + "Note: Epoch counts from 0.", + ) + parser.add_argument( + "--avg", + type=int, + default=5, + help="Number of checkpoints to average. Automatically select " + "consecutive checkpoints before the checkpoint specified by " + "'--epoch'. ", + ) + return parser + + +def get_params() -> AttributeDict: + params = AttributeDict( + { + "exp_dir": Path("tdnn_lstm_ctc/exp/"), + "lang_dir": Path("data/lang"), + "lm_dir": Path("data/lm"), + "feature_dim": 80, + "subsampling_factor": 3, + "search_beam": 20, + "output_beam": 5, + "min_active_states": 30, + "max_active_states": 10000, + "use_double_scores": True, + # Possible values for method: + # - 1best + # - nbest + # - nbest-rescoring + # - whole-lattice-rescoring + "method": "1best", + # num_paths is used when method is "nbest" and "nbest-rescoring" + "num_paths": 30, + } + ) + return params + + +def decode_one_batch( + params: AttributeDict, + model: nn.Module, + HLG: k2.Fsa, + batch: dict, + lexicon: Lexicon, + G: Optional[k2.Fsa] = None, +) -> Dict[str, List[List[int]]]: + """Decode one batch and return the result in a dict. The dict has the + following format: + + - key: It indicates the setting used for decoding. For example, + if no rescoring is used, the key is the string `no_rescore`. + If LM rescoring is used, the key is the string `lm_scale_xxx`, + where `xxx` is the value of `lm_scale`. An example key is + `lm_scale_0.7` + - value: It contains the decoding result. `len(value)` equals to + batch size. `value[i]` is the decoding result for the i-th + utterance in the given batch. + Args: + params: + It's the return value of :func:`get_params`. + + - params.method is "1best", it uses 1best decoding without LM rescoring. + - params.method is "nbest", it uses nbest decoding without LM rescoring. + - params.method is "nbest-rescoring", it uses nbest LM rescoring. + - params.method is "whole-lattice-rescoring", it uses whole lattice LM + rescoring. + + model: + The neural model. + HLG: + The decoding graph. + batch: + It is the return value from iterating + `lhotse.dataset.K2SpeechRecognitionDataset`. See its documentation + for the format of the `batch`. + lexicon: + It contains word symbol table. + G: + An LM. It is not None when params.method is "nbest-rescoring" + or "whole-lattice-rescoring". In general, the G in HLG + is a 3-gram LM, while this G is a 4-gram LM. + Returns: + Return the decoding result. See above description for the format of + the returned dict. + """ + device = HLG.device + feature = batch["inputs"] + assert feature.ndim == 3 + feature = feature.to(device) + # at entry, feature is [N, T, C] + + feature = feature.permute(0, 2, 1) # now feature is [N, C, T] + + nnet_output = model(feature) + # nnet_output is [N, T, C] + + supervisions = batch["supervisions"] + + supervision_segments = torch.stack( + ( + supervisions["sequence_idx"], + supervisions["start_frame"] // params.subsampling_factor, + supervisions["num_frames"] // params.subsampling_factor, + ), + 1, + ).to(torch.int32) + + lattice = get_lattice( + nnet_output=nnet_output, + HLG=HLG, + supervision_segments=supervision_segments, + search_beam=params.search_beam, + output_beam=params.output_beam, + min_active_states=params.min_active_states, + max_active_states=params.max_active_states, + ) + + if params.method in ["1best", "nbest"]: + if params.method == "1best": + best_path = one_best_decoding( + lattice=lattice, use_double_scores=params.use_double_scores + ) + key = "no_rescore" + else: + best_path = nbest_decoding( + lattice=lattice, + num_paths=params.num_paths, + use_double_scores=params.use_double_scores, + ) + key = f"no_rescore-{params.num_paths}" + hyps = get_texts(best_path) + hyps = [[lexicon.word_table[i] for i in ids] for ids in hyps] + return {key: hyps} + + assert params.method in ["nbest-rescoring", "whole-lattice-rescoring"] + + lm_scale_list = [0.8, 0.9, 1.0, 1.1, 1.2, 1.3] + lm_scale_list += [1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2.0] + + if params.method == "nbest-rescoring": + best_path_dict = rescore_with_n_best_list( + lattice=lattice, + G=G, + num_paths=params.num_paths, + lm_scale_list=lm_scale_list, + ) + else: + best_path_dict = rescore_with_whole_lattice( + lattice=lattice, G_with_epsilon_loops=G, lm_scale_list=lm_scale_list + ) + + ans = dict() + for lm_scale_str, best_path in best_path_dict.items(): + hyps = get_texts(best_path) + hyps = [[lexicon.word_table[i] for i in ids] for ids in hyps] + ans[lm_scale_str] = hyps + return ans + + +def decode_dataset( + dl: torch.utils.data.DataLoader, + params: AttributeDict, + model: nn.Module, + HLG: k2.Fsa, + lexicon: Lexicon, + G: Optional[k2.Fsa] = None, +) -> Dict[str, List[Tuple[List[int], List[int]]]]: + """Decode dataset. + + Args: + dl: + PyTorch's dataloader containing the dataset to decode. + params: + It is returned by :func:`get_params`. + model: + The neural model. + HLG: + The decoding graph. + lexicon: + It contains word symbol table. + G: + An LM. It is not None when params.method is "nbest-rescoring" + or "whole-lattice-rescoring". In general, the G in HLG + is a 3-gram LM, while this G is a 4-gram LM. + Returns: + Return a dict, whose key may be "no-rescore" if no LM rescoring + is used, or it may be "lm_scale_0.7" if LM rescoring is used. + Its value is a list of tuples. Each tuple contains two elements: + The first is the reference transcript, and the second is the + predicted result. + """ + results = [] + + num_cuts = 0 + tot_num_cuts = len(dl.dataset.cuts) + + results = defaultdict(list) + for batch_idx, batch in enumerate(dl): + texts = batch["supervisions"]["text"] + + hyps_dict = decode_one_batch( + params=params, + model=model, + HLG=HLG, + batch=batch, + lexicon=lexicon, + G=G, + ) + + for lm_scale, hyps in hyps_dict.items(): + this_batch = [] + assert len(hyps) == len(texts) + for hyp_words, ref_text in zip(hyps, texts): + ref_words = ref_text.split() + this_batch.append((ref_words, hyp_words)) + + results[lm_scale].extend(this_batch) + + num_cuts += len(batch["supervisions"]["text"]) + + if batch_idx % 100 == 0: + logging.info( + f"batch {batch_idx}, cuts processed until now is " + f"{num_cuts}/{tot_num_cuts} " + f"({float(num_cuts)/tot_num_cuts*100:.6f}%)" + ) + return results + + +def save_results( + params: AttributeDict, + test_set_name: str, + results_dict: Dict[str, List[Tuple[List[int], List[int]]]], +): + test_set_wers = dict() + for key, results in results_dict.items(): + recog_path = params.exp_dir / f"recogs-{test_set_name}-{key}.txt" + store_transcripts(filename=recog_path, texts=results) + logging.info(f"The transcripts are stored in {recog_path}") + + # The following prints out WERs, per-word error statistics and aligned + # ref/hyp pairs. + errs_filename = params.exp_dir / f"errs-{test_set_name}-{key}.txt" + with open(errs_filename, "w") as f: + wer = write_error_stats(f, f"{test_set_name}-{key}", results) + test_set_wers[key] = wer + + logging.info("Wrote detailed error stats to {}".format(errs_filename)) + + test_set_wers = sorted(test_set_wers.items(), key=lambda x: x[1]) + errs_info = params.exp_dir / f"wer-summary-{test_set_name}.txt" + with open(errs_info, "w") as f: + print("settings\tWER", file=f) + for key, val in test_set_wers: + print("{}\t{}".format(key, val), file=f) + + s = "\nFor {}, WER of different settings are:\n".format(test_set_name) + note = "\tbest for {}".format(test_set_name) + for key, val in test_set_wers: + s += "{}\t{}{}\n".format(key, val, note) + note = "" + logging.info(s) + + +@torch.no_grad() +def main(): + parser = get_parser() + LibriSpeechAsrDataModule.add_arguments(parser) + args = parser.parse_args() + + params = get_params() + params.update(vars(args)) + + setup_logger(f"{params.exp_dir}/log/log-decode") + logging.info("Decoding started") + logging.info(params) + + lexicon = Lexicon(params.lang_dir) + max_phone_id = max(lexicon.tokens) + + device = torch.device("cpu") + if torch.cuda.is_available(): + device = torch.device("cuda", 0) + + logging.info(f"device: {device}") + + HLG = k2.Fsa.from_dict(torch.load("data/lm/HLG.pt")) + HLG = HLG.to(device) + assert HLG.requires_grad is False + + if not hasattr(HLG, "lm_scores"): + HLG.lm_scores = HLG.scores.clone() + + if params.method in ["nbest-rescoring", "whole-lattice-rescoring"]: + if not (params.lm_dir / "G_4_gram.pt").is_file(): + logging.info("Loading G_4_gram.fst.txt") + logging.warning("It may take 8 minutes.") + with open(params.lm_dir / "G_4_gram.fst.txt") as f: + first_word_disambig_id = lexicon.words["#0"] + + G = k2.Fsa.from_openfst(f.read(), acceptor=False) + # G.aux_labels is not needed in later computations, so + # remove it here. + del G.aux_labels + # CAUTION: The following line is crucial. + # Arcs entering the back-off state have label equal to #0. + # We have to change it to 0 here. + G.labels[G.labels >= first_word_disambig_id] = 0 + G = k2.Fsa.from_fsas([G]).to(device) + G = k2.arc_sort(G) + torch.save(G.as_dict(), params.lm_dir / "G_4_gram.pt") + else: + logging.info("Loading pre-compiled G_4_gram.pt") + d = torch.load(params.lm_dir / "G_4_gram.pt") + G = k2.Fsa.from_dict(d).to(device) + + if params.method == "whole-lattice-rescoring": + # Add epsilon self-loops to G as we will compose + # it with the whole lattice later + G = k2.add_epsilon_self_loops(G) + G = k2.arc_sort(G) + G = G.to(device) + + # G.lm_scores is used to replace HLG.lm_scores during + # LM rescoring. + G.lm_scores = G.scores.clone() + else: + G = None + + model = TdnnLstm( + num_features=params.feature_dim, + num_classes=max_phone_id + 1, # +1 for the blank symbol + subsampling_factor=params.subsampling_factor, + ) + if params.avg == 1: + load_checkpoint(f"{params.exp_dir}/epoch-{params.epoch}.pt", model) + else: + start = params.epoch - params.avg + 1 + filenames = [] + for i in range(start, params.epoch + 1): + if start >= 0: + filenames.append(f"{params.exp_dir}/epoch-{i}.pt") + logging.info(f"averaging {filenames}") + model.load_state_dict(average_checkpoints(filenames)) + + model.to(device) + model.eval() + + librispeech = LibriSpeechAsrDataModule(args) + # CAUTION: `test_sets` is for displaying only. + # If you want to skip test-clean, you have to skip + # it inside the for loop. That is, use + # + # if test_set == 'test-clean': continue + # + test_sets = ["test-clean", "test-other"] + for test_set, test_dl in zip(test_sets, librispeech.test_dataloaders()): + results_dict = decode_dataset( + dl=test_dl, + params=params, + model=model, + HLG=HLG, + lexicon=lexicon, + G=G, + ) + + save_results( + params=params, test_set_name=test_set, results_dict=results_dict + ) + + logging.info("Done!") + + +if __name__ == "__main__": + main() diff --git a/egs/librispeech/ASR/tdnn_lstm_ctc/model.py b/egs/librispeech/ASR/tdnn_lstm_ctc/model.py new file mode 100644 index 0000000000..0dc4228dcc --- /dev/null +++ b/egs/librispeech/ASR/tdnn_lstm_ctc/model.py @@ -0,0 +1,86 @@ +import torch +import torch.nn as nn + + +class TdnnLstm(nn.Module): + def __init__( + self, num_features: int, num_classes: int, subsampling_factor: int = 3 + ) -> None: + """ + Args: + num_features: + The input dimension of the model. + num_classes: + The output dimension of the model. + subsampling_factor: + It reduces the number of output frames by this factor. + """ + super().__init__() + self.num_features = num_features + self.num_classes = num_classes + self.subsampling_factor = subsampling_factor + self.tdnn = nn.Sequential( + nn.Conv1d( + in_channels=num_features, + out_channels=500, + kernel_size=3, + stride=1, + padding=1, + ), + nn.ReLU(inplace=True), + nn.BatchNorm1d(num_features=500, affine=False), + nn.Conv1d( + in_channels=500, + out_channels=500, + kernel_size=3, + stride=1, + padding=1, + ), + nn.ReLU(inplace=True), + nn.BatchNorm1d(num_features=500, affine=False), + nn.Conv1d( + in_channels=500, + out_channels=500, + kernel_size=3, + stride=self.subsampling_factor, # stride: subsampling_factor! + padding=1, + ), + nn.ReLU(inplace=True), + nn.BatchNorm1d(num_features=500, affine=False), + ) + self.lstms = nn.ModuleList( + [ + nn.LSTM(input_size=500, hidden_size=500, num_layers=1) + for _ in range(5) + ] + ) + self.lstm_bnorms = nn.ModuleList( + [nn.BatchNorm1d(num_features=500, affine=False) for _ in range(5)] + ) + self.dropout = nn.Dropout(0.2) + self.linear = nn.Linear(in_features=500, out_features=self.num_classes) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x: + Its shape is [N, C, T] + + Returns: + The output tensor has shape [N, T, C] + """ + x = self.tdnn(x) + x = x.permute(2, 0, 1) # (N, C, T) -> (T, N, C) -> how LSTM expects it + for lstm, bnorm in zip(self.lstms, self.lstm_bnorms): + x_new, _ = lstm(x) + x_new = bnorm(x_new.permute(1, 2, 0)).permute( + 2, 0, 1 + ) # (T, N, C) -> (N, C, T) -> (T, N, C) + x_new = self.dropout(x_new) + x = x_new + x # skip connections + x = x.transpose( + 1, 0 + ) # (T, N, C) -> (N, T, C) -> linear expects "features" in the last dim + x = self.linear(x) + x = nn.functional.log_softmax(x, dim=-1) + return x diff --git a/egs/librispeech/ASR/tdnn_lstm_ctc/train.py b/egs/librispeech/ASR/tdnn_lstm_ctc/train.py new file mode 100755 index 0000000000..d94a2f7258 --- /dev/null +++ b/egs/librispeech/ASR/tdnn_lstm_ctc/train.py @@ -0,0 +1,568 @@ +#!/usr/bin/env python3 + +# This is just at the very beginning ... + +import argparse +import logging +from pathlib import Path +from shutil import copyfile +from typing import Optional + +import k2 +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn as nn +import torch.optim as optim +from lhotse.utils import fix_random_seed +from model import TdnnLstm +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.nn.utils import clip_grad_value_ +from torch.optim.lr_scheduler import StepLR +from torch.utils.tensorboard import SummaryWriter + +from icefall.checkpoint import load_checkpoint +from icefall.checkpoint import save_checkpoint as save_checkpoint_impl +from icefall.dataset.librispeech import LibriSpeechAsrDataModule +from icefall.dist import cleanup_dist, setup_dist +from icefall.graph_compiler import CtcTrainingGraphCompiler +from icefall.lexicon import Lexicon +from icefall.utils import ( + AttributeDict, + encode_supervisions, + setup_logger, + str2bool, +) + + +def get_parser(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + + parser.add_argument( + "--world-size", + type=int, + default=1, + help="Number of GPUs for DDP training.", + ) + + parser.add_argument( + "--master-port", + type=int, + default=12354, + help="Master port to use for DDP training.", + ) + + parser.add_argument( + "--tensorboard", + type=str2bool, + default=True, + help="Should various information be logged in tensorboard.", + ) + + # TODO: add extra arguments and support DDP training. + # Currently, only single GPU training is implemented. Will add + # DDP training once single GPU training is finished. + return parser + + +def get_params() -> AttributeDict: + """Return a dict containing training parameters. + + All training related parameters that are not passed from the commandline + is saved in the variable `params`. + + Commandline options are merged into `params` after they are parsed, so + you can also access them via `params`. + + Explanation of options saved in `params`: + + - exp_dir: It specifies the directory where all training related + files, e.g., checkpoints, log, etc, are saved + + - lang_dir: It contains language related input files such as + "lexicon.txt" + + - lr: It specifies the initial learning rate + + - feature_dim: The model input dim. It has to match the one used + in computing features. + + - weight_decay: The weight_decay for the optimizer. + + - subsampling_factor: The subsampling factor for the model. + + - start_epoch: If it is not zero, load checkpoint `start_epoch-1` + and continue training from that checkpoint. + + - num_epochs: Number of epochs to train. + + - best_train_loss: Best training loss so far. It is used to select + the model that has the lowest training loss. It is + updated during the training. + + - best_valid_loss: Best validation loss so far. It is used to select + the model that has the lowest validation loss. It is + updated during the training. + + - best_train_epoch: It is the epoch that has the best training loss. + + - best_valid_epoch: It is the epoch that has the best validation loss. + + - batch_idx_train: Used to writing statistics to tensorboard. It + contains number of batches trained so far across + epochs. + + - log_interval: Print training loss if batch_idx % log_interval` is 0 + + - valid_interval: Run validation if batch_idx % valid_interval` is 0 + + - beam_size: It is used in k2.ctc_loss + + - reduction: It is used in k2.ctc_loss + + - use_double_scores: It is used in k2.ctc_loss + """ + params = AttributeDict( + { + "exp_dir": Path("tdnn_lstm_ctc/exp"), + "lang_dir": Path("data/lang"), + "lr": 1e-3, + "feature_dim": 80, + "weight_decay": 5e-4, + "subsampling_factor": 3, + "start_epoch": 0, + "num_epochs": 10, + "best_train_loss": float("inf"), + "best_valid_loss": float("inf"), + "best_train_epoch": -1, + "best_valid_epoch": -1, + "batch_idx_train": 0, + "log_interval": 10, + "valid_interval": 1000, + "beam_size": 10, + "reduction": "sum", + "use_double_scores": True, + } + ) + + return params + + +def load_checkpoint_if_available( + params: AttributeDict, + model: nn.Module, + optimizer: Optional[torch.optim.Optimizer] = None, + scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None, +) -> None: + """Load checkpoint from file. + + If params.start_epoch is positive, it will load the checkpoint from + `params.start_epoch - 1`. Otherwise, this function does nothing. + + Apart from loading state dict for `model`, `optimizer` and `scheduler`, + it also updates `best_train_epoch`, `best_train_loss`, `best_valid_epoch`, + and `best_valid_loss` in `params`. + + Args: + params: + The return value of :func:`get_params`. + model: + The training model. + optimizer: + The optimizer that we are using. + scheduler: + The learning rate scheduler we are using. + Returns: + Return None. + """ + if params.start_epoch <= 0: + return + + filename = params.exp_dir / f"epoch-{params.start_epoch-1}.pt" + saved_params = load_checkpoint( + filename, + model=model, + optimizer=optimizer, + scheduler=scheduler, + ) + + keys = [ + "best_train_epoch", + "best_valid_epoch", + "batch_idx_train", + "best_train_loss", + "best_valid_loss", + ] + for k in keys: + params[k] = saved_params[k] + + return saved_params + + +def save_checkpoint( + params: AttributeDict, + model: nn.Module, + optimizer: torch.optim.Optimizer, + scheduler: torch.optim.lr_scheduler._LRScheduler, + rank: int = 0, +) -> None: + """Save model, optimizer, scheduler and training stats to file. + + Args: + params: + It is returned by :func:`get_params`. + model: + The training model. + """ + if rank != 0: + return + filename = params.exp_dir / f"epoch-{params.cur_epoch}.pt" + save_checkpoint_impl( + filename=filename, + model=model, + params=params, + optimizer=optimizer, + scheduler=scheduler, + rank=rank, + ) + + if params.best_train_epoch == params.cur_epoch: + best_train_filename = params.exp_dir / "best-train-loss.pt" + copyfile(src=filename, dst=best_train_filename) + + if params.best_valid_epoch == params.cur_epoch: + best_valid_filename = params.exp_dir / "best-valid-loss.pt" + copyfile(src=filename, dst=best_valid_filename) + + +def compute_loss( + params: AttributeDict, + model: nn.Module, + batch: dict, + graph_compiler: CtcTrainingGraphCompiler, + is_training: bool, +): + """ + Compute CTC loss given the model and its inputs. + + Args: + params: + Parameters for training. See :func:`get_params`. + model: + The model for training. It is an instance of TdnnLstm in our case. + batch: + A batch of data. See `lhotse.dataset.K2SpeechRecognitionDataset()` + for the content in it. + graph_compiler: + It is used to build a decoding graph from a ctc topo and training + transcript. The training transcript is contained in the given `batch`, + while the ctc topo is built when this compiler is instantiated. + is_training: + True for training. False for validation. When it is True, this + function enables autograd during computation; when it is False, it + disables autograd. + """ + device = graph_compiler.device + feature = batch["inputs"] + # at entry, feature is [N, T, C] + feature = feature.permute(0, 2, 1) # now feature is [N, C, T] + assert feature.ndim == 3 + feature = feature.to(device) + + with torch.set_grad_enabled(is_training): + nnet_output = model(feature) + # nnet_output is [N, T, C] + + # NOTE: We need `encode_supervisions` to sort sequences with + # different duration in decreasing order, required by + # `k2.intersect_dense` called in `k2.ctc_loss` + supervisions = batch["supervisions"] + supervision_segments, texts = encode_supervisions( + supervisions, subsampling_factor=params.subsampling_factor + ) + decoding_graph = graph_compiler.compile(texts) + + dense_fsa_vec = k2.DenseFsaVec( + nnet_output, + supervision_segments, + allow_truncate=params.subsampling_factor - 1, + ) + + loss = k2.ctc_loss( + decoding_graph=decoding_graph, + dense_fsa_vec=dense_fsa_vec, + output_beam=params.beam_size, + reduction=params.reduction, + use_double_scores=params.use_double_scores, + ) + + assert loss.requires_grad == is_training + + # train_frames and valid_frames are used for printing. + if is_training: + params.train_frames = supervision_segments[:, 2].sum().item() + else: + params.valid_frames = supervision_segments[:, 2].sum().item() + + return loss + + +def compute_validation_loss( + params: AttributeDict, + model: nn.Module, + graph_compiler: CtcTrainingGraphCompiler, + valid_dl: torch.utils.data.DataLoader, + world_size: int = 1, +) -> None: + """Run the validation process. The validation loss + is saved in `params.valid_loss`. + """ + model.eval() + + tot_loss = 0.0 + tot_frames = 0.0 + for batch_idx, batch in enumerate(valid_dl): + loss = compute_loss( + params=params, + model=model, + batch=batch, + graph_compiler=graph_compiler, + is_training=False, + ) + assert loss.requires_grad is False + + loss_cpu = loss.detach().cpu().item() + tot_loss += loss_cpu + tot_frames += params.valid_frames + + if world_size > 1: + s = torch.tensor([tot_loss, tot_frames], device=loss.device) + dist.all_reduce(s, op=dist.ReduceOp.SUM) + s = s.cpu().tolist() + tot_loss = s[0] + tot_frames = s[1] + + params.valid_loss = tot_loss / tot_frames + + if params.valid_loss < params.best_valid_loss: + params.best_valid_epoch = params.cur_epoch + params.best_valid_loss = params.valid_loss + + +def train_one_epoch( + params: AttributeDict, + model: nn.Module, + optimizer: torch.optim.Optimizer, + graph_compiler: CtcTrainingGraphCompiler, + train_dl: torch.utils.data.DataLoader, + valid_dl: torch.utils.data.DataLoader, + tb_writer: Optional[SummaryWriter] = None, + world_size: int = 1, +) -> None: + """Train the model for one epoch. + + The training loss from the mean of all frames is saved in + `params.train_loss`. It runs the validation process every + `params.valid_interval` batches. + + Args: + params: + It is returned by :func:`get_params`. + model: + The model for training. + optimizer: + The optimizer we are using. + graph_compiler: + It is used to convert transcripts to FSAs. + train_dl: + Dataloader for the training dataset. + valid_dl: + Dataloader for the validation dataset. + tb_writer: + Writer to write log messages to tensorboard. + world_size: + Number of nodes in DDP training. If it is 1, DDP is disabled. + """ + model.train() + + tot_loss = 0.0 # sum of losses over all batches + tot_frames = 0.0 # sum of frames over all batches + for batch_idx, batch in enumerate(train_dl): + params.batch_idx_train += 1 + batch_size = len(batch["supervisions"]["text"]) + + loss = compute_loss( + params=params, + model=model, + batch=batch, + graph_compiler=graph_compiler, + is_training=True, + ) + + # NOTE: We use reduction==sum and loss is computed over utterances + # in the batch and there is no normalization to it so far. + + optimizer.zero_grad() + loss.backward() + clip_grad_value_(model.parameters(), 5.0) + optimizer.step() + + loss_cpu = loss.detach().cpu().item() + + tot_frames += params.train_frames + tot_loss += loss_cpu + tot_avg_loss = tot_loss / tot_frames + + if batch_idx % params.log_interval == 0: + logging.info( + f"Epoch {params.cur_epoch}, batch {batch_idx}, " + f"batch avg loss {loss_cpu/params.train_frames:.4f}, " + f"total avg loss: {tot_avg_loss:.4f}, " + f"batch size: {batch_size}" + ) + + if batch_idx > 0 and batch_idx % params.valid_interval == 0: + compute_validation_loss( + params=params, + model=model, + graph_compiler=graph_compiler, + valid_dl=valid_dl, + world_size=world_size, + ) + model.train() + logging.info( + f"Epoch {params.cur_epoch}, valid loss {params.valid_loss:.4f}," + f" best valid loss: {params.best_valid_loss:.4f} " + f"best valid epoch: {params.best_valid_epoch}" + ) + + params.train_loss = tot_loss / tot_frames + + if params.train_loss < params.best_train_loss: + params.best_train_epoch = params.cur_epoch + params.best_train_loss = params.train_loss + + +def run(rank, world_size, args): + """ + Args: + rank: + It is a value between 0 and `world_size-1`, which is + passed automatically by `mp.spawn()` in :func:`main`. + The node with rank 0 is responsible for saving checkpoint. + world_size: + Number of GPUs for DDP training. + args: + The return value of get_parser().parse_args() + """ + params = get_params() + params.update(vars(args)) + + fix_random_seed(42) + if world_size > 1: + setup_dist(rank, world_size, params.master_port) + + setup_logger(f"{params.exp_dir}/log/log-train") + logging.info("Training started") + logging.info(params) + + if args.tensorboard and rank == 0: + tb_writer = SummaryWriter(log_dir=f"{params.exp_dir}/tensorboard") + else: + tb_writer = None + + lexicon = Lexicon(params.lang_dir) + max_phone_id = max(lexicon.tokens) + + device = torch.device("cpu") + if torch.cuda.is_available(): + device = torch.device("cuda", rank) + + graph_compiler = CtcTrainingGraphCompiler(lexicon=lexicon, device=device) + + model = TdnnLstm( + num_features=params.feature_dim, + num_classes=max_phone_id + 1, # +1 for the blank symbol + subsampling_factor=params.subsampling_factor, + ) + + checkpoints = load_checkpoint_if_available(params=params, model=model) + + model.to(device) + if world_size > 1: + model = DDP(model, device_ids=[rank]) + + optimizer = optim.AdamW( + model.parameters(), + lr=params.lr, + weight_decay=params.weight_decay, + ) + scheduler = StepLR(optimizer, step_size=8, gamma=0.1) + + optimizer.load_state_dict(checkpoints["optimizer"]) + scheduler.load_state_dict(checkpoints["scheduler"]) + + librispeech = LibriSpeechAsrDataModule(args) + train_dl = librispeech.train_dataloaders() + valid_dl = librispeech.valid_dataloaders() + + for epoch in range(params.start_epoch, params.num_epochs): + train_dl.sampler.set_epoch(epoch) + + if epoch > params.start_epoch: + logging.info(f"epoch {epoch}, lr: {scheduler.get_last_lr()[0]}") + + if tb_writer is not None: + tb_writer.add_scalar( + "train/lr", + scheduler.get_last_lr()[0], + params.batch_idx_train, + ) + tb_writer.add_scalar("train/epoch", epoch, params.batch_idx_train) + + params.cur_epoch = epoch + + train_one_epoch( + params=params, + model=model, + optimizer=optimizer, + graph_compiler=graph_compiler, + train_dl=train_dl, + valid_dl=valid_dl, + tb_writer=tb_writer, + world_size=world_size, + ) + + scheduler.step() + + save_checkpoint( + params=params, + model=model, + optimizer=optimizer, + scheduler=scheduler, + rank=rank, + ) + + logging.info("Done!") + if world_size > 1: + torch.distributed.barrier() + cleanup_dist() + + +def main(): + parser = get_parser() + LibriSpeechAsrDataModule.add_arguments(parser) + args = parser.parse_args() + + world_size = args.world_size + assert world_size >= 1 + if world_size > 1: + mp.spawn(run, args=(world_size, args), nprocs=world_size, join=True) + else: + run(rank=0, world_size=1, args=args) + + +if __name__ == "__main__": + main() diff --git a/icefall/__init__.py b/icefall/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/icefall/bpe_graph_compiler.py b/icefall/bpe_graph_compiler.py new file mode 100644 index 0000000000..e22cf4edcb --- /dev/null +++ b/icefall/bpe_graph_compiler.py @@ -0,0 +1,74 @@ +from pathlib import Path +from typing import List, Union + +import k2 +import sentencepiece as spm +import torch + + +class BpeCtcTrainingGraphCompiler(object): + def __init__( + self, + lang_dir: Path, + device: Union[str, torch.device] = "cpu", + sos_token: str = "", + eos_token: str = "", + ) -> None: + """ + Args: + lang_dir: + This directory is expected to contain the following files: + + - bpe.model + - words.txt + device: + It indicates CPU or CUDA. + sos_token: + The word piece that represents sos. + eos_token: + The word piece that represents eos. + """ + lang_dir = Path(lang_dir) + model_file = lang_dir / "bpe.model" + sp = spm.SentencePieceProcessor() + sp.load(str(model_file)) + self.sp = sp + self.word_table = k2.SymbolTable.from_file(lang_dir / "words.txt") + self.device = device + + self.sos_id = self.sp.piece_to_id(sos_token) + self.eos_id = self.sp.piece_to_id(eos_token) + + assert self.sos_id != self.sp.unk_id() + assert self.eos_id != self.sp.unk_id() + + def texts_to_ids(self, texts: List[str]) -> List[List[int]]: + """Convert a list of texts to a list-of-list of piece IDs. + + Args: + texts: + It is a list of strings. Each string consists of space(s) + separated words. An example containing two strings is given below: + + ['HELLO ICEFALL', 'HELLO k2'] + Returns: + Return a list-of-list of piece IDs. + """ + return self.sp.encode(texts, out_type=int) + + def compile( + self, piece_ids: List[List[int]], modified: bool = False, + ) -> k2.Fsa: + """Build a ctc graph from a list-of-list piece IDs. + + Args: + piece_ids: + It is a list-of-list integer IDs. + modified: + See :func:`k2.ctc_graph` for its meaning. + Return: + Return an FsaVec, which is the result of composing a + CTC topology with linear FSAs constructed from the given + piece IDs. + """ + return k2.ctc_graph(piece_ids, modified=modified, device=self.device) diff --git a/icefall/checkpoint.py b/icefall/checkpoint.py new file mode 100644 index 0000000000..e45df4fe47 --- /dev/null +++ b/icefall/checkpoint.py @@ -0,0 +1,131 @@ +import logging +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +import torch +import torch.nn as nn +from torch.cuda.amp import GradScaler +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import Optimizer +from torch.optim.lr_scheduler import _LRScheduler + + +def save_checkpoint( + filename: Path, + model: Union[nn.Module, DDP], + params: Optional[Dict[str, Any]] = None, + optimizer: Optional[Optimizer] = None, + scheduler: Optional[_LRScheduler] = None, + scaler: Optional[GradScaler] = None, + rank: int = 0, +) -> None: + """Save training information to a file. + + Args: + filename: + The checkpoint filename. + model: + The model to be saved. We only save its `state_dict()`. + params: + User defined parameters, e.g., epoch, loss. + optimizer: + The optimizer to be saved. We only save its `state_dict()`. + scheduler: + The scheduler to be saved. We only save its `state_dict()`. + scalar: + The GradScaler to be saved. We only save its `state_dict()`. + rank: + Used in DDP. We save checkpoint only for the node whose rank is 0. + Returns: + Return None. + """ + if rank != 0: + return + + logging.info(f"Saving checkpoint to {filename}") + + if isinstance(model, DDP): + model = model.module + + checkpoint = { + "model": model.state_dict(), + "optimizer": optimizer.state_dict() if optimizer is not None else None, + "scheduler": scheduler.state_dict() if scheduler is not None else None, + "grad_scaler": scaler.state_dict() if scaler is not None else None, + } + + if params: + for k, v in params.items(): + assert k not in checkpoint + checkpoint[k] = v + + torch.save(checkpoint, filename) + + +def load_checkpoint( + filename: Path, + model: nn.Module, + optimizer: Optional[Optimizer] = None, + scheduler: Optional[_LRScheduler] = None, + scaler: Optional[GradScaler] = None, +) -> Dict[str, Any]: + """ + TODO: document it + """ + logging.info(f"Loading checkpoint from {filename}") + checkpoint = torch.load(filename, map_location="cpu") + + if next(iter(checkpoint["model"])).startswith("module."): + logging.info("Loading checkpoint saved by DDP") + + dst_state_dict = model.state_dict() + src_state_dict = checkpoint["model"] + for key in dst_state_dict.keys(): + src_key = "{}.{}".format("module", key) + dst_state_dict[key] = src_state_dict.pop(src_key) + assert len(src_state_dict) == 0 + model.load_state_dict(dst_state_dict, strict=False) + else: + model.load_state_dict(checkpoint["model"], strict=False) + + checkpoint.pop("model") + + def load(name, obj): + s = checkpoint[name] + if obj and s: + obj.load_state_dict(s) + checkpoint.pop(name) + + load("optimizer", optimizer) + load("scheduler", scheduler) + load("grad_scaler", scaler) + + return checkpoint + + +def average_checkpoints(filenames: List[Path]) -> dict: + """Average a list of checkpoints. + + Args: + filenames: + Filenames of the checkpoints to be averaged. We assume all + checkpoints are saved by :func:`save_checkpoint`. + Returns: + Return a dict (i.e., state_dict) which is the average of all + model state dicts contained in the checkpoints. + """ + n = len(filenames) + + avg = torch.load(filenames[0], map_location="cpu")["model"] + for i in range(1, n): + state_dict = torch.load(filenames[i], map_location="cpu")["model"] + for k in avg: + avg[k] += state_dict[k] + + for k in avg: + if avg[k].is_floating_point(): + avg[k] /= n + else: + avg[k] //= n + + return avg diff --git a/icefall/dataset/__init__.py b/icefall/dataset/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/icefall/dataset/asr_datamodule.py b/icefall/dataset/asr_datamodule.py new file mode 100644 index 0000000000..aae7af9ce9 --- /dev/null +++ b/icefall/dataset/asr_datamodule.py @@ -0,0 +1,248 @@ +import argparse +import logging +from pathlib import Path +from typing import List, Union + +from lhotse import Fbank, FbankConfig, load_manifest +from lhotse.dataset import ( + BucketingSampler, + CutConcatenate, + CutMix, + K2SpeechRecognitionDataset, + SingleCutSampler, + SpecAugment, +) +from lhotse.dataset.input_strategies import OnTheFlyFeatures +from torch.utils.data import DataLoader + +from icefall.dataset.datamodule import DataModule +from icefall.utils import str2bool + + +class AsrDataModule(DataModule): + """ + DataModule for K2 ASR experiments. + It assumes there is always one train and valid dataloader, + but there can be multiple test dataloaders (e.g. LibriSpeech test-clean + and test-other). + + It contains all the common data pipeline modules used in ASR + experiments, e.g.: + - dynamic batch size, + - bucketing samplers, + - cut concatenation, + - augmentation, + - on-the-fly feature extraction + + This class should be derived for specific corpora used in ASR tasks. + """ + + @classmethod + def add_arguments(cls, parser: argparse.ArgumentParser): + super().add_arguments(parser) + group = parser.add_argument_group( + title="ASR data related options", + description="These options are used for the preparation of " + "PyTorch DataLoaders from Lhotse CutSet's -- they control the " + "effective batch sizes, sampling strategies, applied data " + "augmentations, etc.", + ) + group.add_argument( + "--feature-dir", + type=Path, + default=Path("data/fbank"), + help="Path to directory with train/valid/test cuts.", + ) + group.add_argument( + "--max-duration", + type=int, + default=500.0, + help="Maximum pooled recordings duration (seconds) in a " + "single batch. You can reduce it if it causes CUDA OOM.", + ) + group.add_argument( + "--bucketing-sampler", + type=str2bool, + default=False, + help="When enabled, the batches will come from buckets of " + "similar duration (saves padding frames).", + ) + group.add_argument( + "--num-buckets", + type=int, + default=30, + help="The number of buckets for the BucketingSampler" + "(you might want to increase it for larger datasets).", + ) + group.add_argument( + "--concatenate-cuts", + type=str2bool, + default=True, + help="When enabled, utterances (cuts) will be concatenated " + "to minimize the amount of padding.", + ) + group.add_argument( + "--duration-factor", + type=float, + default=1.0, + help="Determines the maximum duration of a concatenated cut " + "relative to the duration of the longest cut in a batch.", + ) + group.add_argument( + "--gap", + type=float, + default=1.0, + help="The amount of padding (in seconds) inserted between " + "concatenated cuts. This padding is filled with noise when " + "noise augmentation is used.", + ) + group.add_argument( + "--on-the-fly-feats", + type=str2bool, + default=False, + help="When enabled, use on-the-fly cut mixing and feature " + "extraction. Will drop existing precomputed feature manifests " + "if available.", + ) + + def train_dataloaders(self) -> DataLoader: + logging.info("About to get train cuts") + cuts_train = self.train_cuts() + + logging.info("About to get Musan cuts") + cuts_musan = load_manifest(self.args.feature_dir / "cuts_musan.json.gz") + + logging.info("About to create train dataset") + transforms = [CutMix(cuts=cuts_musan, prob=0.5, snr=(10, 20))] + if self.args.concatenate_cuts: + logging.info( + f"Using cut concatenation with duration factor " + f"{self.args.duration_factor} and gap {self.args.gap}." + ) + # Cut concatenation should be the first transform in the list, + # so that if we e.g. mix noise in, it will fill the gaps between + # different utterances. + transforms = [ + CutConcatenate( + duration_factor=self.args.duration_factor, gap=self.args.gap + ) + ] + transforms + + input_transforms = [ + SpecAugment( + num_frame_masks=2, + features_mask_size=27, + num_feature_masks=2, + frames_mask_size=100, + ) + ] + + train = K2SpeechRecognitionDataset( + cuts_train, + cut_transforms=transforms, + input_transforms=input_transforms, + ) + + if self.args.on_the_fly_feats: + # NOTE: the PerturbSpeed transform should be added only if we + # remove it from data prep stage. + # Add on-the-fly speed perturbation; since originally it would + # have increased epoch size by 3, we will apply prob 2/3 and use + # 3x more epochs. + # Speed perturbation probably should come first before + # concatenation, but in principle the transforms order doesn't have + # to be strict (e.g. could be randomized) + # transforms = [PerturbSpeed(factors=[0.9, 1.1], p=2/3)] + transforms # noqa + # Drop feats to be on the safe side. + cuts_train = cuts_train.drop_features() + train = K2SpeechRecognitionDataset( + cuts=cuts_train, + cut_transforms=transforms, + input_strategy=OnTheFlyFeatures( + Fbank(FbankConfig(num_mel_bins=80)) + ), + input_transforms=input_transforms, + ) + + if self.args.bucketing_sampler: + logging.info("Using BucketingSampler.") + train_sampler = BucketingSampler( + cuts_train, + max_duration=self.args.max_duration, + shuffle=True, + num_buckets=self.args.num_buckets, + ) + else: + logging.info("Using SingleCutSampler.") + train_sampler = SingleCutSampler( + cuts_train, + max_duration=self.args.max_duration, + shuffle=True, + ) + logging.info("About to create train dataloader") + train_dl = DataLoader( + train, + sampler=train_sampler, + batch_size=None, + num_workers=4, + persistent_workers=True, + ) + return train_dl + + def valid_dataloaders(self) -> DataLoader: + logging.info("About to get dev cuts") + cuts_valid = self.valid_cuts() + + logging.info("About to create dev dataset") + if self.args.on_the_fly_feats: + cuts_valid = cuts_valid.drop_features() + validate = K2SpeechRecognitionDataset( + cuts_valid.drop_features(), + input_strategy=OnTheFlyFeatures( + Fbank(FbankConfig(num_mel_bins=80)) + ), + ) + else: + validate = K2SpeechRecognitionDataset(cuts_valid) + valid_sampler = SingleCutSampler( + cuts_valid, + max_duration=self.args.max_duration, + ) + logging.info("About to create dev dataloader") + valid_dl = DataLoader( + validate, + sampler=valid_sampler, + batch_size=None, + num_workers=2, + persistent_workers=True, + ) + return valid_dl + + def test_dataloaders(self) -> Union[DataLoader, List[DataLoader]]: + cuts = self.test_cuts() + is_list = isinstance(cuts, list) + test_loaders = [] + if not is_list: + cuts = [cuts] + + for cuts_test in cuts: + logging.debug("About to create test dataset") + test = K2SpeechRecognitionDataset( + cuts_test, + input_strategy=OnTheFlyFeatures( + Fbank(FbankConfig(num_mel_bins=80)) + ), + ) + sampler = SingleCutSampler( + cuts_test, max_duration=self.args.max_duration + ) + logging.debug("About to create test dataloader") + test_dl = DataLoader( + test, batch_size=None, sampler=sampler, num_workers=1 + ) + test_loaders.append(test_dl) + + if is_list: + return test_loaders + else: + return test_loaders[0] diff --git a/icefall/dataset/datamodule.py b/icefall/dataset/datamodule.py new file mode 100644 index 0000000000..8560c5db0e --- /dev/null +++ b/icefall/dataset/datamodule.py @@ -0,0 +1,43 @@ +import argparse +from typing import List, Union + +from lhotse import CutSet +from torch.utils.data import DataLoader + + +class DataModule: + """ + Contains dataset-related code. It is intended to read/construct Lhotse cuts, + and create Dataset/Sampler/DataLoader out of them. + + There is a separate method to create each of train/valid/test DataLoader. + In principle, there might be multiple DataLoaders for each of + train/valid/test + (e.g. when a corpus has multiple test sets). + The API of this class allows to return lists of CutSets/DataLoaders. + """ + + def __init__(self, args: argparse.Namespace): + self.args = args + + @classmethod + def add_arguments(cls, parser: argparse.ArgumentParser): + pass + + def train_cuts(self) -> Union[CutSet, List[CutSet]]: + raise NotImplementedError() + + def valid_cuts(self) -> Union[CutSet, List[CutSet]]: + raise NotImplementedError() + + def test_cuts(self) -> Union[CutSet, List[CutSet]]: + raise NotImplementedError() + + def train_dataloaders(self) -> Union[DataLoader, List[DataLoader]]: + raise NotImplementedError() + + def valid_dataloaders(self) -> Union[DataLoader, List[DataLoader]]: + raise NotImplementedError() + + def test_dataloaders(self) -> Union[DataLoader, List[DataLoader]]: + raise NotImplementedError() diff --git a/icefall/dataset/librispeech.py b/icefall/dataset/librispeech.py new file mode 100644 index 0000000000..5c18041ed8 --- /dev/null +++ b/icefall/dataset/librispeech.py @@ -0,0 +1,68 @@ +import argparse +import logging +from functools import lru_cache +from typing import List + +from lhotse import CutSet, load_manifest + +from icefall.dataset.asr_datamodule import AsrDataModule +from icefall.utils import str2bool + + +class LibriSpeechAsrDataModule(AsrDataModule): + """ + LibriSpeech ASR data module. Can be used for 100h subset + (``--full-libri false``) or full 960h set. + The train and valid cuts for standard Libri splits are + concatenated into a single CutSet/DataLoader. + """ + + @classmethod + def add_arguments(cls, parser: argparse.ArgumentParser): + super().add_arguments(parser) + group = parser.add_argument_group(title="LibriSpeech specific options") + group.add_argument( + "--full-libri", + type=str2bool, + default=True, + help="When enabled, use 960h LibriSpeech.", + ) + + @lru_cache() + def train_cuts(self) -> CutSet: + logging.info("About to get train cuts") + cuts_train = load_manifest( + self.args.feature_dir / "cuts_train-clean-100.json.gz" + ) + if self.args.full_libri: + cuts_train = ( + cuts_train + + load_manifest( + self.args.feature_dir / "cuts_train-clean-360.json.gz" + ) + + load_manifest( + self.args.feature_dir / "cuts_train-other-500.json.gz" + ) + ) + return cuts_train + + @lru_cache() + def valid_cuts(self) -> CutSet: + logging.info("About to get dev cuts") + cuts_valid = load_manifest( + self.args.feature_dir / "cuts_dev-clean.json.gz" + ) + load_manifest(self.args.feature_dir / "cuts_dev-other.json.gz") + return cuts_valid + + @lru_cache() + def test_cuts(self) -> List[CutSet]: + test_sets = ["test-clean", "test-other"] + cuts = [] + for test_set in test_sets: + logging.debug("About to get test cuts") + cuts.append( + load_manifest( + self.args.feature_dir / f"cuts_{test_set}.json.gz" + ) + ) + return cuts diff --git a/icefall/decode.py b/icefall/decode.py new file mode 100644 index 0000000000..ed08405fa0 --- /dev/null +++ b/icefall/decode.py @@ -0,0 +1,712 @@ +import logging +from typing import Dict, List, Optional, Tuple, Union + +import k2 +import torch +import torch.nn as nn + + +def _intersect_device( + a_fsas: k2.Fsa, + b_fsas: k2.Fsa, + b_to_a_map: torch.Tensor, + sorted_match_a: bool, + batch_size: int = 50, +) -> k2.Fsa: + """This is a wrapper of k2.intersect_device and its purpose is to split + b_fsas into several batches and process each batch separately to avoid + CUDA OOM error. + + The arguments and return value of this function are the same as + k2.intersect_device. + """ + num_fsas = b_fsas.shape[0] + if num_fsas <= batch_size: + return k2.intersect_device( + a_fsas, b_fsas, b_to_a_map=b_to_a_map, sorted_match_a=sorted_match_a + ) + + num_batches = (num_fsas + batch_size - 1) // batch_size + splits = [] + for i in range(num_batches): + start = i * batch_size + end = min(start + batch_size, num_fsas) + splits.append((start, end)) + + ans = [] + for start, end in splits: + indexes = torch.arange(start, end).to(b_to_a_map) + + fsas = k2.index(b_fsas, indexes) + b_to_a = k2.index(b_to_a_map, indexes) + path_lattice = k2.intersect_device( + a_fsas, fsas, b_to_a_map=b_to_a, sorted_match_a=sorted_match_a + ) + ans.append(path_lattice) + + return k2.cat(ans) + + +def get_lattice( + nnet_output: torch.Tensor, + HLG: k2.Fsa, + supervision_segments: torch.Tensor, + search_beam: float, + output_beam: float, + min_active_states: int, + max_active_states: int, + subsampling_factor: int = 1, +) -> k2.Fsa: + """Get the decoding lattice from a decoding graph and neural + network output. + + Args: + nnet_output: + It is the output of a neural model of shape `[N, T, C]`. + HLG: + An Fsa, the decoding graph. See also `compile_HLG.py`. + supervision_segments: + A 2-D **CPU** tensor of dtype `torch.int32` with 3 columns. + Each row contains information for a supervision segment. Column 0 + is the `sequence_index` indicating which sequence this segment + comes from; column 1 specifies the `start_frame` of this segment + within the sequence; column 2 contains the `duration` of this + segment. + search_beam: + Decoding beam, e.g. 20. Smaller is faster, larger is more exact + (less pruning). This is the default value; it may be modified by + `min_active_states` and `max_active_states`. + output_beam: + Beam to prune output, similar to lattice-beam in Kaldi. Relative + to best path of output. + min_active_states: + Minimum number of FSA states that are allowed to be active on any given + frame for any given intersection/composition task. This is advisory, + in that it will try not to have fewer than this number active. + Set it to zero if there is no constraint. + max_active_states: + Maximum number of FSA states that are allowed to be active on any given + frame for any given intersection/composition task. This is advisory, + in that it will try not to exceed that but may not always succeed. + You can use a very large number if no constraint is needed. + subsampling_factor: + The subsampling factor of the model. + Returns: + A lattice containing the decoding result. + """ + dense_fsa_vec = k2.DenseFsaVec( + nnet_output, supervision_segments, allow_truncate=subsampling_factor - 1 + ) + + lattice = k2.intersect_dense_pruned( + HLG, + dense_fsa_vec, + search_beam=search_beam, + output_beam=output_beam, + min_active_states=min_active_states, + max_active_states=max_active_states, + ) + + return lattice + + +def one_best_decoding( + lattice: k2.Fsa, use_double_scores: bool = True +) -> k2.Fsa: + """Get the best path from a lattice. + + Args: + lattice: + The decoding lattice returned by :func:`get_lattice`. + use_double_scores: + True to use double precision floating point in the computation. + False to use single precision. + Return: + An FsaVec containing linear paths. + """ + best_path = k2.shortest_path(lattice, use_double_scores=use_double_scores) + return best_path + + +def nbest_decoding( + lattice: k2.Fsa, num_paths: int, use_double_scores: bool = True +) -> k2.Fsa: + """It implements something like CTC prefix beam search using n-best lists. + + The basic idea is to first extra n-best paths from the given lattice, + build a word seqs from these paths, and compute the total scores + of these sequences in the log-semiring. The one with the max score + is used as the decoding output. + + Caution: + Don't be confused by `best` in the name `n-best`. Paths are selected + randomly, not by ranking their scores. + + Args: + lattice: + The decoding lattice, returned by :func:`get_lattice`. + num_paths: + It specifies the size `n` in n-best. Note: Paths are selected randomly + and those containing identical word sequences are remove dand only one + of them is kept. + use_double_scores: + True to use double precision floating point in the computation. + False to use single precision. + Returns: + An FsaVec containing linear FSAs. + """ + # First, extract `num_paths` paths for each sequence. + # path is a k2.RaggedInt with axes [seq][path][arc_pos] + path = k2.random_paths(lattice, num_paths=num_paths, use_double_scores=True) + + # word_seq is a k2.RaggedInt sharing the same shape as `path` + # but it contains word IDs. Note that it also contains 0s and -1s. + # The last entry in each sublist is -1. + word_seq = k2.index(lattice.aux_labels, path) + # Note: the above operation supports also the case when + # lattice.aux_labels is a ragged tensor. In that case, + # `remove_axis=True` is used inside the pybind11 binding code, + # so the resulting `word_seq` still has 3 axes, like `path`. + # The 3 axes are [seq][path][word_id] + + # Remove 0 (epsilon) and -1 from word_seq + word_seq = k2.ragged.remove_values_leq(word_seq, 0) + + # Remove sequences with identical word sequences. + # + # k2.ragged.unique_sequences will reorder paths within a seq. + # `new2old` is a 1-D torch.Tensor mapping from the output path index + # to the input path index. + # new2old.numel() == unique_word_seqs.tot_size(1) + unique_word_seq, _, new2old = k2.ragged.unique_sequences( + word_seq, need_num_repeats=False, need_new2old_indexes=True + ) + # Note: unique_word_seq still has the same axes as word_seq + + seq_to_path_shape = k2.ragged.get_layer(unique_word_seq.shape(), 0) + + # path_to_seq_map is a 1-D torch.Tensor. + # path_to_seq_map[i] is the seq to which the i-th path belongs + path_to_seq_map = seq_to_path_shape.row_ids(1) + + # Remove the seq axis. + # Now unique_word_seq has only two axes [path][word] + unique_word_seq = k2.ragged.remove_axis(unique_word_seq, 0) + + # word_fsa is an FsaVec with axes [path][state][arc] + word_fsa = k2.linear_fsa(unique_word_seq) + + # add epsilon self loops since we will use + # k2.intersect_device, which treats epsilon as a normal symbol + word_fsa_with_epsilon_loops = k2.add_epsilon_self_loops(word_fsa) + + # lattice has token IDs as labels and word IDs as aux_labels. + # inv_lattice has word IDs as labels and token IDs as aux_labels + inv_lattice = k2.invert(lattice) + inv_lattice = k2.arc_sort(inv_lattice) + + path_lattice = _intersect_device( + inv_lattice, + word_fsa_with_epsilon_loops, + b_to_a_map=path_to_seq_map, + sorted_match_a=True, + ) + # path_lat has word IDs as labels and token IDs as aux_labels + + path_lattice = k2.top_sort(k2.connect(path_lattice)) + + tot_scores = path_lattice.get_tot_scores( + use_double_scores=use_double_scores, log_semiring=False + ) + + # RaggedFloat currently supports float32 only. + # If Ragged is wrapped, we can use k2.RaggedDouble here + ragged_tot_scores = k2.RaggedFloat( + seq_to_path_shape, tot_scores.to(torch.float32) + ) + + argmax_indexes = k2.ragged.argmax_per_sublist(ragged_tot_scores) + + # Since we invoked `k2.ragged.unique_sequences`, which reorders + # the index from `path`, we use `new2old` here to convert argmax_indexes + # to the indexes into `path`. + # + # Use k2.index here since argmax_indexes' dtype is torch.int32 + best_path_indexes = k2.index(new2old, argmax_indexes) + + path_2axes = k2.ragged.remove_axis(path, 0) + + # best_path is a k2.RaggedInt with 2 axes [path][arc_pos] + best_path = k2.index(path_2axes, best_path_indexes) + + # labels is a k2.RaggedInt with 2 axes [path][token_id] + # Note that it contains -1s. + labels = k2.index(lattice.labels.contiguous(), best_path) + + labels = k2.ragged.remove_values_eq(labels, -1) + + # lattice.aux_labels is a k2.RaggedInt tensor with 2 axes, so + # aux_labels is also a k2.RaggedInt with 2 axes + aux_labels = k2.index(lattice.aux_labels, best_path.values()) + + best_path_fsa = k2.linear_fsa(labels) + best_path_fsa.aux_labels = aux_labels + return best_path_fsa + + +def compute_am_and_lm_scores( + lattice: k2.Fsa, + word_fsa_with_epsilon_loops: k2.Fsa, + path_to_seq_map: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute AM scores of n-best lists (represented as word_fsas). + + Args: + lattice: + An FsaVec, e.g., the return value of :func:`get_lattice` + It must have the attribute `lm_scores`. + word_fsa_with_epsilon_loops: + An FsaVec representing an n-best list. Note that it has been processed + by `k2.add_epsilon_self_loops`. + path_to_seq_map: + A 1-D torch.Tensor with dtype torch.int32. path_to_seq_map[i] indicates + which sequence the i-th Fsa in word_fsa_with_epsilon_loops belongs to. + path_to_seq_map.numel() == word_fsas_with_epsilon_loops.arcs.dim0(). + Returns: + Return a tuple containing two 1-D torch.Tensors: (am_scores, lm_scores). + Each tensor's `numel()' equals to `word_fsas_with_epsilon_loops.shape[0]` + """ + assert len(lattice.shape) == 3 + assert hasattr(lattice, "lm_scores") + + # k2.compose() currently does not support b_to_a_map. To void + # replicating `lats`, we use k2.intersect_device here. + # + # lattice has token IDs as `labels` and word IDs as aux_labels, so we + # need to invert it here. + inv_lattice = k2.invert(lattice) + + # Now the `labels` of inv_lattice are word IDs (a 1-D torch.Tensor) + # and its `aux_labels` are token IDs ( a k2.RaggedInt with 2 axes) + + # Remove its `aux_labels` since it is not needed in the + # following computation + del inv_lattice.aux_labels + inv_lattice = k2.arc_sort(inv_lattice) + + path_lattice = _intersect_device( + inv_lattice, + word_fsa_with_epsilon_loops, + b_to_a_map=path_to_seq_map, + sorted_match_a=True, + ) + + path_lattice = k2.top_sort(k2.connect(path_lattice)) + + # The `scores` of every arc consists of `am_scores` and `lm_scores` + path_lattice.scores = path_lattice.scores - path_lattice.lm_scores + + am_scores = path_lattice.get_tot_scores( + use_double_scores=True, log_semiring=False + ) + + path_lattice.scores = path_lattice.lm_scores + + lm_scores = path_lattice.get_tot_scores( + use_double_scores=True, log_semiring=False + ) + + return am_scores.to(torch.float32), lm_scores.to(torch.float32) + + +def rescore_with_n_best_list( + lattice: k2.Fsa, G: k2.Fsa, num_paths: int, lm_scale_list: List[float] +) -> Dict[str, k2.Fsa]: + """Decode using n-best list with LM rescoring. + + `lattice` is a decoding lattice with 3 axes. This function first + extracts `num_paths` paths from `lattice` for each sequence using + `k2.random_paths`. The `am_scores` of these paths are computed. + For each path, its `lm_scores` is computed using `G` (which is an LM). + The final `tot_scores` is the sum of `am_scores` and `lm_scores`. + The path with the largest `tot_scores` within a sequence is used + as the decoding output. + + Args: + lattice: + An FsaVec. It can be the return value of :func:`get_lattice`. + G: + An FsaVec representing the language model (LM). Note that it + is an FsaVec, but it contains only one Fsa. + num_paths: + It is the size `n` in `n-best` list. + lm_scale_list: + A list containing lm_scale values. + Returns: + A dict of FsaVec, whose key is an lm_scale and the value is the + best decoding path for each sequence in the lattice. + """ + device = lattice.device + + assert len(lattice.shape) == 3 + assert hasattr(lattice, "aux_labels") + assert hasattr(lattice, "lm_scores") + + assert G.shape == (1, None, None) + assert G.device == device + assert hasattr(G, "aux_labels") is False + + # First, extract `num_paths` paths for each sequence. + # path is a k2.RaggedInt with axes [seq][path][arc_pos] + path = k2.random_paths(lattice, num_paths=num_paths, use_double_scores=True) + + # word_seq is a k2.RaggedInt sharing the same shape as `path` + # but it contains word IDs. Note that it also contains 0s and -1s. + # The last entry in each sublist is -1. + word_seq = k2.index(lattice.aux_labels, path) + + # Remove epsilons and -1 from word_seq + word_seq = k2.ragged.remove_values_leq(word_seq, 0) + + # Remove paths that has identical word sequences. + # + # unique_word_seq is still a k2.RaggedInt with 3 axes [seq][path][word] + # except that there are no repeated paths with the same word_seq + # within a sequence. + # + # num_repeats is also a k2.RaggedInt with 2 axes containing the + # multiplicities of each path. + # num_repeats.num_elements() == unique_word_seqs.num_elements() + # + # Since k2.ragged.unique_sequences will reorder paths within a seq, + # `new2old` is a 1-D torch.Tensor mapping from the output path index + # to the input path index. + # new2old.numel() == unique_word_seqs.tot_size(1) + unique_word_seq, num_repeats, new2old = k2.ragged.unique_sequences( + word_seq, need_num_repeats=True, need_new2old_indexes=True + ) + + seq_to_path_shape = k2.ragged.get_layer(unique_word_seq.shape(), 0) + + # path_to_seq_map is a 1-D torch.Tensor. + # path_to_seq_map[i] is the seq to which the i-th path + # belongs. + path_to_seq_map = seq_to_path_shape.row_ids(1) + + # Remove the seq axis. + # Now unique_word_seq has only two axes [path][word] + unique_word_seq = k2.ragged.remove_axis(unique_word_seq, 0) + + # word_fsa is an FsaVec with axes [path][state][arc] + word_fsa = k2.linear_fsa(unique_word_seq) + + word_fsa_with_epsilon_loops = k2.add_epsilon_self_loops(word_fsa) + + am_scores, _ = compute_am_and_lm_scores( + lattice, word_fsa_with_epsilon_loops, path_to_seq_map + ) + + # Now compute lm_scores + b_to_a_map = torch.zeros_like(path_to_seq_map) + lm_path_lattice = _intersect_device( + G, + word_fsa_with_epsilon_loops, + b_to_a_map=b_to_a_map, + sorted_match_a=True, + ) + lm_path_lattice = k2.top_sort(k2.connect(lm_path_lattice)) + lm_scores = lm_path_lattice.get_tot_scores( + use_double_scores=True, log_semiring=False + ) + + path_2axes = k2.ragged.remove_axis(path, 0) + + ans = dict() + for lm_scale in lm_scale_list: + tot_scores = am_scores / lm_scale + lm_scores + + # Remember that we used `k2.ragged.unique_sequences` to remove repeated + # paths to avoid redundant computation in `k2.intersect_device`. + # Now we use `num_repeats` to correct the scores for each path. + # + # NOTE(fangjun): It is commented out as it leads to a worse WER + # tot_scores = tot_scores * num_repeats.values() + + ragged_tot_scores = k2.RaggedFloat( + seq_to_path_shape, tot_scores.to(torch.float32) + ) + argmax_indexes = k2.ragged.argmax_per_sublist(ragged_tot_scores) + + # Use k2.index here since argmax_indexes' dtype is torch.int32 + best_path_indexes = k2.index(new2old, argmax_indexes) + + # best_path is a k2.RaggedInt with 2 axes [path][arc_pos] + best_path = k2.index(path_2axes, best_path_indexes) + + # labels is a k2.RaggedInt with 2 axes [path][phone_id] + # Note that it contains -1s. + labels = k2.index(lattice.labels.contiguous(), best_path) + + labels = k2.ragged.remove_values_eq(labels, -1) + + # lattice.aux_labels is a k2.RaggedInt tensor with 2 axes, so + # aux_labels is also a k2.RaggedInt with 2 axes + aux_labels = k2.index(lattice.aux_labels, best_path.values()) + + best_path_fsa = k2.linear_fsa(labels) + best_path_fsa.aux_labels = aux_labels + + key = f"lm_scale_{lm_scale}" + ans[key] = best_path_fsa + + return ans + + +def rescore_with_whole_lattice( + lattice: k2.Fsa, + G_with_epsilon_loops: k2.Fsa, + lm_scale_list: Optional[List[float]] = None, +) -> Union[k2.Fsa, Dict[str, k2.Fsa]]: + """Use whole lattice to rescore. + + Args: + lattice: + An FsaVec It can be the return value of :func:`get_lattice`. + G_with_epsilon_loops: + An FsaVec representing the language model (LM). Note that it + is an FsaVec, but it contains only one Fsa. + lm_scale_list: + A list containing lm_scale values or None. + Returns: + If lm_scale_list is not None, return a dict of FsaVec, whose key + is a lm_scale and the value represents the best decoding path for + each sequence in the lattice. + If lm_scale_list is not None, return a lattice that is rescored + with the given LM. + """ + assert len(lattice.shape) == 3 + assert hasattr(lattice, "lm_scores") + assert G_with_epsilon_loops.shape == (1, None, None) + + device = lattice.device + lattice.scores = lattice.scores - lattice.lm_scores + # We will use lm_scores from G, so remove lats.lm_scores here + del lattice.lm_scores + assert hasattr(lattice, "lm_scores") is False + + # Now, lattice.scores contains only am_scores + + # inv_lattice has word IDs as labels. + # Its aux_labels are token IDs, which is a ragged tensor k2.RaggedInt + inv_lattice = k2.invert(lattice) + num_seqs = lattice.shape[0] + + b_to_a_map = torch.zeros(num_seqs, device=device, dtype=torch.int32) + while True: + try: + rescoring_lattice = k2.intersect_device( + G_with_epsilon_loops, + inv_lattice, + b_to_a_map, + sorted_match_a=True, + ) + rescoring_lattice = k2.top_sort(k2.connect(rescoring_lattice)) + break + except RuntimeError as e: + logging.info(f"Caught exception:\n{e}\n") + logging.info( + f"num_arcs before pruning: {inv_lattice.arcs.num_elements()}" + ) + + # NOTE(fangjun): The choice of the threshold 1e-7 is arbitrary here + # to avoid OOM. We may need to fine tune it. + inv_lattice = k2.prune_on_arc_post(inv_lattice, 1e-7, True) + logging.info( + f"num_arcs after pruning: {inv_lattice.arcs.num_elements()}" + ) + + # lat has token IDs as labels + # and word IDs as aux_labels. + lat = k2.invert(rescoring_lattice) + + if lm_scale_list is None: + return lat + + ans = dict() + # + # The following implements + # scores = (scores - lm_scores)/lm_scale + lm_scores + # = scores/lm_scale + lm_scores*(1 - 1/lm_scale) + # + saved_am_scores = lat.scores - lat.lm_scores + for lm_scale in lm_scale_list: + am_scores = saved_am_scores / lm_scale + lat.scores = am_scores + lat.lm_scores + + best_path = k2.shortest_path(lat, use_double_scores=True) + key = f"lm_scale_{lm_scale}" + ans[key] = best_path + return ans + + +def rescore_with_attention_decoder( + lattice: k2.Fsa, + num_paths: int, + model: nn.Module, + memory: torch.Tensor, + memory_key_padding_mask: torch.Tensor, +) -> Dict[str, k2.Fsa]: + """This function extracts n paths from the given lattice and uses + an attention decoder to rescore them. The path with the highest + score is used as the decoding output. + + lattice: + An FsaVec. It can be the return value of :func:`get_lattice`. + num_paths: + Number of paths to extract from the given lattice for rescoring. + model: + A transformer model. See the class "Transformer" in + conformer_ctc/transformer.py for its interface. + memory: + The encoder memory of the given model. It is the output of + the last torch.nn.TransformerEncoder layer in the given model. + Its shape is `[T, N, C]`. + memory_key_padding_mask: + The padding mask for memory with shape [N, T]. + Returns: + A dict of FsaVec, whose key contains a string + ngram_lm_scale_attention_scale and the value is the + best decoding path for each sequence in the lattice. + """ + # First, extract `num_paths` paths for each sequence. + # path is a k2.RaggedInt with axes [seq][path][arc_pos] + path = k2.random_paths(lattice, num_paths=num_paths, use_double_scores=True) + + # word_seq is a k2.RaggedInt sharing the same shape as `path` + # but it contains word IDs. Note that it also contains 0s and -1s. + # The last entry in each sublist is -1. + word_seq = k2.index(lattice.aux_labels, path) + + # Remove epsilons and -1 from word_seq + word_seq = k2.ragged.remove_values_leq(word_seq, 0) + + # Remove paths that has identical word sequences. + # + # unique_word_seq is still a k2.RaggedInt with 3 axes [seq][path][word] + # except that there are no repeated paths with the same word_seq + # within a sequence. + # + # num_repeats is also a k2.RaggedInt with 2 axes containing the + # multiplicities of each path. + # num_repeats.num_elements() == unique_word_seqs.num_elements() + # + # Since k2.ragged.unique_sequences will reorder paths within a seq, + # `new2old` is a 1-D torch.Tensor mapping from the output path index + # to the input path index. + # new2old.numel() == unique_word_seqs.tot_size(1) + unique_word_seq, num_repeats, new2old = k2.ragged.unique_sequences( + word_seq, need_num_repeats=True, need_new2old_indexes=True + ) + + seq_to_path_shape = k2.ragged.get_layer(unique_word_seq.shape(), 0) + + # path_to_seq_map is a 1-D torch.Tensor. + # path_to_seq_map[i] is the seq to which the i-th path + # belongs. + path_to_seq_map = seq_to_path_shape.row_ids(1) + + # Remove the seq axis. + # Now unique_word_seq has only two axes [path][word] + unique_word_seq = k2.ragged.remove_axis(unique_word_seq, 0) + + # word_fsa is an FsaVec with axes [path][state][arc] + word_fsa = k2.linear_fsa(unique_word_seq) + + word_fsa_with_epsilon_loops = k2.add_epsilon_self_loops(word_fsa) + + am_scores, ngram_lm_scores = compute_am_and_lm_scores( + lattice, word_fsa_with_epsilon_loops, path_to_seq_map + ) + # Now we use the attention decoder to compute another + # score: attention_scores. + # + # To do that, we have to get the input and output for the attention + # decoder. + + # CAUTION: The "tokens" attribute is set in the file + # local/compile_hlg.py + token_seq = k2.index(lattice.tokens, path) + + # Remove epsilons and -1 from token_seq + token_seq = k2.ragged.remove_values_leq(token_seq, 0) + + # Remove the seq axis. + token_seq = k2.ragged.remove_axis(token_seq, 0) + + token_seq, _ = k2.ragged.index( + token_seq, indexes=new2old, axis=0, need_value_indexes=False + ) + + # Now word in unique_word_seq has its corresponding token IDs. + token_ids = k2.ragged.to_list(token_seq) + + num_word_seqs = new2old.numel() + + path_to_seq_map_long = path_to_seq_map.to(torch.long) + expanded_memory = memory.index_select(1, path_to_seq_map_long) + + expanded_memory_key_padding_mask = memory_key_padding_mask.index_select( + 0, path_to_seq_map_long + ) + + # TODO: pass the sos_token_id and eos_token_id via function arguments + nll = model.decoder_nll( + expanded_memory, expanded_memory_key_padding_mask, token_ids, 1, 1 + ) + assert nll.ndim == 2 + assert nll.shape[0] == num_word_seqs + + attention_scores = -nll.sum(dim=1) + assert attention_scores.ndim == 1 + assert attention_scores.numel() == num_word_seqs + + ngram_lm_scale_list = [0.1, 0.3, 0.5, 0.6, 0.7, 0.9, 1.0] + ngram_lm_scale_list += [1.1, 1.2, 1.3, 1.5, 1.7, 1.9, 2.0] + + attention_scale_list = [0.1, 0.3, 0.5, 0.6, 0.7, 0.9, 1.0] + attention_scale_list += [1.1, 1.2, 1.3, 1.5, 1.7, 1.9, 2.0] + + path_2axes = k2.ragged.remove_axis(path, 0) + + ans = dict() + for n_scale in ngram_lm_scale_list: + for a_scale in attention_scale_list: + tot_scores = ( + am_scores + + n_scale * ngram_lm_scores + + a_scale * attention_scores + ) + ragged_tot_scores = k2.RaggedFloat(seq_to_path_shape, tot_scores) + argmax_indexes = k2.ragged.argmax_per_sublist(ragged_tot_scores) + + best_path_indexes = k2.index(new2old, argmax_indexes) + + # best_path is a k2.RaggedInt with 2 axes [path][arc_pos] + best_path = k2.index(path_2axes, best_path_indexes) + + # labels is a k2.RaggedInt with 2 axes [path][token_id] + # Note that it contains -1s. + labels = k2.index(lattice.labels.contiguous(), best_path) + + labels = k2.ragged.remove_values_eq(labels, -1) + + # lattice.aux_labels is a k2.RaggedInt tensor with 2 axes, so + # aux_labels is also a k2.RaggedInt with 2 axes + aux_labels = k2.index(lattice.aux_labels, best_path.values()) + + best_path_fsa = k2.linear_fsa(labels) + best_path_fsa.aux_labels = aux_labels + + key = f"ngram_lm_scale_{n_scale}_attention_scale_{a_scale}" + ans[key] = best_path_fsa + return ans diff --git a/icefall/dist.py b/icefall/dist.py new file mode 100644 index 0000000000..d314d2a436 --- /dev/null +++ b/icefall/dist.py @@ -0,0 +1,17 @@ +import os + +import torch +from torch import distributed as dist + + +def setup_dist(rank, world_size, master_port=None): + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = ( + "12354" if master_port is None else str(master_port) + ) + dist.init_process_group("nccl", rank=rank, world_size=world_size) + torch.cuda.set_device(rank) + + +def cleanup_dist(): + dist.destroy_process_group() diff --git a/icefall/graph_compiler.py b/icefall/graph_compiler.py new file mode 100644 index 0000000000..f7ba3cdaf4 --- /dev/null +++ b/icefall/graph_compiler.py @@ -0,0 +1,106 @@ +from typing import List + +import k2 +import torch + +from icefall.lexicon import Lexicon + + +class CtcTrainingGraphCompiler(object): + def __init__( + self, lexicon: Lexicon, device: torch.device, oov: str = "", + ): + """ + Args: + lexicon: + It is built from `data/lang/lexicon.txt`. + device: + The device to use for operations compiling transcripts to FSAs. + oov: + Out of vocabulary word. When a word in the transcript + does not exist in the lexicon, it is replaced with `oov`. + """ + L_inv = lexicon.L_inv.to(device) + assert L_inv.requires_grad is False + + assert oov in lexicon.word_table + + self.L_inv = k2.arc_sort(L_inv) + self.oov_id = lexicon.word_table[oov] + self.word_table = lexicon.word_table + + max_token_id = max(lexicon.tokens) + ctc_topo = k2.ctc_topo(max_token_id, modified=False) + + self.ctc_topo = ctc_topo.to(device) + self.device = device + + def compile(self, texts: List[str]) -> k2.Fsa: + """Build decoding graphs by composing ctc_topo with + given transcripts. + + Args: + texts: + A list of strings. Each string contains a sentence for an utterance. + A sentence consists of spaces separated words. An example `texts` + looks like: + + ['hello icefall', 'CTC training with k2'] + + Returns: + An FsaVec, the composition result of `self.ctc_topo` and the + transcript FSA. + """ + transcript_fsa = self.convert_transcript_to_fsa(texts) + + # NOTE: k2.compose runs on CUDA only when treat_epsilons_specially + # is False, so we add epsilon self-loops here + fsa_with_self_loops = k2.remove_epsilon_and_add_self_loops( + transcript_fsa + ) + + fsa_with_self_loops = k2.arc_sort(fsa_with_self_loops) + + decoding_graph = k2.compose( + self.ctc_topo, fsa_with_self_loops, treat_epsilons_specially=False + ) + + assert decoding_graph.requires_grad is False + + return decoding_graph + + def convert_transcript_to_fsa(self, texts: List[str]) -> k2.Fsa: + """Convert a list of transcript texts to an FsaVec. + + Args: + texts: + A list of strings. Each string contains a sentence for an utterance. + A sentence consists of spaces separated words. An example `texts` + looks like: + + ['hello icefall', 'CTC training with k2'] + + Returns: + Return an FsaVec, whose `shape[0]` equals to `len(texts)`. + """ + word_ids_list = [] + for text in texts: + word_ids = [] + for word in text.split(" "): + if word in self.word_table: + word_ids.append(self.word_table[word]) + else: + word_ids.append(self.oov_id) + word_ids_list.append(word_ids) + + word_fsa = k2.linear_fsa(word_ids_list, self.device) + + word_fsa_with_self_loops = k2.add_epsilon_self_loops(word_fsa) + + fsa = k2.intersect( + self.L_inv, word_fsa_with_self_loops, treat_epsilons_specially=False + ) + # fsa has word ID as labels and token ID as aux_labels, so + # we need to invert it + ans_fsa = fsa.invert_() + return k2.arc_sort(ans_fsa) diff --git a/icefall/lexicon.py b/icefall/lexicon.py new file mode 100644 index 0000000000..3b52c70c92 --- /dev/null +++ b/icefall/lexicon.py @@ -0,0 +1,193 @@ +import logging +import re +from pathlib import Path +from typing import List, Tuple, Union + +import k2 +import torch + + +def read_lexicon(filename: str) -> List[Tuple[str, List[str]]]: + """Read a lexicon from `filename`. + + Each line in the lexicon contains "word p1 p2 p3 ...". + That is, the first field is a word and the remaining + fields are tokens. Fields are separated by space(s). + + Args: + filename: + Path to the lexicon.txt + + Returns: + A list of tuples., e.g., [('w', ['p1', 'p2']), ('w1', ['p3, 'p4'])] + """ + ans = [] + + with open(filename, "r", encoding="utf-8") as f: + whitespace = re.compile("[ \t]+") + for line in f: + a = whitespace.split(line.strip(" \t\r\n")) + if len(a) == 0: + continue + + if len(a) < 2: + print(f"Found bad line {line} in lexicon file {filename}") + print("Every line is expected to contain at least 2 fields") + sys.exit(1) + word = a[0] + if word == "": + print(f"Found bad line {line} in lexicon file {filename}") + print(" should not be a valid word") + sys.exit(1) + + tokens = a[1:] + ans.append((word, tokens)) + + return ans + + +def write_lexicon(filename: str, lexicon: List[Tuple[str, List[str]]]) -> None: + """Write a lexicon to a file. + + Args: + filename: + Path to the lexicon file to be generated. + lexicon: + It can be the return value of :func:`read_lexicon`. + """ + with open(filename, "w", encoding="utf-8") as f: + for word, tokens in lexicon: + f.write(f"{word} {' '.join(tokens)}\n") + + +class Lexicon(object): + """Phone based lexicon. + + TODO: Add BpeLexicon for BPE models. + """ + + def __init__( + self, lang_dir: Path, disambig_pattern: str = re.compile(r"^#\d+$"), + ): + """ + Args: + lang_dir: + Path to the lang director. It is expected to contain the following + files: + - tokens.txt + - words.txt + - L.pt + The above files are produced by the script `prepare.sh`. You + should have run that before running the training code. + disambig_pattern: + It contains the pattern for disambiguation symbols. + """ + lang_dir = Path(lang_dir) + self.token_table = k2.SymbolTable.from_file(lang_dir / "tokens.txt") + self.word_table = k2.SymbolTable.from_file(lang_dir / "words.txt") + + if (lang_dir / "Linv.pt").exists(): + logging.info(f"Loading pre-compiled {lang_dir}/Linv.pt") + L_inv = k2.Fsa.from_dict(torch.load(lang_dir / "Linv.pt")) + else: + logging.info("Converting L.pt to Linv.pt") + L = k2.Fsa.from_dict(torch.load(lang_dir / "L.pt")) + L_inv = k2.arc_sort(L.invert()) + torch.save(L_inv.as_dict(), lang_dir / "Linv.pt") + + # We save L_inv instead of L because it will be used to intersect with + # transcript, both of whose labels are word IDs. + self.L_inv = L_inv + self.disambig_pattern = disambig_pattern + + @property + def tokens(self) -> List[int]: + """Return a list of token IDs excluding those from + disambiguation symbols. + + Caution: + 0 is not a token ID so it is excluded from the return value. + """ + symbols = self.token_table.symbols + ans = [] + for s in symbols: + if not self.disambig_pattern.match(s): + ans.append(self.token_table[s]) + if 0 in ans: + ans.remove(0) + ans.sort() + return ans + + +class BpeLexicon(Lexicon): + def __init__( + self, lang_dir: Path, disambig_pattern: str = re.compile(r"^#\d+$"), + ): + """ + Refer to the help information in Lexicon.__init__. + """ + super().__init__(lang_dir=lang_dir, disambig_pattern=disambig_pattern) + + self.ragged_lexicon = self.convert_lexicon_to_ragged( + lang_dir / "lexicon.txt" + ) + + def convert_lexicon_to_ragged(self, filename: str) -> k2.RaggedInt: + """Read a BPE lexicon from file and convert it to a + k2 ragged tensor. + + Args: + filename: + Filename of the BPE lexicon, e.g., data/lang/bpe/lexicon.txt + Returns: + A k2 ragged tensor with two axes [word_id] + """ + disambig_id = self.word_table["#0"] + # We reuse the same words.txt from the phone based lexicon + # so that we can share the same G.fst. Here, we have to + # exclude some words present only in the phone based lexicon. + excluded_words = ["", "!SIL", ""] + + # epsilon is not a word, but it occupies on position + # + row_splits = [0] + token_ids = [] + + lexicon = read_lexicon(filename) + lexicon = dict(lexicon) + + for i in range(disambig_id): + w = self.word_table[i] + if w in excluded_words: + row_splits.append(row_splits[-1]) + continue + pieces = lexicon[w] + piece_ids = [self.token_table[k] for k in pieces] + + row_splits.append(row_splits[-1] + len(piece_ids)) + token_ids.extend(piece_ids) + + cached_tot_size = row_splits[-1] + row_splits = torch.tensor(row_splits, dtype=torch.int32) + + shape = k2.ragged.create_ragged_shape2( + row_splits=row_splits, cached_tot_size=cached_tot_size + ) + values = torch.tensor(token_ids, dtype=torch.int32) + + return k2.RaggedInt(shape, values) + + def words_to_piece_ids(self, words: List[str]) -> k2.RaggedInt: + """Convert a list of words to a ragged tensor contained + word piece IDs. + """ + word_ids = [self.word_table[w] for w in words] + word_ids = torch.tensor(word_ids, dtype=torch.int32) + + ragged, _ = k2.ragged.index( + self.ragged_lexicon, + indexes=word_ids, + need_value_indexes=False, + axis=0, + ) + return ragged diff --git a/icefall/utils.py b/icefall/utils.py new file mode 100644 index 0000000000..1f2cf95f34 --- /dev/null +++ b/icefall/utils.py @@ -0,0 +1,383 @@ +import argparse +import logging +import os +import subprocess +from collections import defaultdict +from contextlib import contextmanager +from datetime import datetime +from pathlib import Path +from typing import Dict, Iterable, List, TextIO, Tuple, Union + +import k2 +import k2.ragged as k2r +import kaldialign +import torch +import torch.distributed as dist + +Pathlike = Union[str, Path] + + +@contextmanager +def get_executor(): + # We'll either return a process pool or a distributed worker pool. + # Note that this has to be a context manager because we might use multiple + # context manager ("with" clauses) inside, and this way everything will + # free up the resources at the right time. + try: + # If this is executed on the CLSP grid, we will try to use the + # Grid Engine to distribute the tasks. + # Other clusters can also benefit from that, provided a + # cluster-specific wrapper. + # (see https://github.com/pzelasko/plz for reference) + # + # The following must be installed: + # $ pip install dask distributed + # $ pip install git+https://github.com/pzelasko/plz + name = subprocess.check_output("hostname -f", shell=True, text=True) + if name.strip().endswith(".clsp.jhu.edu"): + import plz + from distributed import Client + + with plz.setup_cluster() as cluster: + cluster.scale(80) + yield Client(cluster) + return + except Exception: + pass + # No need to return anything - compute_and_store_features + # will just instantiate the pool itself. + yield None + + +def str2bool(v): + """Used in argparse.ArgumentParser.add_argument to indicate + that a type is a bool type and user can enter + + - yes, true, t, y, 1, to represent True + - no, false, f, n, 0, to represent False + + See https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse # noqa + """ + if isinstance(v, bool): + return v + if v.lower() in ("yes", "true", "t", "y", "1"): + return True + elif v.lower() in ("no", "false", "f", "n", "0"): + return False + else: + raise argparse.ArgumentTypeError("Boolean value expected.") + + +def setup_logger( + log_filename: Pathlike, log_level: str = "info", use_console: bool = True +) -> None: + """Setup log level. + + Args: + log_filename: + The filename to save the log. + log_level: + The log level to use, e.g., "debug", "info", "warning", "error", + "critical" + """ + now = datetime.now() + date_time = now.strftime("%Y-%m-%d-%H-%M-%S") + + if dist.is_available() and dist.is_initialized(): + world_size = dist.get_world_size() + rank = dist.get_rank() + formatter = f"%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] ({rank}/{world_size}) %(message)s" # noqa + log_filename = f"{log_filename}-{date_time}-{rank}" + else: + formatter = ( + "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s" + ) + log_filename = f"{log_filename}-{date_time}" + + os.makedirs(os.path.dirname(log_filename), exist_ok=True) + + level = logging.ERROR + if log_level == "debug": + level = logging.DEBUG + elif log_level == "info": + level = logging.INFO + elif log_level == "warning": + level = logging.WARNING + elif log_level == "critical": + level = logging.CRITICAL + + logging.basicConfig( + filename=log_filename, format=formatter, level=level, filemode="w" + ) + if use_console: + console = logging.StreamHandler() + console.setLevel(level) + console.setFormatter(logging.Formatter(formatter)) + logging.getLogger("").addHandler(console) + + +def get_env_info(): + """ + TODO: + """ + return { + "k2-git-sha1": None, + "k2-version": None, + "lhotse-version": None, + "torch-version": None, + "icefall-sha1": None, + "icefall-version": None, + } + + +# See +# https://stackoverflow.com/questions/4984647/accessing-dict-keys-like-an-attribute # noqa +class AttributeDict(dict): + __slots__ = () + __getattr__ = dict.__getitem__ + __setattr__ = dict.__setitem__ + + +def encode_supervisions( + supervisions: dict, subsampling_factor: int +) -> Tuple[torch.Tensor, List[str]]: + """ + Encodes Lhotse's ``batch["supervisions"]`` dict into a pair of torch Tensor, + and a list of transcription strings. + + The supervision tensor has shape ``(batch_size, 3)``. + Its second dimension contains information about sequence index [0], + start frames [1] and num frames [2]. + + The batch items might become re-ordered during this operation -- the + returned tensor and list of strings are guaranteed to be consistent with + each other. + """ + supervision_segments = torch.stack( + ( + supervisions["sequence_idx"], + supervisions["start_frame"] // subsampling_factor, + supervisions["num_frames"] // subsampling_factor, + ), + 1, + ).to(torch.int32) + + indices = torch.argsort(supervision_segments[:, 2], descending=True) + supervision_segments = supervision_segments[indices] + texts = supervisions["text"] + texts = [texts[idx] for idx in indices] + + return supervision_segments, texts + + +def get_texts(best_paths: k2.Fsa) -> List[List[int]]: + """Extract the texts (as word IDs) from the best-path FSAs. + Args: + best_paths: + A k2.Fsa with best_paths.arcs.num_axes() == 3, i.e. + containing multiple FSAs, which is expected to be the result + of k2.shortest_path (otherwise the returned values won't + be meaningful). + Returns: + Returns a list of lists of int, containing the label sequences we + decoded. + """ + if isinstance(best_paths.aux_labels, k2.RaggedInt): + # remove 0's and -1's. + aux_labels = k2r.remove_values_leq(best_paths.aux_labels, 0) + aux_shape = k2r.compose_ragged_shapes( + best_paths.arcs.shape(), aux_labels.shape() + ) + + # remove the states and arcs axes. + aux_shape = k2r.remove_axis(aux_shape, 1) + aux_shape = k2r.remove_axis(aux_shape, 1) + aux_labels = k2.RaggedInt(aux_shape, aux_labels.values()) + else: + # remove axis corresponding to states. + aux_shape = k2r.remove_axis(best_paths.arcs.shape(), 1) + aux_labels = k2.RaggedInt(aux_shape, best_paths.aux_labels) + # remove 0's and -1's. + aux_labels = k2r.remove_values_leq(aux_labels, 0) + + assert aux_labels.num_axes() == 2 + return k2r.to_list(aux_labels) + + +def store_transcripts( + filename: Pathlike, texts: Iterable[Tuple[str, str]] +) -> None: + """Save predicted results and reference transcripts to a file. + + Args: + filename: + File to save the results to. + texts: + An iterable of tuples. The first element is the reference transcript + while the second element is the predicted result. + Returns: + Return None. + """ + with open(filename, "w") as f: + for ref, hyp in texts: + print(f"ref={ref}", file=f) + print(f"hyp={hyp}", file=f) + + +def write_error_stats( + f: TextIO, test_set_name: str, results: List[Tuple[str, str]] +) -> float: + """Write statistics based on predicted results and reference transcripts. + + It will write the following to the given file: + + - WER + - number of insertions, deletions, substitutions, corrects and total + reference words. For example:: + + Errors: 23 insertions, 57 deletions, 212 substitutions, over 2606 + reference words (2337 correct) + + - The difference between the reference transcript and predicted results. + An instance is given below:: + + THE ASSOCIATION OF (EDISON->ADDISON) ILLUMINATING COMPANIES + + The above example shows that the reference word is `EDISON`, but it is + predicted to `ADDISON` (a substitution error). + + Another example is:: + + FOR THE FIRST DAY (SIR->*) I THINK + + The reference word `SIR` is missing in the predicted + results (a deletion error). + results: + An iterable of tuples. The first element is the reference transcript + while the second element is the predicted result. + Returns: + Return None. + """ + subs: Dict[Tuple[str, str], int] = defaultdict(int) + ins: Dict[str, int] = defaultdict(int) + dels: Dict[str, int] = defaultdict(int) + + # `words` stores counts per word, as follows: + # corr, ref_sub, hyp_sub, ins, dels + words: Dict[str, List[int]] = defaultdict(lambda: [0, 0, 0, 0, 0]) + num_corr = 0 + ERR = "*" + for ref, hyp in results: + ali = kaldialign.align(ref, hyp, ERR) + for ref_word, hyp_word in ali: + if ref_word == ERR: + ins[hyp_word] += 1 + words[hyp_word][3] += 1 + elif hyp_word == ERR: + dels[ref_word] += 1 + words[ref_word][4] += 1 + elif hyp_word != ref_word: + subs[(ref_word, hyp_word)] += 1 + words[ref_word][1] += 1 + words[hyp_word][2] += 1 + else: + words[ref_word][0] += 1 + num_corr += 1 + ref_len = sum([len(r) for r, _ in results]) + sub_errs = sum(subs.values()) + ins_errs = sum(ins.values()) + del_errs = sum(dels.values()) + tot_errs = sub_errs + ins_errs + del_errs + tot_err_rate = "%.2f" % (100.0 * tot_errs / ref_len) + + logging.info( + f"[{test_set_name}] %WER {tot_errs / ref_len:.2%} " + f"[{tot_errs} / {ref_len}, {ins_errs} ins, " + f"{del_errs} del, {sub_errs} sub ]" + ) + + print(f"%WER = {tot_err_rate}", file=f) + print( + f"Errors: {ins_errs} insertions, {del_errs} deletions, " + f"{sub_errs} substitutions, over {ref_len} reference " + f"words ({num_corr} correct)", + file=f, + ) + print( + "Search below for sections starting with PER-UTT DETAILS:, " + "SUBSTITUTIONS:, DELETIONS:, INSERTIONS:, PER-WORD STATS:", + file=f, + ) + + print("", file=f) + print("PER-UTT DETAILS: corr or (ref->hyp) ", file=f) + for ref, hyp in results: + ali = kaldialign.align(ref, hyp, ERR) + combine_successive_errors = True + if combine_successive_errors: + ali = [[[x], [y]] for x, y in ali] + for i in range(len(ali) - 1): + if ali[i][0] != ali[i][1] and ali[i + 1][0] != ali[i + 1][1]: + ali[i + 1][0] = ali[i][0] + ali[i + 1][0] + ali[i + 1][1] = ali[i][1] + ali[i + 1][1] + ali[i] = [[], []] + ali = [ + [ + list(filter(lambda a: a != ERR, x)), + list(filter(lambda a: a != ERR, y)), + ] + for x, y in ali + ] + ali = list(filter(lambda x: x != [[], []], ali)) + ali = [ + [ + ERR if x == [] else " ".join(x), + ERR if y == [] else " ".join(y), + ] + for x, y in ali + ] + + print( + " ".join( + ( + ref_word + if ref_word == hyp_word + else f"({ref_word}->{hyp_word})" + for ref_word, hyp_word in ali + ) + ), + file=f, + ) + + print("", file=f) + print("SUBSTITUTIONS: count ref -> hyp", file=f) + + for count, (ref, hyp) in sorted( + [(v, k) for k, v in subs.items()], reverse=True + ): + print(f"{count} {ref} -> {hyp}", file=f) + + print("", file=f) + print("DELETIONS: count ref", file=f) + for count, ref in sorted([(v, k) for k, v in dels.items()], reverse=True): + print(f"{count} {ref}", file=f) + + print("", file=f) + print("INSERTIONS: count hyp", file=f) + for count, hyp in sorted([(v, k) for k, v in ins.items()], reverse=True): + print(f"{count} {hyp}", file=f) + + print("", file=f) + print( + "PER-WORD STATS: word corr tot_errs count_in_ref count_in_hyp", file=f + ) + for _, word, counts in sorted( + [(sum(v[1:]), k, v) for k, v in words.items()], reverse=True + ): + (corr, ref_sub, hyp_sub, ins, dels) = counts + tot_errs = ref_sub + hyp_sub + ins + dels + ref_count = corr + ref_sub + dels + hyp_count = corr + hyp_sub + ins + + print(f"{word} {corr} {tot_errs} {ref_count} {hyp_count}", file=f) + return float(tot_err_rate) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000..0d80ed4d22 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,11 @@ +[tool.isort] +profile = "black" + +[tool.black] +line-length = 80 +exclude = ''' +/( + \.git + | \.github +)/ +''' diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000..a54edf118d --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +kaldilm +kaldialign +sentencepiece>=0.1.96 diff --git a/test/test_bpe_graph_compiler.py b/test/test_bpe_graph_compiler.py new file mode 100755 index 0000000000..7b941e5a7a --- /dev/null +++ b/test/test_bpe_graph_compiler.py @@ -0,0 +1,25 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2021 Xiaomi Corporation (authors: Fangjun Kuang) + +from icefall.bpe_graph_compiler import BpeCtcTrainingGraphCompiler +from icefall.lexicon import BpeLexicon +from pathlib import Path + + +def test(): + lang_dir = Path("data/lang/bpe") + if not lang_dir.is_dir(): + return + # TODO: generate data for testing + + compiler = BpeCtcTrainingGraphCompiler(lang_dir) + ids = compiler.texts_to_ids(["HELLO", "WORLD ZZZ"]) + fsa = compiler.compile(ids) + + lexicon = BpeLexicon(lang_dir) + ids0 = lexicon.words_to_piece_ids(["HELLO"]) + assert ids[0] == ids0.values().tolist() + + ids1 = lexicon.words_to_piece_ids(["WORLD", "ZZZ"]) + assert ids[1] == ids1.values().tolist() diff --git a/test/test_checkpoint.py b/test/test_checkpoint.py new file mode 100644 index 0000000000..343768957d --- /dev/null +++ b/test/test_checkpoint.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 + +import pytest +import torch +import torch.nn as nn + +from icefall.checkpoint import ( + average_checkpoints, + load_checkpoint, + save_checkpoint, +) + + +@pytest.fixture +def checkpoints1(tmp_path): + f = tmp_path / "f.pt" + m = nn.Module() + m.p1 = nn.Parameter(torch.tensor([10.0, 20.0]), requires_grad=False) + m.register_buffer("p2", torch.tensor([10, 100])) + + params = {"a": 10, "b": 20} + save_checkpoint(f, m, params=params) + return f + + +@pytest.fixture +def checkpoints2(tmp_path): + f = tmp_path / "f2.pt" + m = nn.Module() + m.p1 = nn.Parameter(torch.Tensor([50, 30.0])) + m.register_buffer("p2", torch.tensor([1, 3])) + params = {"a": 100, "b": 200} + + save_checkpoint(f, m, params=params) + return f + + +def test_load_checkpoints(checkpoints1): + m = nn.Module() + m.p1 = nn.Parameter(torch.Tensor([0, 0.0])) + m.p2 = nn.Parameter(torch.Tensor([0, 0])) + params = load_checkpoint(checkpoints1, m) + assert torch.allclose(m.p1, torch.Tensor([10.0, 20])) + assert params["a"] == 10 + assert params["b"] == 20 + + +def test_average_checkpoints(checkpoints1, checkpoints2): + state_dict = average_checkpoints([checkpoints1, checkpoints2]) + assert torch.allclose(state_dict["p1"], torch.Tensor([30, 25.0])) + assert torch.allclose(state_dict["p2"], torch.tensor([5, 51])) diff --git a/test/test_graph_compiler.py b/test/test_graph_compiler.py new file mode 100644 index 0000000000..4083d79ac5 --- /dev/null +++ b/test/test_graph_compiler.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2021 Xiaomi Corporation (authors: Fangjun Kuang) + +import re + +import k2 +import pytest +import torch + +from icefall.graph_compiler import CtcTrainingGraphCompiler +from icefall.lexicon import Lexicon +from icefall.utils import get_texts + + +@pytest.fixture +def lexicon(): + """ + We use the following test data: + + lexicon.txt + + foo f o o + bar b a r + baz b a z + SPN + + phones.txt + + 0 + a 1 + b 2 + f 3 + o 4 + r 5 + z 6 + SPN 7 + + words.txt: + + 0 + foo 1 + bar 2 + baz 3 + 4 + """ + L = k2.Fsa.from_str( + """ + 0 0 7 4 0 + 0 7 -1 -1 0 + 0 1 3 1 0 + 0 3 2 2 0 + 0 5 2 3 0 + 1 2 4 0 0 + 2 0 4 0 0 + 3 4 1 0 0 + 4 0 5 0 0 + 5 6 1 0 0 + 6 0 6 0 0 + 7 + """, + num_aux_labels=1, + ) + L.labels_sym = k2.SymbolTable.from_str( + """ + a 1 + b 2 + f 3 + o 4 + r 5 + z 6 + SPN 7 + """ + ) + L.aux_labels_sym = k2.SymbolTable.from_str( + """ + foo 1 + bar 2 + baz 3 + 4 + """ + ) + ans = Lexicon.__new__(Lexicon) + ans.token_table = L.labels_sym + ans.word_table = L.aux_labels_sym + ans.L_inv = k2.arc_sort(L.invert_()) + ans.disambig_pattern = re.compile(r"^#\d+$") + + return ans + + +@pytest.fixture +def compiler(lexicon): + return CtcTrainingGraphCompiler(lexicon, device=torch.device("cpu")) + + +class TestCtcTrainingGraphCompiler(object): + @staticmethod + def test_convert_transcript_to_fsa(compiler, lexicon): + texts = ["bar foo", "baz ok"] + fsa = compiler.convert_transcript_to_fsa(texts) + labels0 = fsa[0].labels[:-1].tolist() + aux_labels0 = fsa[0].aux_labels[:-1] + aux_labels0 = aux_labels0[aux_labels0 != 0].tolist() + + labels1 = fsa[1].labels[:-1].tolist() + aux_labels1 = fsa[1].aux_labels[:-1] + aux_labels1 = aux_labels1[aux_labels1 != 0].tolist() + + labels0 = [lexicon.token_table[i] for i in labels0] + labels1 = [lexicon.token_table[i] for i in labels1] + + aux_labels0 = [lexicon.word_table[i] for i in aux_labels0] + aux_labels1 = [lexicon.word_table[i] for i in aux_labels1] + + assert labels0 == ["b", "a", "r", "f", "o", "o"] + assert aux_labels0 == ["bar", "foo"] + + assert labels1 == ["b", "a", "z", "SPN"] + assert aux_labels1 == ["baz", ""] + + @staticmethod + def test_compile(compiler, lexicon): + texts = ["bar foo", "baz ok"] + decoding_graph = compiler.compile(texts) + input1 = ["b", "b", "", "", "a", "a", "r", "", ""] + input1 += ["f", "f", "", "", "o", "o", "", "o", "o"] + + input2 = ["b", "b", "a", "a", "a", "", "", "z", "z"] + input2 += ["", "", "SPN", "SPN", "", ""] + + lexicon.token_table._id2sym[0] == "" + lexicon.token_table._sym2id[""] = 0 + + input1 = [lexicon.token_table[i] for i in input1] + input2 = [lexicon.token_table[i] for i in input2] + + fsa1 = k2.linear_fsa(input1) + fsa2 = k2.linear_fsa(input2) + fsas = k2.Fsa.from_fsas([fsa1, fsa2]) + + decoding_graph = k2.arc_sort(decoding_graph) + lattice = k2.intersect( + decoding_graph, fsas, treat_epsilons_specially=False + ) + lattice = k2.connect(lattice) + + aux_labels0 = lattice[0].aux_labels[:-1] + aux_labels0 = aux_labels0[aux_labels0 != 0].tolist() + aux_labels0 = [lexicon.word_table[i] for i in aux_labels0] + assert aux_labels0 == ["bar", "foo"] + + aux_labels1 = lattice[1].aux_labels[:-1] + aux_labels1 = aux_labels1[aux_labels1 != 0].tolist() + aux_labels1 = [lexicon.word_table[i] for i in aux_labels1] + assert aux_labels1 == ["baz", ""] + + texts = get_texts(lattice) + texts = [[lexicon.word_table[i] for i in words] for words in texts] + assert texts == [["bar", "foo"], ["baz", ""]] diff --git a/test/test_lexicon.py b/test/test_lexicon.py new file mode 100644 index 0000000000..b1284d98aa --- /dev/null +++ b/test/test_lexicon.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 + +from pathlib import Path + +import k2 +import pytest +import torch + +from icefall.lexicon import BpeLexicon, Lexicon + + +@pytest.fixture +def lang_dir(tmp_path): + phone2id = """ + 0 + a 1 + b 2 + f 3 + o 4 + r 5 + z 6 + SPN 7 + #0 8 + """ + word2id = """ + 0 + foo 1 + bar 2 + baz 3 + 4 + #0 5 + """ + + L = k2.Fsa.from_str( + """ + 0 0 7 4 0 + 0 7 -1 -1 0 + 0 1 3 1 0 + 0 3 2 2 0 + 0 5 2 3 0 + 1 2 4 0 0 + 2 0 4 0 0 + 3 4 1 0 0 + 4 0 5 0 0 + 5 6 1 0 0 + 6 0 6 0 0 + 7 + """, + num_aux_labels=1, + ) + + with open(tmp_path / "tokens.txt", "w") as f: + f.write(phone2id) + with open(tmp_path / "words.txt", "w") as f: + f.write(word2id) + + torch.save(L.as_dict(), tmp_path / "L.pt") + + return tmp_path + + +def test_lexicon(lang_dir): + lexicon = Lexicon(lang_dir) + assert lexicon.tokens == list(range(1, 8)) + + +def test_bpe_lexicon(): + lang_dir = Path("data/lang/bpe") + if not lang_dir.is_dir(): + return + # TODO: Generate test data for BpeLexicon + + lexicon = BpeLexicon(lang_dir) + words = ["", "HELLO", "ZZZZ", "WORLD"] + ids = lexicon.words_to_piece_ids(words) + print(ids) + print([lexicon.token_table[i] for i in ids.values().tolist()]) diff --git a/test/test_utils.py b/test/test_utils.py new file mode 100644 index 0000000000..27b1ac203c --- /dev/null +++ b/test/test_utils.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python3 +import k2 +import pytest +import torch + +from icefall.utils import AttributeDict, encode_supervisions, get_texts + + +@pytest.fixture +def sup(): + sequence_idx = torch.tensor([0, 1, 2]) + start_frame = torch.tensor([1, 3, 9]) + num_frames = torch.tensor([20, 30, 10]) + text = ["one", "two", "three"] + return { + "sequence_idx": sequence_idx, + "start_frame": start_frame, + "num_frames": num_frames, + "text": text, + } + + +def test_encode_supervisions(sup): + supervision_segments, texts = encode_supervisions(sup, subsampling_factor=4) + assert torch.all( + torch.eq( + supervision_segments, + torch.tensor( + [[1, 0, 30 // 4], [0, 0, 20 // 4], [2, 9 // 4, 10 // 4]] + ), + ) + ) + assert texts == ["two", "one", "three"] + + +def test_get_texts_ragged(): + fsa1 = k2.Fsa.from_str( + """ + 0 1 1 10 + 1 2 2 20 + 2 3 3 30 + 3 4 -1 0 + 4 + """ + ) + fsa1.aux_labels = k2.RaggedInt("[ [1 3 0 2] [] [4 0 1] [-1]]") + + fsa2 = k2.Fsa.from_str( + """ + 0 1 1 1 + 1 2 2 2 + 2 3 -1 0 + 3 + """ + ) + fsa2.aux_labels = k2.RaggedInt("[[3 0 5 0 8] [0 9 7 0] [-1]]") + fsas = k2.Fsa.from_fsas([fsa1, fsa2]) + texts = get_texts(fsas) + assert texts == [[1, 3, 2, 4, 1], [3, 5, 8, 9, 7]] + + +def test_get_texts_regular(): + fsa1 = k2.Fsa.from_str( + """ + 0 1 1 3 10 + 1 2 2 0 20 + 2 3 3 2 30 + 3 4 -1 -1 0 + 4 + """, + num_aux_labels=1, + ) + + fsa2 = k2.Fsa.from_str( + """ + 0 1 1 10 1 + 1 2 2 5 2 + 2 3 -1 -1 0 + 3 + """, + num_aux_labels=1, + ) + fsas = k2.Fsa.from_fsas([fsa1, fsa2]) + texts = get_texts(fsas) + assert texts == [[3, 2], [10, 5]] + + +def test_attribute_dict(): + s = AttributeDict({"a": 10, "b": 20}) + assert s.a == 10 + assert s["b"] == 20 + s.c = 100 + assert s["c"] == 100