diff --git a/.flake8 b/.flake8
index a0f44263c..41d8799c8 100644
--- a/.flake8
+++ b/.flake8
@@ -1,7 +1,7 @@
[flake8]
show-source=true
statistics=true
-max-line-length = 80
+max-line-length = 88
per-file-ignores =
# line too long
icefall/diagnostics.py: E501,
@@ -12,6 +12,7 @@ per-file-ignores =
egs/librispeech/ASR/lstm_transducer_stateless*/*.py: E501, E203
egs/librispeech/ASR/conv_emformer_transducer_stateless*/*.py: E501, E203
egs/librispeech/ASR/conformer_ctc*/*py: E501,
+ egs/librispeech/ASR/zipformer_mmi/*.py: E501, E203
egs/librispeech/ASR/RESULTS.md: E999,
# invalid escape sequence (cause by tex formular), W605
diff --git a/.github/scripts/run-librispeech-conformer-ctc3-2022-11-28.sh b/.github/scripts/run-librispeech-conformer-ctc3-2022-11-28.sh
index 27944807f..df29f188e 100755
--- a/.github/scripts/run-librispeech-conformer-ctc3-2022-11-28.sh
+++ b/.github/scripts/run-librispeech-conformer-ctc3-2022-11-28.sh
@@ -13,7 +13,6 @@ cd egs/librispeech/ASR
repo_url=https://huggingface.co/Zengwei/icefall-asr-librispeech-conformer-ctc3-2022-11-27
log "Downloading pre-trained model from $repo_url"
-git lfs install
GIT_LFS_SKIP_SMUDGE=1 git clone $repo_url
repo=$(basename $repo_url)
@@ -23,7 +22,12 @@ soxi $repo/test_wavs/*.wav
ls -lh $repo/test_wavs/*.wav
pushd $repo/exp
-git lfs pull --include "data/*"
+git lfs pull --include "data/lang_bpe_500/HLG.pt"
+git lfs pull --include "data/lang_bpe_500/L.pt"
+git lfs pull --include "data/lang_bpe_500/LG.pt"
+git lfs pull --include "data/lang_bpe_500/Linv.pt"
+git lfs pull --include "data/lang_bpe_500/bpe.model"
+git lfs pull --include "data/lm/G_4_gram.pt"
git lfs pull --include "exp/jit_trace.pt"
git lfs pull --include "exp/pretrained.pt"
ln -s pretrained.pt epoch-99.pt
diff --git a/.github/scripts/run-librispeech-lstm-transducer-stateless2-2022-09-03.sh b/.github/scripts/run-librispeech-lstm-transducer-stateless2-2022-09-03.sh
index ac5b15979..9b883f889 100755
--- a/.github/scripts/run-librispeech-lstm-transducer-stateless2-2022-09-03.sh
+++ b/.github/scripts/run-librispeech-lstm-transducer-stateless2-2022-09-03.sh
@@ -193,7 +193,7 @@ if [[ x"${GITHUB_EVENT_LABEL_NAME}" == x"shallow-fusion" ]]; then
ls -lh data
ls -lh lstm_transducer_stateless2/exp
- log "Decoding test-clean and test-other"
+ log "Decoding test-clean and test-other with RNN LM"
./lstm_transducer_stateless2/decode.py \
--use-averaged-model 0 \
@@ -201,12 +201,14 @@ if [[ x"${GITHUB_EVENT_LABEL_NAME}" == x"shallow-fusion" ]]; then
--avg 1 \
--exp-dir lstm_transducer_stateless2/exp \
--max-duration 600 \
- --decoding-method modified_beam_search_rnnlm_shallow_fusion \
+ --decoding-method modified_beam_search_lm_shallow_fusion \
--beam 4 \
- --rnn-lm-scale 0.3 \
- --rnn-lm-exp-dir $lm_repo/exp \
- --rnn-lm-epoch 88 \
- --rnn-lm-avg 1 \
+ --use-shallow-fusion 1 \
+ --lm-type rnn \
+ --lm-exp-dir $lm_repo/exp \
+ --lm-epoch 88 \
+ --lm-avg 1 \
+ --lm-scale 0.3 \
--rnn-lm-num-layers 3 \
--rnn-lm-tie-weights 1
fi
@@ -245,11 +247,13 @@ if [[ x"${GITHUB_EVENT_LABEL_NAME}" == x"LODR" ]]; then
--avg 1 \
--exp-dir lstm_transducer_stateless2/exp \
--max-duration 600 \
- --decoding-method modified_beam_search_rnnlm_LODR \
+ --decoding-method modified_beam_search_LODR \
--beam 4 \
- --rnn-lm-scale 0.3 \
- --rnn-lm-exp-dir $lm_repo/exp \
- --rnn-lm-epoch 88 \
+ --use-shallow-fusion 1 \
+ --lm-type rnn \
+ --lm-exp-dir $lm_repo/exp \
+ --lm-scale 0.4 \
+ --lm-epoch 88 \
--rnn-lm-avg 1 \
--rnn-lm-num-layers 3 \
--rnn-lm-tie-weights 1 \
diff --git a/.github/scripts/run-librispeech-pruned-transducer-stateless7-2022-11-11.sh b/.github/scripts/run-librispeech-pruned-transducer-stateless7-2022-11-11.sh
index 8e485d2e6..999841b80 100755
--- a/.github/scripts/run-librispeech-pruned-transducer-stateless7-2022-11-11.sh
+++ b/.github/scripts/run-librispeech-pruned-transducer-stateless7-2022-11-11.sh
@@ -30,6 +30,15 @@ ln -s pretrained.pt epoch-99.pt
ls -lh *.pt
popd
+log "Test exporting to ONNX format"
+./pruned_transducer_stateless7/export.py \
+ --exp-dir $repo/exp \
+ --use-averaged-model false \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ --epoch 99 \
+ --avg 1 \
+ --onnx 1
+
log "Export to torchscript model"
./pruned_transducer_stateless7/export.py \
--exp-dir $repo/exp \
@@ -41,6 +50,27 @@ log "Export to torchscript model"
ls -lh $repo/exp/*.pt
+log "Decode with ONNX models"
+
+./pruned_transducer_stateless7/onnx_check.py \
+ --jit-filename $repo/exp/cpu_jit.pt \
+ --onnx-encoder-filename $repo/exp/encoder.onnx \
+ --onnx-decoder-filename $repo/exp/decoder.onnx \
+ --onnx-joiner-filename $repo/exp/joiner.onnx \
+ --onnx-joiner-encoder-proj-filename $repo/exp/joiner_encoder_proj.onnx \
+ --onnx-joiner-decoder-proj-filename $repo/exp/joiner_decoder_proj.onnx
+
+./pruned_transducer_stateless7/onnx_pretrained.py \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ --encoder-model-filename $repo/exp/encoder.onnx \
+ --decoder-model-filename $repo/exp/decoder.onnx \
+ --joiner-model-filename $repo/exp/joiner.onnx \
+ --joiner-encoder-proj-model-filename $repo/exp/joiner_encoder_proj.onnx \
+ --joiner-decoder-proj-model-filename $repo/exp/joiner_decoder_proj.onnx \
+ $repo/test_wavs/1089-134686-0001.wav \
+ $repo/test_wavs/1221-135766-0001.wav \
+ $repo/test_wavs/1221-135766-0002.wav
+
log "Decode with models exported by torch.jit.script()"
./pruned_transducer_stateless7/jit_pretrained.py \
diff --git a/.github/scripts/run-librispeech-pruned-transducer-stateless7-ctc-2022-12-01.sh b/.github/scripts/run-librispeech-pruned-transducer-stateless7-ctc-2022-12-01.sh
index 6642d5f67..3cbb480f6 100755
--- a/.github/scripts/run-librispeech-pruned-transducer-stateless7-ctc-2022-12-01.sh
+++ b/.github/scripts/run-librispeech-pruned-transducer-stateless7-ctc-2022-12-01.sh
@@ -13,7 +13,6 @@ cd egs/librispeech/ASR
repo_url=https://huggingface.co/Zengwei/icefall-asr-librispeech-pruned-transducer-stateless7-ctc-2022-12-01
log "Downloading pre-trained model from $repo_url"
-git lfs install
GIT_LFS_SKIP_SMUDGE=1 git clone $repo_url
repo=$(basename $repo_url)
@@ -23,7 +22,12 @@ soxi $repo/test_wavs/*.wav
ls -lh $repo/test_wavs/*.wav
pushd $repo/exp
-git lfs pull --include "data/*"
+git lfs pull --include "data/lang_bpe_500/HLG.pt"
+git lfs pull --include "data/lang_bpe_500/L.pt"
+git lfs pull --include "data/lang_bpe_500/LG.pt"
+git lfs pull --include "data/lang_bpe_500/Linv.pt"
+git lfs pull --include "data/lang_bpe_500/bpe.model"
+git lfs pull --include "data/lm/G_4_gram.pt"
git lfs pull --include "exp/cpu_jit.pt"
git lfs pull --include "exp/pretrained.pt"
ln -s pretrained.pt epoch-99.pt
@@ -144,4 +148,4 @@ if [[ x"${GITHUB_EVENT_NAME}" == x"schedule" || x"${GITHUB_EVENT_LABEL_NAME}" ==
done
rm pruned_transducer_stateless7_ctc/exp/*.pt
-fi
+fi
\ No newline at end of file
diff --git a/.github/scripts/run-librispeech-pruned-transducer-stateless7-ctc-bs-2022-12-15.sh b/.github/scripts/run-librispeech-pruned-transducer-stateless7-ctc-bs-2022-12-15.sh
new file mode 100755
index 000000000..ed66a728e
--- /dev/null
+++ b/.github/scripts/run-librispeech-pruned-transducer-stateless7-ctc-bs-2022-12-15.sh
@@ -0,0 +1,148 @@
+#!/usr/bin/env bash
+
+set -e
+
+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]}) $*"
+}
+
+cd egs/librispeech/ASR
+
+repo_url=https://huggingface.co/yfyeung/icefall-asr-librispeech-pruned_transducer_stateless7_ctc_bs-2022-12-14
+
+log "Downloading pre-trained model from $repo_url"
+GIT_LFS_SKIP_SMUDGE=1 git clone $repo_url
+repo=$(basename $repo_url)
+
+log "Display test files"
+tree $repo/
+soxi $repo/test_wavs/*.wav
+ls -lh $repo/test_wavs/*.wav
+
+pushd $repo/exp
+git lfs pull --include "data/lang_bpe_500/HLG.pt"
+git lfs pull --include "data/lang_bpe_500/L.pt"
+git lfs pull --include "data/lang_bpe_500/LG.pt"
+git lfs pull --include "data/lang_bpe_500/Linv.pt"
+git lfs pull --include "data/lang_bpe_500/bpe.model"
+git lfs pull --include "exp/cpu_jit.pt"
+git lfs pull --include "exp/pretrained.pt"
+ln -s pretrained.pt epoch-99.pt
+ls -lh *.pt
+popd
+
+log "Export to torchscript model"
+./pruned_transducer_stateless7_ctc_bs/export.py \
+ --exp-dir $repo/exp \
+ --use-averaged-model false \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ --epoch 99 \
+ --avg 1 \
+ --jit 1
+
+ls -lh $repo/exp/*.pt
+
+log "Decode with models exported by torch.jit.script()"
+
+./pruned_transducer_stateless7_ctc_bs/jit_pretrained.py \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ --nn-model-filename $repo/exp/cpu_jit.pt \
+ $repo/test_wavs/1089-134686-0001.wav \
+ $repo/test_wavs/1221-135766-0001.wav \
+ $repo/test_wavs/1221-135766-0002.wav
+
+for m in ctc-decoding 1best; do
+ ./pruned_transducer_stateless7_ctc_bs/jit_pretrained_ctc.py \
+ --model-filename $repo/exp/cpu_jit.pt \
+ --words-file $repo/data/lang_bpe_500/words.txt \
+ --HLG $repo/data/lang_bpe_500/HLG.pt \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ --method $m \
+ --sample-rate 16000 \
+ $repo/test_wavs/1089-134686-0001.wav \
+ $repo/test_wavs/1221-135766-0001.wav \
+ $repo/test_wavs/1221-135766-0002.wav
+done
+
+for sym in 1 2 3; do
+ log "Greedy search with --max-sym-per-frame $sym"
+
+ ./pruned_transducer_stateless7_ctc_bs/pretrained.py \
+ --method greedy_search \
+ --max-sym-per-frame $sym \
+ --checkpoint $repo/exp/pretrained.pt \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ $repo/test_wavs/1089-134686-0001.wav \
+ $repo/test_wavs/1221-135766-0001.wav \
+ $repo/test_wavs/1221-135766-0002.wav
+done
+
+for method in modified_beam_search beam_search fast_beam_search; do
+ log "$method"
+
+ ./pruned_transducer_stateless7_ctc_bs/pretrained.py \
+ --method $method \
+ --beam-size 4 \
+ --checkpoint $repo/exp/pretrained.pt \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ $repo/test_wavs/1089-134686-0001.wav \
+ $repo/test_wavs/1221-135766-0001.wav \
+ $repo/test_wavs/1221-135766-0002.wav
+done
+
+for m in ctc-decoding 1best; do
+ ./pruned_transducer_stateless7_ctc_bs/pretrained_ctc.py \
+ --checkpoint $repo/exp/pretrained.pt \
+ --words-file $repo/data/lang_bpe_500/words.txt \
+ --HLG $repo/data/lang_bpe_500/HLG.pt \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ --method $m \
+ --sample-rate 16000 \
+ $repo/test_wavs/1089-134686-0001.wav \
+ $repo/test_wavs/1221-135766-0001.wav \
+ $repo/test_wavs/1221-135766-0002.wav
+done
+
+echo "GITHUB_EVENT_NAME: ${GITHUB_EVENT_NAME}"
+echo "GITHUB_EVENT_LABEL_NAME: ${GITHUB_EVENT_LABEL_NAME}"
+
+if [[ x"${GITHUB_EVENT_NAME}" == x"schedule" || x"${GITHUB_EVENT_LABEL_NAME}" == x"run-decode" ]]; then
+ mkdir -p pruned_transducer_stateless7_ctc_bs/exp
+ ln -s $PWD/$repo/exp/pretrained.pt pruned_transducer_stateless7_ctc_bs/exp/epoch-999.pt
+ ln -s $PWD/$repo/data/lang_bpe_500 data/
+
+ ls -lh data
+ ls -lh pruned_transducer_stateless7_ctc_bs/exp
+
+ log "Decoding test-clean and test-other"
+
+ # use a small value for decoding with CPU
+ max_duration=100
+
+ for method in greedy_search fast_beam_search modified_beam_search; do
+ log "Decoding with $method"
+
+ ./pruned_transducer_stateless7_ctc_bs/decode.py \
+ --decoding-method $method \
+ --epoch 999 \
+ --avg 1 \
+ --use-averaged-model 0 \
+ --max-duration $max_duration \
+ --exp-dir pruned_transducer_stateless7_ctc_bs/exp
+ done
+
+ for m in ctc-decoding 1best; do
+ ./pruned_transducer_stateless7_ctc_bs/ctc_decode.py \
+ --epoch 999 \
+ --avg 1 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration $max_duration \
+ --use-averaged-model 0 \
+ --decoding-method $m \
+ --hlg-scale 0.6
+ done
+
+ rm pruned_transducer_stateless7_ctc_bs/exp/*.pt
+fi
diff --git a/.github/scripts/run-librispeech-pruned-transducer-stateless7-streaming-2022-12-29.sh b/.github/scripts/run-librispeech-pruned-transducer-stateless7-streaming-2022-12-29.sh
new file mode 100755
index 000000000..afb0dc05a
--- /dev/null
+++ b/.github/scripts/run-librispeech-pruned-transducer-stateless7-streaming-2022-12-29.sh
@@ -0,0 +1,148 @@
+#!/usr/bin/env bash
+
+set -e
+
+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]}) $*"
+}
+
+cd egs/librispeech/ASR
+
+repo_url=https://huggingface.co/Zengwei/icefall-asr-librispeech-pruned-transducer-stateless7-streaming-2022-12-29
+
+log "Downloading pre-trained model from $repo_url"
+git lfs install
+GIT_LFS_SKIP_SMUDGE=1 git clone $repo_url
+repo=$(basename $repo_url)
+
+log "Display test files"
+tree $repo/
+soxi $repo/test_wavs/*.wav
+ls -lh $repo/test_wavs/*.wav
+
+pushd $repo/exp
+git lfs pull --include "data/lang_bpe_500/bpe.model"
+git lfs pull --include "exp/cpu_jit.pt"
+git lfs pull --include "exp/pretrained.pt"
+git lfs pull --include "exp/encoder_jit_trace.pt"
+git lfs pull --include "exp/decoder_jit_trace.pt"
+git lfs pull --include "exp/joiner_jit_trace.pt"
+ln -s pretrained.pt epoch-99.pt
+ls -lh *.pt
+popd
+
+log "Export to torchscript model"
+./pruned_transducer_stateless7_streaming/export.py \
+ --exp-dir $repo/exp \
+ --use-averaged-model false \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ --decode-chunk-len 32 \
+ --epoch 99 \
+ --avg 1 \
+ --jit 1
+
+ls -lh $repo/exp/*.pt
+
+log "Decode with models exported by torch.jit.script()"
+
+./pruned_transducer_stateless7_streaming/jit_pretrained.py \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ --nn-model-filename $repo/exp/cpu_jit.pt \
+ --decode-chunk-len 32 \
+ $repo/test_wavs/1089-134686-0001.wav \
+ $repo/test_wavs/1221-135766-0001.wav \
+ $repo/test_wavs/1221-135766-0002.wav
+
+log "Export to torchscript model by torch.jit.trace()"
+./pruned_transducer_stateless7_streaming/jit_trace_export.py \
+ --exp-dir $repo/exp \
+ --use-averaged-model false \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ --decode-chunk-len 32 \
+ --epoch 99 \
+ --avg 1
+
+log "Decode with models exported by torch.jit.trace()"
+
+./pruned_transducer_stateless7_streaming/jit_trace_pretrained.py \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ --encoder-model-filename $repo/exp/encoder_jit_trace.pt \
+ --decoder-model-filename $repo/exp/decoder_jit_trace.pt \
+ --joiner-model-filename $repo/exp/joiner_jit_trace.pt \
+ --decode-chunk-len 32 \
+ $repo/test_wavs/1089-134686-0001.wav
+
+for sym in 1 2 3; do
+ log "Greedy search with --max-sym-per-frame $sym"
+
+ ./pruned_transducer_stateless7_streaming/pretrained.py \
+ --method greedy_search \
+ --max-sym-per-frame $sym \
+ --checkpoint $repo/exp/pretrained.pt \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ --decode-chunk-len 32 \
+ $repo/test_wavs/1089-134686-0001.wav \
+ $repo/test_wavs/1221-135766-0001.wav \
+ $repo/test_wavs/1221-135766-0002.wav
+done
+
+for method in modified_beam_search beam_search fast_beam_search; do
+ log "$method"
+
+ ./pruned_transducer_stateless7_streaming/pretrained.py \
+ --method $method \
+ --beam-size 4 \
+ --checkpoint $repo/exp/pretrained.pt \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ --decode-chunk-len 32 \
+ $repo/test_wavs/1089-134686-0001.wav \
+ $repo/test_wavs/1221-135766-0001.wav \
+ $repo/test_wavs/1221-135766-0002.wav
+done
+
+echo "GITHUB_EVENT_NAME: ${GITHUB_EVENT_NAME}"
+echo "GITHUB_EVENT_LABEL_NAME: ${GITHUB_EVENT_LABEL_NAME}"
+if [[ x"${GITHUB_EVENT_NAME}" == x"schedule" || x"${GITHUB_EVENT_LABEL_NAME}" == x"run-decode" ]]; then
+ mkdir -p pruned_transducer_stateless7_streaming/exp
+ ln -s $PWD/$repo/exp/pretrained.pt pruned_transducer_stateless7_streaming/exp/epoch-999.pt
+ ln -s $PWD/$repo/data/lang_bpe_500 data/
+
+ ls -lh data
+ ls -lh pruned_transducer_stateless7_streaming/exp
+
+ log "Decoding test-clean and test-other"
+
+ # use a small value for decoding with CPU
+ max_duration=100
+ num_decode_stream=200
+
+ for method in greedy_search fast_beam_search modified_beam_search; do
+ log "decoding with $method"
+
+ ./pruned_transducer_stateless7_streaming/decode.py \
+ --decoding-method $method \
+ --epoch 999 \
+ --avg 1 \
+ --use-averaged-model 0 \
+ --max-duration $max_duration \
+ --decode-chunk-len 32 \
+ --exp-dir pruned_transducer_stateless7_streaming/exp
+ done
+
+ for method in greedy_search fast_beam_search modified_beam_search; do
+ log "Decoding with $method"
+
+ ./pruned_transducer_stateless7_streaming/streaming_decode.py \
+ --decoding-method $method \
+ --epoch 999 \
+ --avg 1 \
+ --use-averaged-model 0 \
+ --decode-chunk-len 32 \
+ --num-decode-streams $num_decode_stream
+ --exp-dir pruned_transducer_stateless7_streaming/exp
+ done
+
+ rm pruned_transducer_stateless7_streaming/exp/*.pt
+fi
diff --git a/.github/scripts/run-librispeech-zipformer-mmi-2022-12-08.sh b/.github/scripts/run-librispeech-zipformer-mmi-2022-12-08.sh
new file mode 100755
index 000000000..77f28b054
--- /dev/null
+++ b/.github/scripts/run-librispeech-zipformer-mmi-2022-12-08.sh
@@ -0,0 +1,103 @@
+#!/usr/bin/env bash
+
+set -e
+
+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]}) $*"
+}
+
+cd egs/librispeech/ASR
+
+repo_url=https://huggingface.co/Zengwei/icefall-asr-librispeech-zipformer-mmi-2022-12-08
+
+log "Downloading pre-trained model from $repo_url"
+GIT_LFS_SKIP_SMUDGE=1 git clone $repo_url
+repo=$(basename $repo_url)
+
+log "Display test files"
+tree $repo/
+soxi $repo/test_wavs/*.wav
+ls -lh $repo/test_wavs/*.wav
+
+pushd $repo/exp
+git lfs pull --include "data/lang_bpe_500/3gram.pt"
+git lfs pull --include "data/lang_bpe_500/4gram.pt"
+git lfs pull --include "data/lang_bpe_500/L.pt"
+git lfs pull --include "data/lang_bpe_500/LG.pt"
+git lfs pull --include "data/lang_bpe_500/Linv.pt"
+git lfs pull --include "data/lang_bpe_500/bpe.model"
+git lfs pull --include "exp/cpu_jit.pt"
+git lfs pull --include "exp/pretrained.pt"
+ln -s pretrained.pt epoch-99.pt
+ls -lh *.pt
+popd
+
+log "Export to torchscript model"
+./zipformer_mmi/export.py \
+ --exp-dir $repo/exp \
+ --use-averaged-model false \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ --epoch 99 \
+ --avg 1 \
+ --jit 1
+
+ls -lh $repo/exp/*.pt
+
+log "Decode with models exported by torch.jit.script()"
+
+./zipformer_mmi/jit_pretrained.py \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ --nn-model-filename $repo/exp/cpu_jit.pt \
+ --lang-dir $repo/data/lang_bpe_500 \
+ $repo/test_wavs/1089-134686-0001.wav \
+ $repo/test_wavs/1221-135766-0001.wav \
+ $repo/test_wavs/1221-135766-0002.wav
+
+for method in 1best nbest nbest-rescoring-LG nbest-rescoring-3-gram nbest-rescoring-4-gram; do
+ log "$method"
+
+ ./zipformer_mmi/pretrained.py \
+ --method $method \
+ --checkpoint $repo/exp/pretrained.pt \
+ --lang-dir $repo/data/lang_bpe_500 \
+ --bpe-model $repo/data/lang_bpe_500/bpe.model \
+ $repo/test_wavs/1089-134686-0001.wav \
+ $repo/test_wavs/1221-135766-0001.wav \
+ $repo/test_wavs/1221-135766-0002.wav
+done
+
+
+echo "GITHUB_EVENT_NAME: ${GITHUB_EVENT_NAME}"
+echo "GITHUB_EVENT_LABEL_NAME: ${GITHUB_EVENT_LABEL_NAME}"
+if [[ x"${GITHUB_EVENT_NAME}" == x"schedule" || x"${GITHUB_EVENT_LABEL_NAME}" == x"run-decode" ]]; then
+ mkdir -p zipformer_mmi/exp
+ ln -s $PWD/$repo/exp/pretrained.pt zipformer_mmi/exp/epoch-999.pt
+ ln -s $PWD/$repo/data/lang_bpe_500 data/
+
+ ls -lh data
+ ls -lh zipformer_mmi/exp
+
+ log "Decoding test-clean and test-other"
+
+ # use a small value for decoding with CPU
+ max_duration=100
+
+ for method in 1best nbest nbest-rescoring-LG nbest-rescoring-3-gram nbest-rescoring-4-gram; do
+ log "Decoding with $method"
+
+ ./zipformer_mmi/decode.py \
+ --decoding-method $method \
+ --epoch 999 \
+ --avg 1 \
+ --use-averaged-model 0 \
+ --nbest-scale 1.2 \
+ --hp-scale 1.0 \
+ --max-duration $max_duration \
+ --lang-dir $repo/data/lang_bpe_500 \
+ --exp-dir zipformer_mmi/exp
+ done
+
+ rm zipformer_mmi/exp/*.pt
+fi
diff --git a/.github/workflows/run-librispeech-2022-11-11-stateless7.yml b/.github/workflows/run-librispeech-2022-11-11-stateless7.yml
index 365e2761a..7694e8bf5 100644
--- a/.github/workflows/run-librispeech-2022-11-11-stateless7.yml
+++ b/.github/workflows/run-librispeech-2022-11-11-stateless7.yml
@@ -39,7 +39,7 @@ concurrency:
jobs:
run_librispeech_2022_11_11_zipformer:
- if: github.event.label.name == 'ready' || github.event.label.name == 'run-decode' || github.event_name == 'push' || github.event_name == 'schedule'
+ if: github.event.label.name == 'onnx' || github.event.label.name == 'ready' || github.event.label.name == 'run-decode' || github.event_name == 'push' || github.event_name == 'schedule'
runs-on: ${{ matrix.os }}
strategy:
matrix:
diff --git a/.github/workflows/run-librispeech-2022-12-08-zipformer-mmi.yml b/.github/workflows/run-librispeech-2022-12-08-zipformer-mmi.yml
new file mode 100644
index 000000000..5472ca59b
--- /dev/null
+++ b/.github/workflows/run-librispeech-2022-12-08-zipformer-mmi.yml
@@ -0,0 +1,167 @@
+# Copyright 2022 Zengwei Yao
+
+# 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: run-librispeech-2022-12-08-zipformer-mmi
+# zipformer
+
+on:
+ push:
+ branches:
+ - master
+ pull_request:
+ types: [labeled]
+
+ schedule:
+ # minute (0-59)
+ # hour (0-23)
+ # day of the month (1-31)
+ # month (1-12)
+ # day of the week (0-6)
+ # nightly build at 15:50 UTC time every day
+ - cron: "50 15 * * *"
+
+concurrency:
+ group: run_librispeech_2022_12_08_zipformer-${{ github.ref }}
+ cancel-in-progress: true
+
+jobs:
+ run_librispeech_2022_12_08_zipformer:
+ if: github.event.label.name == 'ready' || github.event.label.name == 'run-decode' || github.event_name == 'push' || github.event_name == 'schedule'
+ runs-on: ${{ matrix.os }}
+ strategy:
+ matrix:
+ os: [ubuntu-latest]
+ python-version: [3.8]
+
+ fail-fast: false
+
+ steps:
+ - uses: actions/checkout@v2
+ with:
+ fetch-depth: 0
+
+ - name: Setup Python ${{ matrix.python-version }}
+ uses: actions/setup-python@v2
+ with:
+ python-version: ${{ matrix.python-version }}
+ cache: 'pip'
+ cache-dependency-path: '**/requirements-ci.txt'
+
+ - name: Install Python dependencies
+ run: |
+ grep -v '^#' ./requirements-ci.txt | xargs -n 1 -L 1 pip install
+ pip uninstall -y protobuf
+ pip install --no-binary protobuf protobuf
+
+ - name: Cache kaldifeat
+ id: my-cache
+ uses: actions/cache@v2
+ with:
+ path: |
+ ~/tmp/kaldifeat
+ key: cache-tmp-${{ matrix.python-version }}-2022-09-25
+
+ - name: Install kaldifeat
+ if: steps.my-cache.outputs.cache-hit != 'true'
+ shell: bash
+ run: |
+ .github/scripts/install-kaldifeat.sh
+
+ - name: Cache LibriSpeech test-clean and test-other datasets
+ id: libri-test-clean-and-test-other-data
+ uses: actions/cache@v2
+ with:
+ path: |
+ ~/tmp/download
+ key: cache-libri-test-clean-and-test-other
+
+ - name: Download LibriSpeech test-clean and test-other
+ if: steps.libri-test-clean-and-test-other-data.outputs.cache-hit != 'true'
+ shell: bash
+ run: |
+ .github/scripts/download-librispeech-test-clean-and-test-other-dataset.sh
+
+ - name: Prepare manifests for LibriSpeech test-clean and test-other
+ shell: bash
+ run: |
+ .github/scripts/prepare-librispeech-test-clean-and-test-other-manifests.sh
+
+ - name: Cache LibriSpeech test-clean and test-other fbank features
+ id: libri-test-clean-and-test-other-fbank
+ uses: actions/cache@v2
+ with:
+ path: |
+ ~/tmp/fbank-libri
+ key: cache-libri-fbank-test-clean-and-test-other-v2
+
+ - name: Compute fbank for LibriSpeech test-clean and test-other
+ if: steps.libri-test-clean-and-test-other-fbank.outputs.cache-hit != 'true'
+ shell: bash
+ run: |
+ .github/scripts/compute-fbank-librispeech-test-clean-and-test-other.sh
+
+ - name: Inference with pre-trained model
+ shell: bash
+ env:
+ GITHUB_EVENT_NAME: ${{ github.event_name }}
+ GITHUB_EVENT_LABEL_NAME: ${{ github.event.label.name }}
+ run: |
+ mkdir -p egs/librispeech/ASR/data
+ ln -sfv ~/tmp/fbank-libri egs/librispeech/ASR/data/fbank
+ ls -lh egs/librispeech/ASR/data/*
+
+ sudo apt-get -qq install git-lfs tree sox
+ export PYTHONPATH=$PWD:$PYTHONPATH
+ export PYTHONPATH=~/tmp/kaldifeat/kaldifeat/python:$PYTHONPATH
+ export PYTHONPATH=~/tmp/kaldifeat/build/lib:$PYTHONPATH
+
+ .github/scripts/run-librispeech-zipformer-mmi-2022-12-08.sh
+
+ - name: Display decoding results for librispeech zipformer-mmi
+ if: github.event_name == 'schedule' || github.event.label.name == 'run-decode'
+ shell: bash
+ run: |
+ cd egs/librispeech/ASR/
+ tree ./zipformer-mmi/exp
+
+ cd zipformer-mmi
+ echo "results for zipformer-mmi"
+ echo "===1best==="
+ find exp/1best -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/1best -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+ echo "===nbest==="
+ find exp/nbest -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/nbest -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+ echo "===nbest-rescoring-LG==="
+ find exp/nbest-rescoring-LG -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/nbest-rescoring-LG -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+ echo "===nbest-rescoring-3-gram==="
+ find exp/nbest-rescoring-3-gram -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/nbest-rescoring-3-gram -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+ echo "===nbest-rescoring-4-gram==="
+ find exp/nbest-rescoring-4-gram -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/nbest-rescoring-4-gram -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+ - name: Upload decoding results for librispeech zipformer-mmi
+ uses: actions/upload-artifact@v2
+ if: github.event_name == 'schedule' || github.event.label.name == 'run-decode'
+ with:
+ name: torch-${{ matrix.torch }}-python-${{ matrix.python-version }}-ubuntu-18.04-cpu-zipformer_mmi-2022-12-08
+ path: egs/librispeech/ASR/zipformer_mmi/exp/
diff --git a/.github/workflows/run-librispeech-2022-12-15-stateless7-ctc-bs.yml b/.github/workflows/run-librispeech-2022-12-15-stateless7-ctc-bs.yml
new file mode 100644
index 000000000..6e2b40cf3
--- /dev/null
+++ b/.github/workflows/run-librispeech-2022-12-15-stateless7-ctc-bs.yml
@@ -0,0 +1,163 @@
+# Copyright 2022 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: run-librispeech-2022-12-15-stateless7-ctc-bs
+# zipformer
+
+on:
+ push:
+ branches:
+ - master
+ pull_request:
+ types: [labeled]
+
+ schedule:
+ # minute (0-59)
+ # hour (0-23)
+ # day of the month (1-31)
+ # month (1-12)
+ # day of the week (0-6)
+ # nightly build at 15:50 UTC time every day
+ - cron: "50 15 * * *"
+
+jobs:
+ run_librispeech_2022_12_15_zipformer_ctc_bs:
+ if: github.event.label.name == 'ready' || github.event.label.name == 'run-decode' || github.event.label.name == 'blank-skip' || github.event_name == 'push' || github.event_name == 'schedule'
+ runs-on: ${{ matrix.os }}
+ strategy:
+ matrix:
+ os: [ubuntu-latest]
+ python-version: [3.8]
+
+ fail-fast: false
+
+ steps:
+ - uses: actions/checkout@v2
+ with:
+ fetch-depth: 0
+
+ - name: Setup Python ${{ matrix.python-version }}
+ uses: actions/setup-python@v2
+ with:
+ python-version: ${{ matrix.python-version }}
+ cache: 'pip'
+ cache-dependency-path: '**/requirements-ci.txt'
+
+ - name: Install Python dependencies
+ run: |
+ grep -v '^#' ./requirements-ci.txt | xargs -n 1 -L 1 pip install
+ pip uninstall -y protobuf
+ pip install --no-binary protobuf protobuf
+
+ - name: Cache kaldifeat
+ id: my-cache
+ uses: actions/cache@v2
+ with:
+ path: |
+ ~/tmp/kaldifeat
+ key: cache-tmp-${{ matrix.python-version }}-2022-09-25
+
+ - name: Install kaldifeat
+ if: steps.my-cache.outputs.cache-hit != 'true'
+ shell: bash
+ run: |
+ .github/scripts/install-kaldifeat.sh
+
+ - name: Cache LibriSpeech test-clean and test-other datasets
+ id: libri-test-clean-and-test-other-data
+ uses: actions/cache@v2
+ with:
+ path: |
+ ~/tmp/download
+ key: cache-libri-test-clean-and-test-other
+
+ - name: Download LibriSpeech test-clean and test-other
+ if: steps.libri-test-clean-and-test-other-data.outputs.cache-hit != 'true'
+ shell: bash
+ run: |
+ .github/scripts/download-librispeech-test-clean-and-test-other-dataset.sh
+
+ - name: Prepare manifests for LibriSpeech test-clean and test-other
+ shell: bash
+ run: |
+ .github/scripts/prepare-librispeech-test-clean-and-test-other-manifests.sh
+
+ - name: Cache LibriSpeech test-clean and test-other fbank features
+ id: libri-test-clean-and-test-other-fbank
+ uses: actions/cache@v2
+ with:
+ path: |
+ ~/tmp/fbank-libri
+ key: cache-libri-fbank-test-clean-and-test-other-v2
+
+ - name: Compute fbank for LibriSpeech test-clean and test-other
+ if: steps.libri-test-clean-and-test-other-fbank.outputs.cache-hit != 'true'
+ shell: bash
+ run: |
+ .github/scripts/compute-fbank-librispeech-test-clean-and-test-other.sh
+
+ - name: Inference with pre-trained model
+ shell: bash
+ env:
+ GITHUB_EVENT_NAME: ${{ github.event_name }}
+ GITHUB_EVENT_LABEL_NAME: ${{ github.event.label.name }}
+ run: |
+ mkdir -p egs/librispeech/ASR/data
+ ln -sfv ~/tmp/fbank-libri egs/librispeech/ASR/data/fbank
+ ls -lh egs/librispeech/ASR/data/*
+
+ sudo apt-get -qq install git-lfs tree sox
+ export PYTHONPATH=$PWD:$PYTHONPATH
+ export PYTHONPATH=~/tmp/kaldifeat/kaldifeat/python:$PYTHONPATH
+ export PYTHONPATH=~/tmp/kaldifeat/build/lib:$PYTHONPATH
+
+ .github/scripts/run-librispeech-pruned-transducer-stateless7-ctc-bs-2022-12-15.sh
+
+ - name: Display decoding results for librispeech pruned_transducer_stateless7_ctc_bs
+ if: github.event_name == 'schedule' || github.event.label.name == 'run-decode'
+ shell: bash
+ run: |
+ cd egs/librispeech/ASR/
+ tree ./pruned_transducer_stateless7_ctc_bs/exp
+
+ cd pruned_transducer_stateless7_ctc_bs
+ echo "results for pruned_transducer_stateless7_ctc_bs"
+ echo "===greedy search==="
+ find exp/greedy_search -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/greedy_search -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+ echo "===fast_beam_search==="
+ find exp/fast_beam_search -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/fast_beam_search -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+ echo "===modified beam search==="
+ find exp/modified_beam_search -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/modified_beam_search -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+ echo "===ctc decoding==="
+ find exp/ctc-decoding -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/ctc-decoding -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+ echo "===1best==="
+ find exp/1best -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/1best -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+ - name: Upload decoding results for librispeech pruned_transducer_stateless7_ctc_bs
+ uses: actions/upload-artifact@v2
+ if: github.event_name == 'schedule' || github.event.label.name == 'run-decode'
+ with:
+ name: torch-${{ matrix.torch }}-python-${{ matrix.python-version }}-ubuntu-18.04-cpu-pruned_transducer_stateless7-ctc-bs-2022-12-15
+ path: egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/exp/
diff --git a/.github/workflows/run-librispeech-2022-12-29-stateless7-streaming.yml b/.github/workflows/run-librispeech-2022-12-29-stateless7-streaming.yml
new file mode 100644
index 000000000..6dd93946a
--- /dev/null
+++ b/.github/workflows/run-librispeech-2022-12-29-stateless7-streaming.yml
@@ -0,0 +1,172 @@
+# Copyright 2022 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: run-librispeech-2022-12-29-stateless7-streaming
+# zipformer
+
+on:
+ push:
+ branches:
+ - master
+ pull_request:
+ types: [labeled]
+
+ schedule:
+ # minute (0-59)
+ # hour (0-23)
+ # day of the month (1-31)
+ # month (1-12)
+ # day of the week (0-6)
+ # nightly build at 15:50 UTC time every day
+ - cron: "50 15 * * *"
+
+concurrency:
+ group: run_librispeech_2022_12_29_zipformer_streaming-${{ github.ref }}
+ cancel-in-progress: true
+
+jobs:
+ run_librispeech_2022_12_29_zipformer_streaming:
+ if: github.event.label.name == 'ready' || github.event.label.name == 'run-decode' || github.event.label.name == 'streaming-zipformer' || github.event_name == 'push' || github.event_name == 'schedule'
+ runs-on: ${{ matrix.os }}
+ strategy:
+ matrix:
+ os: [ubuntu-latest]
+ python-version: [3.8]
+
+ fail-fast: false
+
+ steps:
+ - uses: actions/checkout@v2
+ with:
+ fetch-depth: 0
+
+ - name: Setup Python ${{ matrix.python-version }}
+ uses: actions/setup-python@v2
+ with:
+ python-version: ${{ matrix.python-version }}
+ cache: 'pip'
+ cache-dependency-path: '**/requirements-ci.txt'
+
+ - name: Install Python dependencies
+ run: |
+ grep -v '^#' ./requirements-ci.txt | xargs -n 1 -L 1 pip install
+ pip uninstall -y protobuf
+ pip install --no-binary protobuf protobuf
+
+ - name: Cache kaldifeat
+ id: my-cache
+ uses: actions/cache@v2
+ with:
+ path: |
+ ~/tmp/kaldifeat
+ key: cache-tmp-${{ matrix.python-version }}-2022-09-25
+
+ - name: Install kaldifeat
+ if: steps.my-cache.outputs.cache-hit != 'true'
+ shell: bash
+ run: |
+ .github/scripts/install-kaldifeat.sh
+
+ - name: Cache LibriSpeech test-clean and test-other datasets
+ id: libri-test-clean-and-test-other-data
+ uses: actions/cache@v2
+ with:
+ path: |
+ ~/tmp/download
+ key: cache-libri-test-clean-and-test-other
+
+ - name: Download LibriSpeech test-clean and test-other
+ if: steps.libri-test-clean-and-test-other-data.outputs.cache-hit != 'true'
+ shell: bash
+ run: |
+ .github/scripts/download-librispeech-test-clean-and-test-other-dataset.sh
+
+ - name: Prepare manifests for LibriSpeech test-clean and test-other
+ shell: bash
+ run: |
+ .github/scripts/prepare-librispeech-test-clean-and-test-other-manifests.sh
+
+ - name: Cache LibriSpeech test-clean and test-other fbank features
+ id: libri-test-clean-and-test-other-fbank
+ uses: actions/cache@v2
+ with:
+ path: |
+ ~/tmp/fbank-libri
+ key: cache-libri-fbank-test-clean-and-test-other-v2
+
+ - name: Compute fbank for LibriSpeech test-clean and test-other
+ if: steps.libri-test-clean-and-test-other-fbank.outputs.cache-hit != 'true'
+ shell: bash
+ run: |
+ .github/scripts/compute-fbank-librispeech-test-clean-and-test-other.sh
+
+ - name: Inference with pre-trained model
+ shell: bash
+ env:
+ GITHUB_EVENT_NAME: ${{ github.event_name }}
+ GITHUB_EVENT_LABEL_NAME: ${{ github.event.label.name }}
+ run: |
+ mkdir -p egs/librispeech/ASR/data
+ ln -sfv ~/tmp/fbank-libri egs/librispeech/ASR/data/fbank
+ ls -lh egs/librispeech/ASR/data/*
+
+ sudo apt-get -qq install git-lfs tree sox
+ export PYTHONPATH=$PWD:$PYTHONPATH
+ export PYTHONPATH=~/tmp/kaldifeat/kaldifeat/python:$PYTHONPATH
+ export PYTHONPATH=~/tmp/kaldifeat/build/lib:$PYTHONPATH
+
+ .github/scripts/run-librispeech-pruned-transducer-stateless7-streaming-2022-12-29.sh
+
+ - name: Display decoding results for librispeech pruned_transducer_stateless7_streaming
+ if: github.event_name == 'schedule' || github.event.label.name == 'run-decode'
+ shell: bash
+ run: |
+ cd egs/librispeech/ASR/
+ tree ./pruned_transducer_stateless7_streaming/exp
+
+ cd pruned_transducer_stateless7_streaming
+ echo "results for pruned_transducer_stateless7_streaming"
+ echo "===greedy search==="
+ find exp/greedy_search -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/greedy_search -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+ echo "===fast_beam_search==="
+ find exp/fast_beam_search -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/fast_beam_search -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+ echo "===modified beam search==="
+ find exp/modified_beam_search -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/modified_beam_search -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+ echo "===streaming greedy search==="
+ find exp/streaming/greedy_search -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/streaming/greedy_search -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+ echo "===streaming fast_beam_search==="
+ find exp/streaming/fast_beam_search -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/streaming/fast_beam_search -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+ echo "===streaming modified beam search==="
+ find exp/streaming/modified_beam_search -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find exp/streaming/modified_beam_search -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+
+ - name: Upload decoding results for librispeech pruned_transducer_stateless7_streaming
+ uses: actions/upload-artifact@v2
+ if: github.event_name == 'schedule' || github.event.label.name == 'run-decode'
+ with:
+ name: torch-${{ matrix.torch }}-python-${{ matrix.python-version }}-ubuntu-18.04-cpu-pruned_transducer_stateless7-streaming-2022-12-29
+ path: egs/librispeech/ASR/pruned_transducer_stateless7_streaming/exp/
diff --git a/.github/workflows/run-librispeech-lstm-transducer-stateless2-2022-09-03.yml b/.github/workflows/run-librispeech-lstm-transducer-stateless2-2022-09-03.yml
index f5ee09e16..3752f67e3 100644
--- a/.github/workflows/run-librispeech-lstm-transducer-stateless2-2022-09-03.yml
+++ b/.github/workflows/run-librispeech-lstm-transducer-stateless2-2022-09-03.yml
@@ -139,9 +139,10 @@ jobs:
cd egs/librispeech/ASR
tree lstm_transducer_stateless2/exp
cd lstm_transducer_stateless2/exp
- echo "===modified_beam_search_rnnlm_shallow_fusion==="
- find modified_beam_search_rnnlm_shallow_fusion -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
- find modified_beam_search_rnnlm_shallow_fusion -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+ echo "===modified_beam_search_lm_shallow_fusion==="
+ echo "===Using RNNLM==="
+ find modified_beam_search_lm_shallow_fusion -name "log-*rnn*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find modified_beam_search_lm_shallow_fusion -name "log-*rnn*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
- name: Display decoding results for lstm_transducer_stateless2
if: github.event.label.name == 'LODR'
@@ -151,8 +152,8 @@ jobs:
tree lstm_transducer_stateless2/exp
cd lstm_transducer_stateless2/exp
echo "===modified_beam_search_rnnlm_LODR==="
- find modified_beam_search_rnnlm_LODR -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
- find modified_beam_search_rnnlm_LODR -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+ find modified_beam_search_LODR -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+ find modified_beam_search_LODR -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
- name: Upload decoding results for lstm_transducer_stateless2
uses: actions/upload-artifact@v2
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 4dbe99827..c062a2a3d 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -113,6 +113,9 @@ jobs:
cd ../pruned_transducer_stateless4
pytest -v -s
+ cd ../pruned_transducer_stateless7
+ pytest -v -s
+
cd ../transducer_stateless
pytest -v -s
diff --git a/.gitignore b/.gitignore
index 583410f45..8af05d884 100644
--- a/.gitignore
+++ b/.gitignore
@@ -33,3 +33,4 @@ node_modules
*.param
*.bin
+.DS_Store
diff --git a/docs/README.md b/docs/README.md
new file mode 100644
index 000000000..3abb38f8b
--- /dev/null
+++ b/docs/README.md
@@ -0,0 +1,24 @@
+
+## Usage
+
+```bash
+cd /path/to/icefall/docs
+pip install -r requirements.txt
+make clean
+make html
+cd build/html
+python3 -m http.server 8000
+```
+
+It prints:
+
+```
+Serving HTTP on 0.0.0.0 port 8000 (http://0.0.0.0:8000/) ...
+```
+
+Open your browser and go to to view the generated
+documentation.
+
+Done!
+
+**Hint**: You can change the port number when starting the server.
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 221d9d734..ef9fe1445 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -78,3 +78,12 @@ html_context = {
}
todo_include_todos = True
+
+rst_epilog = """
+.. _sherpa-ncnn: https://github.com/k2-fsa/sherpa-ncnn
+.. _icefall: https://github.com/k2-fsa/icefall
+.. _git-lfs: https://git-lfs.com/
+.. _ncnn: https://github.com/tencent/ncnn
+.. _LibriSpeech: https://www.openslr.org/12
+.. _musan: http://www.openslr.org/17/
+"""
diff --git a/docs/source/faqs.rst b/docs/source/faqs.rst
new file mode 100644
index 000000000..72b0302d7
--- /dev/null
+++ b/docs/source/faqs.rst
@@ -0,0 +1,107 @@
+Frequently Asked Questions (FAQs)
+=================================
+
+In this section, we collect issues reported by users and post the corresponding
+solutions.
+
+
+OSError: libtorch_hip.so: cannot open shared object file: no such file or directory
+-----------------------------------------------------------------------------------
+
+One user is using the following code to install ``torch`` and ``torchaudio``:
+
+.. code-block:: bash
+
+ pip install \
+ torch==1.10.0+cu111 \
+ torchvision==0.11.0+cu111 \
+ torchaudio==0.10.0 \
+ -f https://download.pytorch.org/whl/torch_stable.html
+
+and it throws the following error when running ``tdnn/train.py``:
+
+.. code-block::
+
+ OSError: libtorch_hip.so: cannot open shared object file: no such file or directory
+
+The fix is to specify the CUDA version while installing ``torchaudio``. That
+is, change ``torchaudio==0.10.0`` to ``torchaudio==0.10.0+cu11```. Therefore,
+the correct command is:
+
+.. code-block:: bash
+
+ pip install \
+ torch==1.10.0+cu111 \
+ torchvision==0.11.0+cu111 \
+ torchaudio==0.10.0+cu111 \
+ -f https://download.pytorch.org/whl/torch_stable.html
+
+AttributeError: module 'distutils' has no attribute 'version'
+-------------------------------------------------------------
+
+The error log is:
+
+.. code-block::
+
+ Traceback (most recent call last):
+ File "./tdnn/train.py", line 14, in
+ from asr_datamodule import YesNoAsrDataModule
+ File "/home/xxx/code/next-gen-kaldi/icefall/egs/yesno/ASR/tdnn/asr_datamodule.py", line 34, in
+ from icefall.dataset.datamodule import DataModule
+ File "/home/xxx/code/next-gen-kaldi/icefall/icefall/__init__.py", line 3, in
+ from . import (
+ File "/home/xxx/code/next-gen-kaldi/icefall/icefall/decode.py", line 23, in
+ from icefall.utils import add_eos, add_sos, get_texts
+ File "/home/xxx/code/next-gen-kaldi/icefall/icefall/utils.py", line 39, in
+ from torch.utils.tensorboard import SummaryWriter
+ File "/home/xxx/tool/miniconda3/envs/yyy/lib/python3.8/site-packages/torch/utils/tensorboard/__init__.py", line 4, in
+ LooseVersion = distutils.version.LooseVersion
+ AttributeError: module 'distutils' has no attribute 'version'
+
+The fix is:
+
+.. code-block:: bash
+
+ pip uninstall setuptools
+
+ pip install setuptools==58.0.4
+
+ImportError: libpython3.10.so.1.0: cannot open shared object file: No such file or directory
+--------------------------------------------------------------------------------------------
+
+If you are using ``conda`` and encounter the following issue:
+
+.. code-block::
+
+ Traceback (most recent call last):
+ File "/k2-dev/yangyifan/anaconda3/envs/icefall/lib/python3.10/site-packages/k2-1.23.3.dev20230112+cuda11.6.torch1.13.1-py3.10-linux-x86_64.egg/k2/__init__.py", line 24, in
+ from _k2 import DeterminizeWeightPushingType
+ ImportError: libpython3.10.so.1.0: cannot open shared object file: No such file or directory
+
+ During handling of the above exception, another exception occurred:
+
+ Traceback (most recent call last):
+ File "/k2-dev/yangyifan/icefall/egs/librispeech/ASR/./pruned_transducer_stateless7_ctc_bs/decode.py", line 104, in
+ import k2
+ File "/k2-dev/yangyifan/anaconda3/envs/icefall/lib/python3.10/site-packages/k2-1.23.3.dev20230112+cuda11.6.torch1.13.1-py3.10-linux-x86_64.egg/k2/__init__.py", line 30, in
+ raise ImportError(
+ ImportError: libpython3.10.so.1.0: cannot open shared object file: No such file or directory
+ Note: If you're using anaconda and importing k2 on MacOS,
+ you can probably fix this by setting the environment variable:
+ export DYLD_LIBRARY_PATH=$CONDA_PREFIX/lib/python3.10/site-packages:$DYLD_LIBRARY_PATH
+
+Please first try to find where ``libpython3.10.so.1.0`` locates.
+
+For instance,
+
+.. code-block:: bash
+
+ cd $CONDA_PREFIX/lib
+ find . -name "libpython*"
+
+If you are able to find it inside ``$CODNA_PREFIX/lib``, please set the
+following environment variable:
+
+.. code-block:: bash
+
+ export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH
diff --git a/docs/source/index.rst b/docs/source/index.rst
index be9977ca9..8d76eb68b 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -21,7 +21,16 @@ speech recognition recipes using `k2 `_.
:caption: Contents:
installation/index
+ faqs
model-export/index
+
+.. toctree::
+ :maxdepth: 3
+
recipes/index
+
+.. toctree::
+ :maxdepth: 2
+
contributing/index
huggingface/index
diff --git a/docs/source/model-export/code/export-conv-emformer-transducer-for-ncnn-output.txt b/docs/source/model-export/code/export-conv-emformer-transducer-for-ncnn-output.txt
new file mode 100644
index 000000000..ecbdd4b31
--- /dev/null
+++ b/docs/source/model-export/code/export-conv-emformer-transducer-for-ncnn-output.txt
@@ -0,0 +1,21 @@
+2023-01-11 12:15:38,677 INFO [export-for-ncnn.py:220] device: cpu
+2023-01-11 12:15:38,681 INFO [export-for-ncnn.py:229] {'best_train_loss': inf, 'best_valid_loss': inf, 'best_train_epoch': -1, 'best_v
+alid_epoch': -1, 'batch_idx_train': 0, 'log_interval': 50, 'reset_interval': 200, 'valid_interval': 3000, 'feature_dim': 80, 'subsampl
+ing_factor': 4, 'decoder_dim': 512, 'joiner_dim': 512, 'model_warm_step': 3000, 'env_info': {'k2-version': '1.23.2', 'k2-build-type':
+'Release', 'k2-with-cuda': True, 'k2-git-sha1': 'a34171ed85605b0926eebbd0463d059431f4f74a', 'k2-git-date': 'Wed Dec 14 00:06:38 2022',
+ 'lhotse-version': '1.12.0.dev+missing.version.file', 'torch-version': '1.10.0+cu102', 'torch-cuda-available': False, 'torch-cuda-vers
+ion': '10.2', 'python-version': '3.8', 'icefall-git-branch': 'fix-stateless3-train-2022-12-27', 'icefall-git-sha1': '530e8a1-dirty', '
+icefall-git-date': 'Tue Dec 27 13:59:18 2022', 'icefall-path': '/star-fj/fangjun/open-source/icefall', 'k2-path': '/star-fj/fangjun/op
+en-source/k2/k2/python/k2/__init__.py', 'lhotse-path': '/star-fj/fangjun/open-source/lhotse/lhotse/__init__.py', 'hostname': 'de-74279
+-k2-train-3-1220120619-7695ff496b-s9n4w', 'IP address': '127.0.0.1'}, 'epoch': 30, 'iter': 0, 'avg': 1, 'exp_dir': PosixPath('icefa
+ll-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp'), 'bpe_model': './icefall-asr-librispeech-conv-emformer-transdu
+cer-stateless2-2022-07-05//data/lang_bpe_500/bpe.model', 'jit': False, 'context_size': 2, 'use_averaged_model': False, 'encoder_dim':
+512, 'nhead': 8, 'dim_feedforward': 2048, 'num_encoder_layers': 12, 'cnn_module_kernel': 31, 'left_context_length': 32, 'chunk_length'
+: 32, 'right_context_length': 8, 'memory_size': 32, 'blank_id': 0, 'vocab_size': 500}
+2023-01-11 12:15:38,681 INFO [export-for-ncnn.py:231] About to create model
+2023-01-11 12:15:40,053 INFO [checkpoint.py:112] Loading checkpoint from icefall-asr-librispeech-conv-emformer-transducer-stateless2-2
+022-07-05/exp/epoch-30.pt
+2023-01-11 12:15:40,708 INFO [export-for-ncnn.py:315] Number of model parameters: 75490012
+2023-01-11 12:15:41,681 INFO [export-for-ncnn.py:318] Using torch.jit.trace()
+2023-01-11 12:15:41,681 INFO [export-for-ncnn.py:320] Exporting encoder
+2023-01-11 12:15:41,682 INFO [export-for-ncnn.py:149] chunk_length: 32, right_context_length: 8
diff --git a/docs/source/model-export/code/generate-int-8-scale-table-for-conv-emformer.txt b/docs/source/model-export/code/generate-int-8-scale-table-for-conv-emformer.txt
new file mode 100644
index 000000000..347e7e51a
--- /dev/null
+++ b/docs/source/model-export/code/generate-int-8-scale-table-for-conv-emformer.txt
@@ -0,0 +1,104 @@
+Don't Use GPU. has_gpu: 0, config.use_vulkan_compute: 1
+num encoder conv layers: 88
+num joiner conv layers: 3
+num files: 3
+Processing ../test_wavs/1089-134686-0001.wav
+Processing ../test_wavs/1221-135766-0001.wav
+Processing ../test_wavs/1221-135766-0002.wav
+Processing ../test_wavs/1089-134686-0001.wav
+Processing ../test_wavs/1221-135766-0001.wav
+Processing ../test_wavs/1221-135766-0002.wav
+----------encoder----------
+conv_87 : max = 15.942385 threshold = 15.938493 scale = 7.968131
+conv_88 : max = 35.442448 threshold = 15.549335 scale = 8.167552
+conv_89 : max = 23.228289 threshold = 8.001738 scale = 15.871552
+linear_90 : max = 3.976146 threshold = 1.101789 scale = 115.267128
+linear_91 : max = 6.962030 threshold = 5.162033 scale = 24.602713
+linear_92 : max = 12.323041 threshold = 3.853959 scale = 32.953129
+linear_94 : max = 6.905416 threshold = 4.648006 scale = 27.323545
+linear_93 : max = 6.905416 threshold = 5.474093 scale = 23.200188
+linear_95 : max = 1.888012 threshold = 1.403563 scale = 90.483986
+linear_96 : max = 6.856741 threshold = 5.398679 scale = 23.524273
+linear_97 : max = 9.635942 threshold = 2.613655 scale = 48.590950
+linear_98 : max = 6.460340 threshold = 5.670146 scale = 22.398010
+linear_99 : max = 9.532276 threshold = 2.585537 scale = 49.119396
+linear_101 : max = 6.585871 threshold = 5.719224 scale = 22.205809
+linear_100 : max = 6.585871 threshold = 5.751382 scale = 22.081648
+linear_102 : max = 1.593344 threshold = 1.450581 scale = 87.551147
+linear_103 : max = 6.592681 threshold = 5.705824 scale = 22.257959
+linear_104 : max = 8.752957 threshold = 1.980955 scale = 64.110489
+linear_105 : max = 6.696240 threshold = 5.877193 scale = 21.608953
+linear_106 : max = 9.059659 threshold = 2.643138 scale = 48.048950
+linear_108 : max = 6.975461 threshold = 4.589567 scale = 27.671457
+linear_107 : max = 6.975461 threshold = 6.190381 scale = 20.515701
+linear_109 : max = 3.710759 threshold = 2.305635 scale = 55.082436
+linear_110 : max = 7.531228 threshold = 5.731162 scale = 22.159557
+linear_111 : max = 10.528083 threshold = 2.259322 scale = 56.211544
+linear_112 : max = 8.148807 threshold = 5.500842 scale = 23.087374
+linear_113 : max = 8.592566 threshold = 1.948851 scale = 65.166611
+linear_115 : max = 8.437109 threshold = 5.608947 scale = 22.642395
+linear_114 : max = 8.437109 threshold = 6.193942 scale = 20.503904
+linear_116 : max = 3.966980 threshold = 3.200896 scale = 39.676392
+linear_117 : max = 9.451303 threshold = 6.061664 scale = 20.951344
+linear_118 : max = 12.077262 threshold = 3.965800 scale = 32.023804
+linear_119 : max = 9.671615 threshold = 4.847613 scale = 26.198460
+linear_120 : max = 8.625638 threshold = 3.131427 scale = 40.556595
+linear_122 : max = 10.274080 threshold = 4.888716 scale = 25.978189
+linear_121 : max = 10.274080 threshold = 5.420480 scale = 23.429659
+linear_123 : max = 4.826197 threshold = 3.599617 scale = 35.281532
+linear_124 : max = 11.396383 threshold = 7.325849 scale = 17.335875
+linear_125 : max = 9.337198 threshold = 3.941410 scale = 32.221970
+linear_126 : max = 9.699965 threshold = 4.842878 scale = 26.224073
+linear_127 : max = 8.775370 threshold = 3.884215 scale = 32.696438
+linear_129 : max = 9.872276 threshold = 4.837319 scale = 26.254213
+linear_128 : max = 9.872276 threshold = 7.180057 scale = 17.687883
+linear_130 : max = 4.150427 threshold = 3.454298 scale = 36.765789
+linear_131 : max = 11.112692 threshold = 7.924847 scale = 16.025545
+linear_132 : max = 11.852893 threshold = 3.116593 scale = 40.749626
+linear_133 : max = 11.517084 threshold = 5.024665 scale = 25.275314
+linear_134 : max = 10.683807 threshold = 3.878618 scale = 32.743618
+linear_136 : max = 12.421055 threshold = 6.322729 scale = 20.086264
+linear_135 : max = 12.421055 threshold = 5.309880 scale = 23.917679
+linear_137 : max = 4.827781 threshold = 3.744595 scale = 33.915554
+linear_138 : max = 14.422395 threshold = 7.742882 scale = 16.402161
+linear_139 : max = 8.527538 threshold = 3.866123 scale = 32.849449
+linear_140 : max = 12.128619 threshold = 4.657793 scale = 27.266134
+linear_141 : max = 9.839593 threshold = 3.845993 scale = 33.021378
+linear_143 : max = 12.442304 threshold = 7.099039 scale = 17.889746
+linear_142 : max = 12.442304 threshold = 5.325038 scale = 23.849592
+linear_144 : max = 5.929444 threshold = 5.618206 scale = 22.605080
+linear_145 : max = 13.382126 threshold = 9.321095 scale = 13.625010
+linear_146 : max = 9.894987 threshold = 3.867645 scale = 32.836517
+linear_147 : max = 10.915313 threshold = 4.906028 scale = 25.886522
+linear_148 : max = 9.614287 threshold = 3.908151 scale = 32.496181
+linear_150 : max = 11.724932 threshold = 4.485588 scale = 28.312899
+linear_149 : max = 11.724932 threshold = 5.161146 scale = 24.606939
+linear_151 : max = 7.164453 threshold = 5.847355 scale = 21.719223
+linear_152 : max = 13.086471 threshold = 5.984121 scale = 21.222834
+linear_153 : max = 11.099524 threshold = 3.991601 scale = 31.816805
+linear_154 : max = 10.054585 threshold = 4.489706 scale = 28.286930
+linear_155 : max = 12.389185 threshold = 3.100321 scale = 40.963501
+linear_157 : max = 9.982999 threshold = 5.154796 scale = 24.637253
+linear_156 : max = 9.982999 threshold = 8.537706 scale = 14.875190
+linear_158 : max = 8.420287 threshold = 6.502287 scale = 19.531588
+linear_159 : max = 25.014746 threshold = 9.423280 scale = 13.477261
+linear_160 : max = 45.633553 threshold = 5.715335 scale = 22.220921
+linear_161 : max = 20.371849 threshold = 5.117830 scale = 24.815203
+linear_162 : max = 12.492933 threshold = 3.126283 scale = 40.623318
+linear_164 : max = 20.697504 threshold = 4.825712 scale = 26.317358
+linear_163 : max = 20.697504 threshold = 5.078367 scale = 25.008038
+linear_165 : max = 9.023975 threshold = 6.836278 scale = 18.577358
+linear_166 : max = 34.860619 threshold = 7.259792 scale = 17.493614
+linear_167 : max = 30.380934 threshold = 5.496160 scale = 23.107042
+linear_168 : max = 20.691216 threshold = 4.733317 scale = 26.831076
+linear_169 : max = 9.723948 threshold = 3.952728 scale = 32.129707
+linear_171 : max = 21.034811 threshold = 5.366547 scale = 23.665123
+linear_170 : max = 21.034811 threshold = 5.356277 scale = 23.710501
+linear_172 : max = 10.556884 threshold = 5.729481 scale = 22.166058
+linear_173 : max = 20.033039 threshold = 10.207264 scale = 12.442120
+linear_174 : max = 11.597379 threshold = 2.658676 scale = 47.768131
+----------joiner----------
+linear_2 : max = 19.293503 threshold = 14.305265 scale = 8.877850
+linear_1 : max = 10.812222 threshold = 8.766452 scale = 14.487047
+linear_3 : max = 0.999999 threshold = 0.999755 scale = 127.031174
+ncnn int8 calibration table create success, best wish for your int8 inference has a low accuracy loss...\(^0^)/...233...
diff --git a/docs/source/model-export/code/test-stremaing-ncnn-decode-conv-emformer-transducer-libri.txt b/docs/source/model-export/code/test-stremaing-ncnn-decode-conv-emformer-transducer-libri.txt
new file mode 100644
index 000000000..114fe7342
--- /dev/null
+++ b/docs/source/model-export/code/test-stremaing-ncnn-decode-conv-emformer-transducer-libri.txt
@@ -0,0 +1,7 @@
+2023-01-11 14:02:12,216 INFO [streaming-ncnn-decode.py:320] {'tokens': './icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/data/lang_bpe_500/tokens.txt', 'encoder_param_filename': './icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/encoder_jit_trace-pnnx.ncnn.param', 'encoder_bin_filename': './icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/encoder_jit_trace-pnnx.ncnn.bin', 'decoder_param_filename': './icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/decoder_jit_trace-pnnx.ncnn.param', 'decoder_bin_filename': './icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/decoder_jit_trace-pnnx.ncnn.bin', 'joiner_param_filename': './icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/joiner_jit_trace-pnnx.ncnn.param', 'joiner_bin_filename': './icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/joiner_jit_trace-pnnx.ncnn.bin', 'sound_filename': './icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/test_wavs/1089-134686-0001.wav'}
+T 51 32
+2023-01-11 14:02:13,141 INFO [streaming-ncnn-decode.py:328] Constructing Fbank computer
+2023-01-11 14:02:13,151 INFO [streaming-ncnn-decode.py:331] Reading sound files: ./icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/test_wavs/1089-134686-0001.wav
+2023-01-11 14:02:13,176 INFO [streaming-ncnn-decode.py:336] torch.Size([106000])
+2023-01-11 14:02:17,581 INFO [streaming-ncnn-decode.py:380] ./icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/test_wavs/1089-134686-0001.wav
+2023-01-11 14:02:17,581 INFO [streaming-ncnn-decode.py:381] AFTER EARLY NIGHTFALL THE YELLOW LAMPS WOULD LIGHT UP HERE AND THERE THE SQUALID QUARTER OF THE BROTHELS
diff --git a/docs/source/model-export/export-ncnn.rst b/docs/source/model-export/export-ncnn.rst
index 3dbb8b514..ed0264089 100644
--- a/docs/source/model-export/export-ncnn.rst
+++ b/docs/source/model-export/export-ncnn.rst
@@ -1,12 +1,771 @@
Export to ncnn
==============
-We support exporting LSTM transducer models to `ncnn `_.
-
-Please refer to :ref:`export-model-for-ncnn` for details.
+We support exporting both
+`LSTM transducer models `_
+and
+`ConvEmformer transducer models `_
+to `ncnn `_.
We also provide ``_
performing speech recognition using ``ncnn`` with exported models.
-It has been tested on Linux, macOS, Windows, and Raspberry Pi. The project is
-self-contained and can be statically linked to produce a binary containing
-everything needed.
+It has been tested on Linux, macOS, Windows, ``Android``, and ``Raspberry Pi``.
+
+`sherpa-ncnn`_ is self-contained and can be statically linked to produce
+a binary containing everything needed. Please refer
+to its documentation for details:
+
+ - ``_
+
+
+Export LSTM transducer models
+-----------------------------
+
+Please refer to :ref:`export-lstm-transducer-model-for-ncnn` for details.
+
+
+
+Export ConvEmformer transducer models
+-------------------------------------
+
+We use the pre-trained model from the following repository as an example:
+
+ - ``_
+
+We will show you step by step how to export it to `ncnn`_ and run it with `sherpa-ncnn`_.
+
+.. hint::
+
+ We use ``Ubuntu 18.04``, ``torch 1.10``, and ``Python 3.8`` for testing.
+
+.. caution::
+
+ Please use a more recent version of PyTorch. For instance, ``torch 1.8``
+ may ``not`` work.
+
+1. Download the pre-trained model
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. hint::
+
+ You can also refer to ``_ to download the pre-trained model.
+
+ You have to install `git-lfs`_ before you continue.
+
+.. code-block:: bash
+
+ cd egs/librispeech/ASR
+
+ GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/Zengwei/icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05
+ cd icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05
+
+ git lfs pull --include "exp/pretrained-epoch-30-avg-10-averaged.pt"
+ git lfs pull --include "data/lang_bpe_500/bpe.model"
+
+ cd ..
+
+.. note::
+
+ We download ``exp/pretrained-xxx.pt``, not ``exp/cpu-jit_xxx.pt``.
+
+
+In the above code, we download the pre-trained model into the directory
+``egs/librispeech/ASR/icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05``.
+
+2. Install ncnn and pnnx
+^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. code-block:: bash
+
+ # We put ncnn into $HOME/open-source/ncnn
+ # You can change it to anywhere you like
+
+ cd $HOME
+ mkdir -p open-source
+ cd open-source
+
+ git clone https://github.com/csukuangfj/ncnn
+ cd ncnn
+ git submodule update --recursive --init
+
+ # Note: We don't use "python setup.py install" or "pip install ." here
+
+ mkdir -p build-wheel
+ cd build-wheel
+
+ cmake \
+ -DCMAKE_BUILD_TYPE=Release \
+ -DNCNN_PYTHON=ON \
+ -DNCNN_BUILD_BENCHMARK=OFF \
+ -DNCNN_BUILD_EXAMPLES=OFF \
+ -DNCNN_BUILD_TOOLS=ON \
+ ..
+
+ make -j4
+
+ cd ..
+
+ # Note: $PWD here is $HOME/open-source/ncnn
+
+ export PYTHONPATH=$PWD/python:$PYTHONPATH
+ export PATH=$PWD/tools/pnnx/build/src:$PATH
+ export PATH=$PWD/build-wheel/tools/quantize:$PATH
+
+ # Now build pnnx
+ cd tools/pnnx
+ mkdir build
+ cd build
+ cmake ..
+ make -j4
+
+ ./src/pnnx
+
+Congratulations! You have successfully installed the following components:
+
+ - ``pnxx``, which is an executable located in
+ ``$HOME/open-source/ncnn/tools/pnnx/build/src``. We will use
+ it to convert models exported by ``torch.jit.trace()``.
+ - ``ncnn2int8``, which is an executable located in
+ ``$HOME/open-source/ncnn/build-wheel/tools/quantize``. We will use
+ it to quantize our models to ``int8``.
+ - ``ncnn.cpython-38-x86_64-linux-gnu.so``, which is a Python module located
+ in ``$HOME/open-source/ncnn/python/ncnn``.
+
+ .. note::
+
+ I am using ``Python 3.8``, so it
+ is ``ncnn.cpython-38-x86_64-linux-gnu.so``. If you use a different
+ version, say, ``Python 3.9``, the name would be
+ ``ncnn.cpython-39-x86_64-linux-gnu.so``.
+
+ Also, if you are not using Linux, the file name would also be different.
+ But that does not matter. As long as you can compile it, it should work.
+
+We have set up ``PYTHONPATH`` so that you can use ``import ncnn`` in your
+Python code. We have also set up ``PATH`` so that you can use
+``pnnx`` and ``ncnn2int8`` later in your terminal.
+
+.. caution::
+
+ Please don't use ``_.
+ We have made some modifications to the offical `ncnn`_.
+
+ We will synchronize ``_ periodically
+ with the official one.
+
+3. Export the model via torch.jit.trace()
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+First, let us rename our pre-trained model:
+
+.. code-block::
+
+ cd egs/librispeech/ASR
+
+ cd icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp
+
+ ln -s pretrained-epoch-30-avg-10-averaged.pt epoch-30.pt
+
+ cd ../..
+
+Next, we use the following code to export our model:
+
+.. code-block:: bash
+
+ dir=./icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/
+
+ ./conv_emformer_transducer_stateless2/export-for-ncnn.py \
+ --exp-dir $dir/exp \
+ --bpe-model $dir/data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 1 \
+ --use-averaged-model 0 \
+ \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --encoder-dim 512
+
+.. hint::
+
+ We have renamed our model to ``epoch-30.pt`` so that we can use ``--epoch 30``.
+ There is only one pre-trained model, so we use ``--avg 1 --use-averaged-model 0``.
+
+ If you have trained a model by yourself and if you have all checkpoints
+ available, please first use ``decode.py`` to tune ``--epoch --avg``
+ and select the best combination with with ``--use-averaged-model 1``.
+
+.. note::
+
+ You will see the following log output:
+
+ .. literalinclude:: ./code/export-conv-emformer-transducer-for-ncnn-output.txt
+
+ The log shows the model has ``75490012`` parameters, i.e., ``~75 M``.
+
+ .. code-block::
+
+ ls -lh icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/pretrained-epoch-30-avg-10-averaged.pt
+
+ -rw-r--r-- 1 kuangfangjun root 289M Jan 11 12:05 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/pretrained-epoch-30-avg-10-averaged.pt
+
+ You can see that the file size of the pre-trained model is ``289 MB``, which
+ is roughly ``75490012*4/1024/1024 = 287.97 MB``.
+
+After running ``conv_emformer_transducer_stateless2/export-for-ncnn.py``,
+we will get the following files:
+
+.. code-block:: bash
+
+ ls -lh icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/*pnnx*
+
+ -rw-r--r-- 1 kuangfangjun root 1010K Jan 11 12:15 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/decoder_jit_trace-pnnx.pt
+ -rw-r--r-- 1 kuangfangjun root 283M Jan 11 12:15 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/encoder_jit_trace-pnnx.pt
+ -rw-r--r-- 1 kuangfangjun root 3.0M Jan 11 12:15 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/joiner_jit_trace-pnnx.pt
+
+
+.. _conv-emformer-step-3-export-torchscript-model-via-pnnx:
+
+3. Export torchscript model via pnnx
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. hint::
+
+ Make sure you have set up the ``PATH`` environment variable. Otherwise,
+ it will throw an error saying that ``pnnx`` could not be found.
+
+Now, it's time to export our models to `ncnn`_ via ``pnnx``.
+
+.. code-block::
+
+ cd icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/
+
+ pnnx ./encoder_jit_trace-pnnx.pt
+ pnnx ./decoder_jit_trace-pnnx.pt
+ pnnx ./joiner_jit_trace-pnnx.pt
+
+It will generate the following files:
+
+.. code-block:: bash
+
+ ls -lh icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/*ncnn*{bin,param}
+
+ -rw-r--r-- 1 kuangfangjun root 503K Jan 11 12:38 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/decoder_jit_trace-pnnx.ncnn.bin
+ -rw-r--r-- 1 kuangfangjun root 437 Jan 11 12:38 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/decoder_jit_trace-pnnx.ncnn.param
+ -rw-r--r-- 1 kuangfangjun root 142M Jan 11 12:36 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/encoder_jit_trace-pnnx.ncnn.bin
+ -rw-r--r-- 1 kuangfangjun root 79K Jan 11 12:36 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/encoder_jit_trace-pnnx.ncnn.param
+ -rw-r--r-- 1 kuangfangjun root 1.5M Jan 11 12:38 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/joiner_jit_trace-pnnx.ncnn.bin
+ -rw-r--r-- 1 kuangfangjun root 488 Jan 11 12:38 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/joiner_jit_trace-pnnx.ncnn.param
+
+There are two types of files:
+
+- ``param``: It is a text file containing the model architectures. You can
+ use a text editor to view its content.
+- ``bin``: It is a binary file containing the model parameters.
+
+We compare the file sizes of the models below before and after converting via ``pnnx``:
+
+.. see https://tableconvert.com/restructuredtext-generator
+
++----------------------------------+------------+
+| File name | File size |
++==================================+============+
+| encoder_jit_trace-pnnx.pt | 283 MB |
++----------------------------------+------------+
+| decoder_jit_trace-pnnx.pt | 1010 KB |
++----------------------------------+------------+
+| joiner_jit_trace-pnnx.pt | 3.0 MB |
++----------------------------------+------------+
+| encoder_jit_trace-pnnx.ncnn.bin | 142 MB |
++----------------------------------+------------+
+| decoder_jit_trace-pnnx.ncnn.bin | 503 KB |
++----------------------------------+------------+
+| joiner_jit_trace-pnnx.ncnn.bin | 1.5 MB |
++----------------------------------+------------+
+
+You can see that the file sizes of the models after conversion are about one half
+of the models before conversion:
+
+ - encoder: 283 MB vs 142 MB
+ - decoder: 1010 KB vs 503 KB
+ - joiner: 3.0 MB vs 1.5 MB
+
+The reason is that by default ``pnnx`` converts ``float32`` parameters
+to ``float16``. A ``float32`` parameter occupies 4 bytes, while it is 2 bytes
+for ``float16``. Thus, it is ``twice smaller`` after conversion.
+
+.. hint::
+
+ If you use ``pnnx ./encoder_jit_trace-pnnx.pt fp16=0``, then ``pnnx``
+ won't convert ``float32`` to ``float16``.
+
+4. Test the exported models in icefall
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. note::
+
+ We assume you have set up the environment variable ``PYTHONPATH`` when
+ building `ncnn`_.
+
+Now we have successfully converted our pre-trained model to `ncnn`_ format.
+The generated 6 files are what we need. You can use the following code to
+test the converted models:
+
+.. code-block:: bash
+
+ ./conv_emformer_transducer_stateless2/streaming-ncnn-decode.py \
+ --tokens ./icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/data/lang_bpe_500/tokens.txt \
+ --encoder-param-filename ./icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/encoder_jit_trace-pnnx.ncnn.param \
+ --encoder-bin-filename ./icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/encoder_jit_trace-pnnx.ncnn.bin \
+ --decoder-param-filename ./icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/decoder_jit_trace-pnnx.ncnn.param \
+ --decoder-bin-filename ./icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/decoder_jit_trace-pnnx.ncnn.bin \
+ --joiner-param-filename ./icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/joiner_jit_trace-pnnx.ncnn.param \
+ --joiner-bin-filename ./icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/joiner_jit_trace-pnnx.ncnn.bin \
+ ./icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/test_wavs/1089-134686-0001.wav
+
+.. hint::
+
+ `ncnn`_ supports only ``batch size == 1``, so ``streaming-ncnn-decode.py`` accepts
+ only 1 wave file as input.
+
+The output is given below:
+
+.. literalinclude:: ./code/test-stremaing-ncnn-decode-conv-emformer-transducer-libri.txt
+
+Congratulations! You have successfully exported a model from PyTorch to `ncnn`_!
+
+
+.. _conv-emformer-modify-the-exported-encoder-for-sherpa-ncnn:
+
+5. Modify the exported encoder for sherpa-ncnn
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+In order to use the exported models in `sherpa-ncnn`_, we have to modify
+``encoder_jit_trace-pnnx.ncnn.param``.
+
+Let us have a look at the first few lines of ``encoder_jit_trace-pnnx.ncnn.param``:
+
+.. code-block::
+
+ 7767517
+ 1060 1342
+ Input in0 0 1 in0
+
+**Explanation** of the above three lines:
+
+ 1. ``7767517``, it is a magic number and should not be changed.
+ 2. ``1060 1342``, the first number ``1060`` specifies the number of layers
+ in this file, while ``1342`` specifies the number of intermediate outputs
+ of this file
+ 3. ``Input in0 0 1 in0``, ``Input`` is the layer type of this layer; ``in0``
+ is the layer name of this layer; ``0`` means this layer has no input;
+ ``1`` means this layer has one output; ``in0`` is the output name of
+ this layer.
+
+We need to add 1 extra line and also increment the number of layers.
+The result looks like below:
+
+.. code-block:: bash
+
+ 7767517
+ 1061 1342
+ SherpaMetaData sherpa_meta_data1 0 0 0=1 1=12 2=32 3=31 4=8 5=32 6=8 7=512
+ Input in0 0 1 in0
+
+**Explanation**
+
+ 1. ``7767517``, it is still the same
+ 2. ``1061 1342``, we have added an extra layer, so we need to update ``1060`` to ``1061``.
+ We don't need to change ``1342`` since the newly added layer has no inputs or outputs.
+ 3. ``SherpaMetaData sherpa_meta_data1 0 0 0=1 1=12 2=32 3=31 4=8 5=32 6=8 7=512``
+ This line is newly added. Its explanation is given below:
+
+ - ``SherpaMetaData`` is the type of this layer. Must be ``SherpaMetaData``.
+ - ``sherpa_meta_data1`` is the name of this layer. Must be ``sherpa_meta_data1``.
+ - ``0 0`` means this layer has no inputs or output. Must be ``0 0``
+ - ``0=1``, 0 is the key and 1 is the value. MUST be ``0=1``
+ - ``1=12``, 1 is the key and 12 is the value of the
+ parameter ``--num-encoder-layers`` that you provided when running
+ ``conv_emformer_transducer_stateless2/export-for-ncnn.py``.
+ - ``2=32``, 2 is the key and 32 is the value of the
+ parameter ``--memory-size`` that you provided when running
+ ``conv_emformer_transducer_stateless2/export-for-ncnn.py``.
+ - ``3=31``, 3 is the key and 31 is the value of the
+ parameter ``--cnn-module-kernel`` that you provided when running
+ ``conv_emformer_transducer_stateless2/export-for-ncnn.py``.
+ - ``4=8``, 4 is the key and 8 is the value of the
+ parameter ``--left-context-length`` that you provided when running
+ ``conv_emformer_transducer_stateless2/export-for-ncnn.py``.
+ - ``5=32``, 5 is the key and 32 is the value of the
+ parameter ``--chunk-length`` that you provided when running
+ ``conv_emformer_transducer_stateless2/export-for-ncnn.py``.
+ - ``6=8``, 6 is the key and 8 is the value of the
+ parameter ``--right-context-length`` that you provided when running
+ ``conv_emformer_transducer_stateless2/export-for-ncnn.py``.
+ - ``7=512``, 7 is the key and 512 is the value of the
+ parameter ``--encoder-dim`` that you provided when running
+ ``conv_emformer_transducer_stateless2/export-for-ncnn.py``.
+
+ For ease of reference, we list the key-value pairs that you need to add
+ in the following table. If your model has a different setting, please
+ change the values for ``SherpaMetaData`` accordingly. Otherwise, you
+ will be ``SAD``.
+
+ +------+-----------------------------+
+ | key | value |
+ +======+=============================+
+ | 0 | 1 (fixed) |
+ +------+-----------------------------+
+ | 1 | ``--num-encoder-layers`` |
+ +------+-----------------------------+
+ | 2 | ``--memory-size`` |
+ +------+-----------------------------+
+ | 3 | ``--cnn-module-kernel`` |
+ +------+-----------------------------+
+ | 4 | ``--left-context-length`` |
+ +------+-----------------------------+
+ | 5 | ``--chunk-length`` |
+ +------+-----------------------------+
+ | 6 | ``--right-context-length`` |
+ +------+-----------------------------+
+ | 7 | ``--encoder-dim`` |
+ +------+-----------------------------+
+
+ 4. ``Input in0 0 1 in0``. No need to change it.
+
+.. caution::
+
+ When you add a new layer ``SherpaMetaData``, please remember to update the
+ number of layers. In our case, update ``1060`` to ``1061``. Otherwise,
+ you will be SAD later.
+
+.. hint::
+
+ After adding the new layer ``SherpaMetaData``, you cannot use this model
+ with ``streaming-ncnn-decode.py`` anymore since ``SherpaMetaData`` is
+ supported only in `sherpa-ncnn`_.
+
+.. hint::
+
+ `ncnn`_ is very flexible. You can add new layers to it just by text-editing
+ the ``param`` file! You don't need to change the ``bin`` file.
+
+Now you can use this model in `sherpa-ncnn`_.
+Please refer to the following documentation:
+
+ - Linux/macOS/Windows/arm/aarch64: ``_
+ - Android: ``_
+ - Python: ``_
+
+We have a list of pre-trained models that have been exported for `sherpa-ncnn`_:
+
+ - ``_
+
+ You can find more usages there.
+
+6. (Optional) int8 quantization with sherpa-ncnn
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+This step is optional.
+
+In this step, we describe how to quantize our model with ``int8``.
+
+Change :ref:`conv-emformer-step-3-export-torchscript-model-via-pnnx` to
+disable ``fp16`` when using ``pnnx``:
+
+.. code-block::
+
+ cd icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/
+
+ pnnx ./encoder_jit_trace-pnnx.pt fp16=0
+ pnnx ./decoder_jit_trace-pnnx.pt
+ pnnx ./joiner_jit_trace-pnnx.pt fp16=0
+
+.. note::
+
+ We add ``fp16=0`` when exporting the encoder and joiner. `ncnn`_ does not
+ support quantizing the decoder model yet. We will update this documentation
+ once `ncnn`_ supports it. (Maybe in this year, 2023).
+
+It will generate the following files
+
+.. code-block:: bash
+
+ ls -lh icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/*_jit_trace-pnnx.ncnn.{param,bin}
+
+ -rw-r--r-- 1 kuangfangjun root 503K Jan 11 15:56 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/decoder_jit_trace-pnnx.ncnn.bin
+ -rw-r--r-- 1 kuangfangjun root 437 Jan 11 15:56 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/decoder_jit_trace-pnnx.ncnn.param
+ -rw-r--r-- 1 kuangfangjun root 283M Jan 11 15:56 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/encoder_jit_trace-pnnx.ncnn.bin
+ -rw-r--r-- 1 kuangfangjun root 79K Jan 11 15:56 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/encoder_jit_trace-pnnx.ncnn.param
+ -rw-r--r-- 1 kuangfangjun root 3.0M Jan 11 15:56 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/joiner_jit_trace-pnnx.ncnn.bin
+ -rw-r--r-- 1 kuangfangjun root 488 Jan 11 15:56 icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/joiner_jit_trace-pnnx.ncnn.param
+
+Let us compare again the file sizes:
+
++----------------------------------------+------------+
+| File name | File size |
++----------------------------------------+------------+
+| encoder_jit_trace-pnnx.pt | 283 MB |
++----------------------------------------+------------+
+| decoder_jit_trace-pnnx.pt | 1010 KB |
++----------------------------------------+------------+
+| joiner_jit_trace-pnnx.pt | 3.0 MB |
++----------------------------------------+------------+
+| encoder_jit_trace-pnnx.ncnn.bin (fp16) | 142 MB |
++----------------------------------------+------------+
+| decoder_jit_trace-pnnx.ncnn.bin (fp16) | 503 KB |
++----------------------------------------+------------+
+| joiner_jit_trace-pnnx.ncnn.bin (fp16) | 1.5 MB |
++----------------------------------------+------------+
+| encoder_jit_trace-pnnx.ncnn.bin (fp32) | 283 MB |
++----------------------------------------+------------+
+| joiner_jit_trace-pnnx.ncnn.bin (fp32) | 3.0 MB |
++----------------------------------------+------------+
+
+You can see that the file sizes are doubled when we disable ``fp16``.
+
+.. note::
+
+ You can again use ``streaming-ncnn-decode.py`` to test the exported models.
+
+Next, follow :ref:`conv-emformer-modify-the-exported-encoder-for-sherpa-ncnn`
+to modify ``encoder_jit_trace-pnnx.ncnn.param``.
+
+Change
+
+.. code-block:: bash
+
+ 7767517
+ 1060 1342
+ Input in0 0 1 in0
+
+to
+
+.. code-block:: bash
+
+ 7767517
+ 1061 1342
+ SherpaMetaData sherpa_meta_data1 0 0 0=1 1=12 2=32 3=31 4=8 5=32 6=8 7=512
+ Input in0 0 1 in0
+
+.. caution::
+
+ Please follow :ref:`conv-emformer-modify-the-exported-encoder-for-sherpa-ncnn`
+ to change the values for ``SherpaMetaData`` if your model uses a different setting.
+
+
+Next, let us compile `sherpa-ncnn`_ since we will quantize our models within
+`sherpa-ncnn`_.
+
+.. code-block:: bash
+
+ # We will download sherpa-ncnn to $HOME/open-source/
+ # You can change it to anywhere you like.
+ cd $HOME
+ mkdir -p open-source
+
+ cd open-source
+ git clone https://github.com/k2-fsa/sherpa-ncnn
+ cd sherpa-ncnn
+ mkdir build
+ cd build
+ cmake ..
+ make -j 4
+
+ ./bin/generate-int8-scale-table
+
+ export PATH=$HOME/open-source/sherpa-ncnn/build/bin:$PATH
+
+The output of the above commands are:
+
+.. code-block:: bash
+
+ (py38) kuangfangjun:build$ generate-int8-scale-table
+ Please provide 10 arg. Currently given: 1
+ Usage:
+ generate-int8-scale-table encoder.param encoder.bin decoder.param decoder.bin joiner.param joiner.bin encoder-scale-table.txt joiner-scale-table.txt wave_filenames.txt
+
+ Each line in wave_filenames.txt is a path to some 16k Hz mono wave file.
+
+We need to create a file ``wave_filenames.txt``, in which we need to put
+some calibration wave files. For testing purpose, we put the ``test_wavs``
+from the pre-trained model repository ``_
+
+.. code-block:: bash
+
+ cd egs/librispeech/ASR
+ cd icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/
+
+ cat < wave_filenames.txt
+ ../test_wavs/1089-134686-0001.wav
+ ../test_wavs/1221-135766-0001.wav
+ ../test_wavs/1221-135766-0002.wav
+ EOF
+
+Now we can calculate the scales needed for quantization with the calibration data:
+
+.. code-block:: bash
+
+ cd egs/librispeech/ASR
+ cd icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/
+
+ generate-int8-scale-table \
+ ./encoder_jit_trace-pnnx.ncnn.param \
+ ./encoder_jit_trace-pnnx.ncnn.bin \
+ ./decoder_jit_trace-pnnx.ncnn.param \
+ ./decoder_jit_trace-pnnx.ncnn.bin \
+ ./joiner_jit_trace-pnnx.ncnn.param \
+ ./joiner_jit_trace-pnnx.ncnn.bin \
+ ./encoder-scale-table.txt \
+ ./joiner-scale-table.txt \
+ ./wave_filenames.txt
+
+The output logs are in the following:
+
+.. literalinclude:: ./code/generate-int-8-scale-table-for-conv-emformer.txt
+
+It generates the following two files:
+
+.. code-block:: bash
+
+ $ ls -lh encoder-scale-table.txt joiner-scale-table.txt
+ -rw-r--r-- 1 kuangfangjun root 955K Jan 11 17:28 encoder-scale-table.txt
+ -rw-r--r-- 1 kuangfangjun root 18K Jan 11 17:28 joiner-scale-table.txt
+
+.. caution::
+
+ Definitely, you need more calibration data to compute the scale table.
+
+Finally, let us use the scale table to quantize our models into ``int8``.
+
+.. code-block:: bash
+
+ ncnn2int8
+
+ usage: ncnn2int8 [inparam] [inbin] [outparam] [outbin] [calibration table]
+
+First, we quantize the encoder model:
+
+.. code-block:: bash
+
+ cd egs/librispeech/ASR
+ cd icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/
+
+ ncnn2int8 \
+ ./encoder_jit_trace-pnnx.ncnn.param \
+ ./encoder_jit_trace-pnnx.ncnn.bin \
+ ./encoder_jit_trace-pnnx.ncnn.int8.param \
+ ./encoder_jit_trace-pnnx.ncnn.int8.bin \
+ ./encoder-scale-table.txt
+
+Next, we quantize the joiner model:
+
+.. code-block:: bash
+
+ ncnn2int8 \
+ ./joiner_jit_trace-pnnx.ncnn.param \
+ ./joiner_jit_trace-pnnx.ncnn.bin \
+ ./joiner_jit_trace-pnnx.ncnn.int8.param \
+ ./joiner_jit_trace-pnnx.ncnn.int8.bin \
+ ./joiner-scale-table.txt
+
+The above two commands generate the following 4 files:
+
+.. code-block:: bash
+
+ -rw-r--r-- 1 kuangfangjun root 99M Jan 11 17:34 encoder_jit_trace-pnnx.ncnn.int8.bin
+ -rw-r--r-- 1 kuangfangjun root 78K Jan 11 17:34 encoder_jit_trace-pnnx.ncnn.int8.param
+ -rw-r--r-- 1 kuangfangjun root 774K Jan 11 17:35 joiner_jit_trace-pnnx.ncnn.int8.bin
+ -rw-r--r-- 1 kuangfangjun root 496 Jan 11 17:35 joiner_jit_trace-pnnx.ncnn.int8.param
+
+Congratulations! You have successfully quantized your model from ``float32`` to ``int8``.
+
+.. caution::
+
+ ``ncnn.int8.param`` and ``ncnn.int8.bin`` must be used in pairs.
+
+ You can replace ``ncnn.param`` and ``ncnn.bin`` with ``ncnn.int8.param``
+ and ``ncnn.int8.bin`` in `sherpa-ncnn`_ if you like.
+
+ For instance, to use only the ``int8`` encoder in ``sherpa-ncnn``, you can
+ replace the following invocation:
+
+ .. code-block::
+
+ cd egs/librispeech/ASR
+ cd icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/
+
+ sherpa-ncnn \
+ ../data/lang_bpe_500/tokens.txt \
+ ./encoder_jit_trace-pnnx.ncnn.param \
+ ./encoder_jit_trace-pnnx.ncnn.bin \
+ ./decoder_jit_trace-pnnx.ncnn.param \
+ ./decoder_jit_trace-pnnx.ncnn.bin \
+ ./joiner_jit_trace-pnnx.ncnn.param \
+ ./joiner_jit_trace-pnnx.ncnn.bin \
+ ../test_wavs/1089-134686-0001.wav
+
+ with
+
+ .. code-block::
+
+ cd egs/librispeech/ASR
+ cd icefall-asr-librispeech-conv-emformer-transducer-stateless2-2022-07-05/exp/
+
+ sherpa-ncnn \
+ ../data/lang_bpe_500/tokens.txt \
+ ./encoder_jit_trace-pnnx.ncnn.int8.param \
+ ./encoder_jit_trace-pnnx.ncnn.int8.bin \
+ ./decoder_jit_trace-pnnx.ncnn.param \
+ ./decoder_jit_trace-pnnx.ncnn.bin \
+ ./joiner_jit_trace-pnnx.ncnn.param \
+ ./joiner_jit_trace-pnnx.ncnn.bin \
+ ../test_wavs/1089-134686-0001.wav
+
+
+The following table compares again the file sizes:
+
+
++----------------------------------------+------------+
+| File name | File size |
++----------------------------------------+------------+
+| encoder_jit_trace-pnnx.pt | 283 MB |
++----------------------------------------+------------+
+| decoder_jit_trace-pnnx.pt | 1010 KB |
++----------------------------------------+------------+
+| joiner_jit_trace-pnnx.pt | 3.0 MB |
++----------------------------------------+------------+
+| encoder_jit_trace-pnnx.ncnn.bin (fp16) | 142 MB |
++----------------------------------------+------------+
+| decoder_jit_trace-pnnx.ncnn.bin (fp16) | 503 KB |
++----------------------------------------+------------+
+| joiner_jit_trace-pnnx.ncnn.bin (fp16) | 1.5 MB |
++----------------------------------------+------------+
+| encoder_jit_trace-pnnx.ncnn.bin (fp32) | 283 MB |
++----------------------------------------+------------+
+| joiner_jit_trace-pnnx.ncnn.bin (fp32) | 3.0 MB |
++----------------------------------------+------------+
+| encoder_jit_trace-pnnx.ncnn.int8.bin | 99 MB |
++----------------------------------------+------------+
+| joiner_jit_trace-pnnx.ncnn.int8.bin | 774 KB |
++----------------------------------------+------------+
+
+You can see that the file sizes of the model after ``int8`` quantization
+are much smaller.
+
+.. hint::
+
+ Currently, only linear layers and convolutional layers are quantized
+ with ``int8``, so you don't see an exact ``4x`` reduction in file sizes.
+
+.. note::
+
+ You need to test the recognition accuracy after ``int8`` quantization.
+
+You can find the speed comparison at ``_.
+
+
+That's it! Have fun with `sherpa-ncnn`_!
diff --git a/docs/source/model-export/export-with-torch-jit-script.rst b/docs/source/model-export/export-with-torch-jit-script.rst
index a041dc1d5..efd7dc2e1 100644
--- a/docs/source/model-export/export-with-torch-jit-script.rst
+++ b/docs/source/model-export/export-with-torch-jit-script.rst
@@ -1,7 +1,7 @@
.. _export-model-with-torch-jit-script:
Export model with torch.jit.script()
-===================================
+====================================
In this section, we describe how to export a model via
``torch.jit.script()``.
diff --git a/docs/source/recipes/aishell/conformer_ctc.rst b/docs/source/recipes/Non-streaming-ASR/aishell/conformer_ctc.rst
similarity index 99%
rename from docs/source/recipes/aishell/conformer_ctc.rst
rename to docs/source/recipes/Non-streaming-ASR/aishell/conformer_ctc.rst
index 72690e102..6e30ce397 100644
--- a/docs/source/recipes/aishell/conformer_ctc.rst
+++ b/docs/source/recipes/Non-streaming-ASR/aishell/conformer_ctc.rst
@@ -703,7 +703,7 @@ It will show you the following message:
HLG decoding
-^^^^^^^^^^^^
+~~~~~~~~~~~~
.. code-block:: bash
diff --git a/docs/source/recipes/aishell/images/aishell-conformer-ctc-tensorboard-log.jpg b/docs/source/recipes/Non-streaming-ASR/aishell/images/aishell-conformer-ctc-tensorboard-log.jpg
similarity index 100%
rename from docs/source/recipes/aishell/images/aishell-conformer-ctc-tensorboard-log.jpg
rename to docs/source/recipes/Non-streaming-ASR/aishell/images/aishell-conformer-ctc-tensorboard-log.jpg
diff --git a/docs/source/recipes/aishell/images/aishell-tdnn-lstm-ctc-tensorboard-log.jpg b/docs/source/recipes/Non-streaming-ASR/aishell/images/aishell-tdnn-lstm-ctc-tensorboard-log.jpg
similarity index 100%
rename from docs/source/recipes/aishell/images/aishell-tdnn-lstm-ctc-tensorboard-log.jpg
rename to docs/source/recipes/Non-streaming-ASR/aishell/images/aishell-tdnn-lstm-ctc-tensorboard-log.jpg
diff --git a/docs/source/recipes/aishell/images/aishell-transducer_stateless_modified-tensorboard-log.png b/docs/source/recipes/Non-streaming-ASR/aishell/images/aishell-transducer_stateless_modified-tensorboard-log.png
similarity index 100%
rename from docs/source/recipes/aishell/images/aishell-transducer_stateless_modified-tensorboard-log.png
rename to docs/source/recipes/Non-streaming-ASR/aishell/images/aishell-transducer_stateless_modified-tensorboard-log.png
diff --git a/docs/source/recipes/aishell/index.rst b/docs/source/recipes/Non-streaming-ASR/aishell/index.rst
similarity index 100%
rename from docs/source/recipes/aishell/index.rst
rename to docs/source/recipes/Non-streaming-ASR/aishell/index.rst
diff --git a/docs/source/recipes/aishell/stateless_transducer.rst b/docs/source/recipes/Non-streaming-ASR/aishell/stateless_transducer.rst
similarity index 100%
rename from docs/source/recipes/aishell/stateless_transducer.rst
rename to docs/source/recipes/Non-streaming-ASR/aishell/stateless_transducer.rst
diff --git a/docs/source/recipes/aishell/tdnn_lstm_ctc.rst b/docs/source/recipes/Non-streaming-ASR/aishell/tdnn_lstm_ctc.rst
similarity index 100%
rename from docs/source/recipes/aishell/tdnn_lstm_ctc.rst
rename to docs/source/recipes/Non-streaming-ASR/aishell/tdnn_lstm_ctc.rst
diff --git a/docs/source/recipes/Non-streaming-ASR/index.rst b/docs/source/recipes/Non-streaming-ASR/index.rst
new file mode 100644
index 000000000..67123a648
--- /dev/null
+++ b/docs/source/recipes/Non-streaming-ASR/index.rst
@@ -0,0 +1,10 @@
+Non Streaming ASR
+=================
+
+.. toctree::
+ :maxdepth: 2
+
+ aishell/index
+ librispeech/index
+ timit/index
+ yesno/index
diff --git a/docs/source/recipes/librispeech/conformer_ctc.rst b/docs/source/recipes/Non-streaming-ASR/librispeech/conformer_ctc.rst
similarity index 99%
rename from docs/source/recipes/librispeech/conformer_ctc.rst
rename to docs/source/recipes/Non-streaming-ASR/librispeech/conformer_ctc.rst
index 4656acfd6..b7f89c89f 100644
--- a/docs/source/recipes/librispeech/conformer_ctc.rst
+++ b/docs/source/recipes/Non-streaming-ASR/librispeech/conformer_ctc.rst
@@ -888,7 +888,7 @@ It will show you the following message:
CTC decoding
-^^^^^^^^^^^^
+~~~~~~~~~~~~
.. code-block:: bash
@@ -926,7 +926,7 @@ Its output is:
YET THESE THOUGHTS AFFECTED HESTER PRYNNE LESS WITH HOPE THAN APPREHENSION
HLG decoding
-^^^^^^^^^^^^
+~~~~~~~~~~~~
.. code-block:: bash
@@ -966,7 +966,7 @@ The output is:
HLG decoding + n-gram LM rescoring
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. code-block:: bash
@@ -1012,7 +1012,7 @@ The output is:
HLG decoding + n-gram LM rescoring + attention decoder rescoring
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. code-block:: bash
diff --git a/docs/source/recipes/Non-streaming-ASR/librispeech/distillation.rst b/docs/source/recipes/Non-streaming-ASR/librispeech/distillation.rst
new file mode 100644
index 000000000..ea9f350cd
--- /dev/null
+++ b/docs/source/recipes/Non-streaming-ASR/librispeech/distillation.rst
@@ -0,0 +1,223 @@
+Distillation with HuBERT
+========================
+
+This tutorial shows you how to perform knowledge distillation in `icefall`_
+with the `LibriSpeech`_ dataset. The distillation method
+used here is called "Multi Vector Quantization Knowledge Distillation" (MVQ-KD).
+Please have a look at our paper `Predicting Multi-Codebook Vector Quantization Indexes for Knowledge Distillation `_
+for more details about MVQ-KD.
+
+.. note::
+
+ This tutorial is based on recipe
+ `pruned_transducer_stateless4 `_.
+ Currently, we only implement MVQ-KD in this recipe. However, MVQ-KD is theoretically applicable to all recipes
+ with only minor changes needed. Feel free to try out MVQ-KD in different recipes. If you
+ encounter any problems, please open an issue here `icefall `_.
+
+.. note::
+
+ We assume you have read the page :ref:`install icefall` and have setup
+ the environment for `icefall`_.
+
+.. HINT::
+
+ We recommend you to use a GPU or several GPUs to run this recipe.
+
+Data preparation
+----------------
+
+We first prepare necessary training data for `LibriSpeech`_.
+This is the same as in :ref:`non_streaming_librispeech_pruned_transducer_stateless`.
+
+.. hint::
+
+ The data preparation is the same as other recipes on LibriSpeech dataset,
+ if you have finished this step, you can skip to :ref:`codebook_index_preparation` directly.
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./prepare.sh
+
+The script ``./prepare.sh`` handles the data preparation for you, **automagically**.
+All you need to do is to run it.
+
+The data preparation contains several stages, you can use the following two
+options:
+
+ - ``--stage``
+ - ``--stop-stage``
+
+to control which stage(s) should be run. By default, all stages are executed.
+
+For example,
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./prepare.sh --stage 0 --stop-stage 0 # run only stage 0
+ $ ./prepare.sh --stage 2 --stop-stage 5 # run from stage 2 to stage 5
+
+.. HINT::
+
+ If you have pre-downloaded the `LibriSpeech`_
+ dataset and the `musan`_ dataset, say,
+ they are saved in ``/tmp/LibriSpeech`` and ``/tmp/musan``, you can modify
+ the ``dl_dir`` variable in ``./prepare.sh`` to point to ``/tmp`` so that
+ ``./prepare.sh`` won't re-download them.
+
+.. NOTE::
+
+ All generated files by ``./prepare.sh``, e.g., features, lexicon, etc,
+ are saved in ``./data`` directory.
+
+We provide the following YouTube video showing how to run ``./prepare.sh``.
+
+.. note::
+
+ To get the latest news of `next-gen Kaldi `_, please subscribe
+ the following YouTube channel by `Nadira Povey `_:
+
+ ``_
+
+.. youtube:: ofEIoJL-mGM
+
+
+.. _codebook_index_preparation:
+
+Codebook index preparation
+--------------------------
+
+Here, we prepare necessary data for MVQ-KD. This requires the generation
+of codebook indexes (please read our `paper `_.
+if you are interested in details). In this tutorial, we use the pre-computed
+codebook indexes for convenience. The only thing you need to do is to
+run `./distillation_with_hubert.sh `_.
+
+.. note::
+
+ There are 5 stages in total, the first and second stage will be automatically skipped
+ when choosing to downloaded codebook indexes prepared by `icefall`_.
+ Of course, you can extract and compute the codebook indexes by yourself. This
+ will require you downloading a HuBERT-XL model and it can take a while for
+ the extraction of codebook indexes.
+
+
+As usual, you can control the stages you want to run by specifying the following
+two options:
+
+ - ``--stage``
+ - ``--stop-stage``
+
+For example,
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./distillation_with_hubert.sh --stage 0 --stop-stage 0 # run only stage 0
+ $ ./distillation_with_hubert.sh --stage 2 --stop-stage 4 # run from stage 2 to stage 5
+
+Here are a few options in `./distillation_with_hubert.sh `_
+you need to know before you proceed.
+
+- ``--full_libri`` If True, use full 960h data. Otherwise only ``train-clean-100`` will be used
+- ``--use_extracted_codebook`` If True, the first two stages will be skipped and the codebook
+ indexes uploaded by us will be downloaded.
+
+Since we are using the pre-computed codebook indexes, we set
+``use_extracted_codebook=True``. If you want to do full `LibriSpeech`_
+experiments, please set ``full_libri=True``.
+
+The following command downloads the pre-computed codebook indexes
+and prepares MVQ-augmented training manifests.
+
+.. code-block:: bash
+
+ $ ./distillation_with_hubert.sh --stage 2 --stop-stage 2 # run only stage 2
+
+Please see the
+following screenshot for the output of an example execution.
+
+.. figure:: ./images/distillation_codebook.png
+ :width: 800
+ :alt: Downloading codebook indexes and preparing training manifest.
+ :align: center
+
+ Downloading codebook indexes and preparing training manifest.
+
+.. hint::
+
+ The codebook indexes we prepared for you in this tutorial
+ are extracted from the 36-th layer of a fine-tuned HuBERT-XL model
+ with 8 codebooks. If you want to try other configurations, please
+ set ``use_extracted_codebook=False`` and set ``embedding_layer`` and
+ ``num_codebooks`` by yourself.
+
+Now, you should see the following files under the directory ``./data/vq_fbank_layer36_cb8``.
+
+.. figure:: ./images/distillation_directory.png
+ :width: 800
+ :alt: MVQ-augmented training manifests
+ :align: center
+
+ MVQ-augmented training manifests.
+
+Whola! You are ready to perform knowledge distillation training now!
+
+Training
+--------
+
+To perform training, please run stage 3 by executing the following command.
+
+.. code-block:: bash
+
+ $ ./prepare.sh --stage 3 --stop-stage 3 # run MVQ training
+
+Here is the code snippet for training:
+
+.. code-block:: bash
+
+ WORLD_SIZE=$(echo ${CUDA_VISIBLE_DEVICES} | awk '{n=split($1, _, ","); print n}')
+
+ ./pruned_transducer_stateless6/train.py \
+ --manifest-dir ./data/vq_fbank_layer36_cb8 \
+ --master-port 12359 \
+ --full-libri $full_libri \
+ --spec-aug-time-warp-factor -1 \
+ --max-duration 300 \
+ --world-size ${WORLD_SIZE} \
+ --num-epochs 30 \
+ --exp-dir $exp_dir \
+ --enable-distillation True \
+ --codebook-loss-scale 0.01
+
+There are a few training arguments in the following
+training commands that should be paid attention to.
+
+ - ``--enable-distillation`` If True, knowledge distillation training is enabled.
+ - ``--codebook-loss-scale`` The scale of the knowledge distillation loss.
+ - ``--manifest-dir`` The path to the MVQ-augmented manifest.
+
+
+Decoding
+--------
+
+After training finished, you can test the performance on using
+the following command.
+
+.. code-block:: bash
+
+ export CUDA_VISIBLE_DEVICES=0
+ ./pruned_transducer_stateless6/train.py \
+ --decoding-method "modified_beam_search" \
+ --epoch 30 \
+ --avg 10 \
+ --max-duration 200 \
+ --exp-dir $exp_dir \
+ --enable-distillation True
+
+You should get similar results as `here `_.
+
+That's all! Feel free to experiment with your own setups and report your results.
+If you encounter any problems during training, please open up an issue `here `_.
diff --git a/docs/source/recipes/Non-streaming-ASR/librispeech/images/distillation_codebook.png b/docs/source/recipes/Non-streaming-ASR/librispeech/images/distillation_codebook.png
new file mode 100644
index 000000000..1a40d6c6e
Binary files /dev/null and b/docs/source/recipes/Non-streaming-ASR/librispeech/images/distillation_codebook.png differ
diff --git a/docs/source/recipes/Non-streaming-ASR/librispeech/images/distillation_directory.png b/docs/source/recipes/Non-streaming-ASR/librispeech/images/distillation_directory.png
new file mode 100644
index 000000000..30763046f
Binary files /dev/null and b/docs/source/recipes/Non-streaming-ASR/librispeech/images/distillation_directory.png differ
diff --git a/docs/source/recipes/librispeech/images/librispeech-conformer-ctc-tensorboard-log.png b/docs/source/recipes/Non-streaming-ASR/librispeech/images/librispeech-conformer-ctc-tensorboard-log.png
similarity index 100%
rename from docs/source/recipes/librispeech/images/librispeech-conformer-ctc-tensorboard-log.png
rename to docs/source/recipes/Non-streaming-ASR/librispeech/images/librispeech-conformer-ctc-tensorboard-log.png
diff --git a/docs/source/recipes/Non-streaming-ASR/librispeech/images/librispeech-pruned-transducer-tensorboard-log.jpg b/docs/source/recipes/Non-streaming-ASR/librispeech/images/librispeech-pruned-transducer-tensorboard-log.jpg
new file mode 100644
index 000000000..800835749
Binary files /dev/null and b/docs/source/recipes/Non-streaming-ASR/librispeech/images/librispeech-pruned-transducer-tensorboard-log.jpg differ
diff --git a/docs/source/recipes/Non-streaming-ASR/librispeech/index.rst b/docs/source/recipes/Non-streaming-ASR/librispeech/index.rst
new file mode 100644
index 000000000..bf439861a
--- /dev/null
+++ b/docs/source/recipes/Non-streaming-ASR/librispeech/index.rst
@@ -0,0 +1,12 @@
+LibriSpeech
+===========
+
+.. toctree::
+ :maxdepth: 1
+
+ tdnn_lstm_ctc
+ conformer_ctc
+ pruned_transducer_stateless
+ zipformer_mmi
+ zipformer_ctc_blankskip
+ distillation
diff --git a/docs/source/recipes/Non-streaming-ASR/librispeech/pruned_transducer_stateless.rst b/docs/source/recipes/Non-streaming-ASR/librispeech/pruned_transducer_stateless.rst
new file mode 100644
index 000000000..42fd3df77
--- /dev/null
+++ b/docs/source/recipes/Non-streaming-ASR/librispeech/pruned_transducer_stateless.rst
@@ -0,0 +1,548 @@
+.. _non_streaming_librispeech_pruned_transducer_stateless:
+
+Pruned transducer statelessX
+============================
+
+This tutorial shows you how to run a conformer transducer model
+with the `LibriSpeech `_ dataset.
+
+.. Note::
+
+ The tutorial is suitable for `pruned_transducer_stateless `_,
+ `pruned_transducer_stateless2 `_,
+ `pruned_transducer_stateless4 `_,
+ `pruned_transducer_stateless5 `_,
+ We will take pruned_transducer_stateless4 as an example in this tutorial.
+
+.. HINT::
+
+ We assume you have read the page :ref:`install icefall` and have setup
+ the environment for ``icefall``.
+
+.. HINT::
+
+ We recommend you to use a GPU or several GPUs to run this recipe.
+
+.. hint::
+
+ Please scroll down to the bottom of this page to find download links
+ for pretrained models if you don't want to train a model from scratch.
+
+
+We use pruned RNN-T to compute the loss.
+
+.. note::
+
+ You can find the paper about pruned RNN-T at the following address:
+
+ ``_
+
+The transducer model consists of 3 parts:
+
+ - Encoder, a.k.a, the transcription network. We use a Conformer model (the reworked version by Daniel Povey)
+ - Decoder, a.k.a, the prediction network. We use a stateless model consisting of
+ ``nn.Embedding`` and ``nn.Conv1d``
+ - Joiner, a.k.a, the joint network.
+
+.. caution::
+
+ Contrary to the conventional RNN-T models, we use a stateless decoder.
+ That is, it has no recurrent connections.
+
+
+Data preparation
+----------------
+
+.. hint::
+
+ The data preparation is the same as other recipes on LibriSpeech dataset,
+ if you have finished this step, you can skip to ``Training`` directly.
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./prepare.sh
+
+The script ``./prepare.sh`` handles the data preparation for you, **automagically**.
+All you need to do is to run it.
+
+The data preparation contains several stages, you can use the following two
+options:
+
+ - ``--stage``
+ - ``--stop-stage``
+
+to control which stage(s) should be run. By default, all stages are executed.
+
+
+For example,
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./prepare.sh --stage 0 --stop-stage 0
+
+means to run only stage 0.
+
+To run stage 2 to stage 5, use:
+
+.. code-block:: bash
+
+ $ ./prepare.sh --stage 2 --stop-stage 5
+
+.. HINT::
+
+ If you have pre-downloaded the `LibriSpeech `_
+ dataset and the `musan `_ dataset, say,
+ they are saved in ``/tmp/LibriSpeech`` and ``/tmp/musan``, you can modify
+ the ``dl_dir`` variable in ``./prepare.sh`` to point to ``/tmp`` so that
+ ``./prepare.sh`` won't re-download them.
+
+.. NOTE::
+
+ All generated files by ``./prepare.sh``, e.g., features, lexicon, etc,
+ are saved in ``./data`` directory.
+
+We provide the following YouTube video showing how to run ``./prepare.sh``.
+
+.. note::
+
+ To get the latest news of `next-gen Kaldi `_, please subscribe
+ the following YouTube channel by `Nadira Povey `_:
+
+ ``_
+
+.. youtube:: ofEIoJL-mGM
+
+
+Training
+--------
+
+Configurable options
+~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./pruned_transducer_stateless4/train.py --help
+
+
+shows you the training options that can be passed from the commandline.
+The following options are used quite often:
+
+ - ``--exp-dir``
+
+ The directory to save checkpoints, training logs and tensorboard.
+
+ - ``--full-libri``
+
+ If it's True, the training part uses all the training data, i.e.,
+ 960 hours. Otherwise, the training part uses only the subset
+ ``train-clean-100``, which has 100 hours of training data.
+
+ .. CAUTION::
+ The training set is perturbed by speed with two factors: 0.9 and 1.1.
+ If ``--full-libri`` is True, each epoch actually processes
+ ``3x960 == 2880`` hours of data.
+
+ - ``--num-epochs``
+
+ It is the number of epochs to train. For instance,
+ ``./pruned_transducer_stateless4/train.py --num-epochs 30`` trains for 30 epochs
+ and generates ``epoch-1.pt``, ``epoch-2.pt``, ..., ``epoch-30.pt``
+ in the folder ``./pruned_transducer_stateless4/exp``.
+
+ - ``--start-epoch``
+
+ It's used to resume training.
+ ``./pruned_transducer_stateless4/train.py --start-epoch 10`` loads the
+ checkpoint ``./pruned_transducer_stateless4/exp/epoch-9.pt`` and starts
+ training from epoch 10, based on the state from epoch 9.
+
+ - ``--world-size``
+
+ It is used for multi-GPU single-machine DDP training.
+
+ - (a) If it is 1, then no DDP training is used.
+
+ - (b) If it is 2, then GPU 0 and GPU 1 are used for DDP training.
+
+ The following shows some use cases with it.
+
+ **Use case 1**: You have 4 GPUs, but you only want to use GPU 0 and
+ GPU 2 for training. You can do the following:
+
+ .. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ export CUDA_VISIBLE_DEVICES="0,2"
+ $ ./pruned_transducer_stateless4/train.py --world-size 2
+
+ **Use case 2**: You have 4 GPUs and you want to use all of them
+ for training. You can do the following:
+
+ .. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./pruned_transducer_stateless4/train.py --world-size 4
+
+ **Use case 3**: You have 4 GPUs but you only want to use GPU 3
+ for training. You can do the following:
+
+ .. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ export CUDA_VISIBLE_DEVICES="3"
+ $ ./pruned_transducer_stateless4/train.py --world-size 1
+
+ .. caution::
+
+ Only multi-GPU single-machine DDP training is implemented at present.
+ Multi-GPU multi-machine DDP training will be added later.
+
+ - ``--max-duration``
+
+ It specifies the number of seconds over all utterances in a
+ batch, before **padding**.
+ If you encounter CUDA OOM, please reduce it.
+
+ .. HINT::
+
+ Due to padding, the number of seconds of all utterances in a
+ batch will usually be larger than ``--max-duration``.
+
+ A larger value for ``--max-duration`` may cause OOM during training,
+ while a smaller value may increase the training time. You have to
+ tune it.
+
+ - ``--use-fp16``
+
+ If it is True, the model will train with half precision, from our experiment
+ results, by using half precision you can train with two times larger ``--max-duration``
+ so as to get almost 2X speed up.
+
+
+Pre-configured options
+~~~~~~~~~~~~~~~~~~~~~~
+
+There are some training options, e.g., number of encoder layers,
+encoder dimension, decoder dimension, number of warmup steps etc,
+that are not passed from the commandline.
+They are pre-configured by the function ``get_params()`` in
+`pruned_transducer_stateless4/train.py `_
+
+You don't need to change these pre-configured parameters. If you really need to change
+them, please modify ``./pruned_transducer_stateless4/train.py`` directly.
+
+
+.. NOTE::
+
+ The options for `pruned_transducer_stateless5 `_ are a little different from
+ other recipes. It allows you to configure ``--num-encoder-layers``, ``--dim-feedforward``, ``--nhead``, ``--encoder-dim``, ``--decoder-dim``, ``--joiner-dim`` from commandline, so that you can train models with different size with pruned_transducer_stateless5.
+
+
+Training logs
+~~~~~~~~~~~~~
+
+Training logs and checkpoints are saved in ``--exp-dir`` (e.g. ``pruned_transducer_stateless4/exp``.
+You will find the following files in that directory:
+
+ - ``epoch-1.pt``, ``epoch-2.pt``, ...
+
+ These are checkpoint files saved at the end of each epoch, containing model
+ ``state_dict`` and optimizer ``state_dict``.
+ To resume training from some checkpoint, say ``epoch-10.pt``, you can use:
+
+ .. code-block:: bash
+
+ $ ./pruned_transducer_stateless4/train.py --start-epoch 11
+
+ - ``checkpoint-436000.pt``, ``checkpoint-438000.pt``, ...
+
+ These are checkpoint files saved every ``--save-every-n`` batches,
+ containing model ``state_dict`` and optimizer ``state_dict``.
+ To resume training from some checkpoint, say ``checkpoint-436000``, you can use:
+
+ .. code-block:: bash
+
+ $ ./pruned_transducer_stateless4/train.py --start-batch 436000
+
+ - ``tensorboard/``
+
+ This folder contains tensorBoard logs. Training loss, validation loss, learning
+ rate, etc, are recorded in these logs. You can visualize them by:
+
+ .. code-block:: bash
+
+ $ cd pruned_transducer_stateless4/exp/tensorboard
+ $ tensorboard dev upload --logdir . --description "pruned transducer training for LibriSpeech with icefall"
+
+ It will print something like below:
+
+ .. code-block::
+
+ TensorFlow installation not found - running with reduced feature set.
+ Upload started and will continue reading any new data as it's added to the logdir.
+
+ To stop uploading, press Ctrl-C.
+
+ New experiment created. View your TensorBoard at: https://tensorboard.dev/experiment/QOGSPBgsR8KzcRMmie9JGw/
+
+ [2022-11-20T15:50:50] Started scanning logdir.
+ Uploading 4468 scalars...
+ [2022-11-20T15:53:02] Total uploaded: 210171 scalars, 0 tensors, 0 binary objects
+ Listening for new data in logdir...
+
+ Note there is a URL in the above output. Click it and you will see
+ the following screenshot:
+
+ .. figure:: images/librispeech-pruned-transducer-tensorboard-log.jpg
+ :width: 600
+ :alt: TensorBoard screenshot
+ :align: center
+ :target: https://tensorboard.dev/experiment/QOGSPBgsR8KzcRMmie9JGw/
+
+ TensorBoard screenshot.
+
+ .. hint::
+
+ If you don't have access to google, you can use the following command
+ to view the tensorboard log locally:
+
+ .. code-block:: bash
+
+ cd pruned_transducer_stateless4/exp/tensorboard
+ tensorboard --logdir . --port 6008
+
+ It will print the following message:
+
+ .. code-block::
+
+ Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
+ TensorBoard 2.8.0 at http://localhost:6008/ (Press CTRL+C to quit)
+
+ Now start your browser and go to ``_ to view the tensorboard
+ logs.
+
+
+ - ``log/log-train-xxxx``
+
+ It is the detailed training log in text format, same as the one
+ you saw printed to the console during training.
+
+Usage example
+~~~~~~~~~~~~~
+
+You can use the following command to start the training using 6 GPUs:
+
+.. code-block:: bash
+
+ export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5"
+ ./pruned_transducer_stateless4/train.py \
+ --world-size 6 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --exp-dir pruned_transducer_stateless4/exp \
+ --full-libri 1 \
+ --max-duration 300
+
+
+Decoding
+--------
+
+The decoding part uses checkpoints saved by the training part, so you have
+to run the training part first.
+
+.. hint::
+
+ There are two kinds of checkpoints:
+
+ - (1) ``epoch-1.pt``, ``epoch-2.pt``, ..., which are saved at the end
+ of each epoch. You can pass ``--epoch`` to
+ ``pruned_transducer_stateless4/decode.py`` to use them.
+
+ - (2) ``checkpoints-436000.pt``, ``epoch-438000.pt``, ..., which are saved
+ every ``--save-every-n`` batches. You can pass ``--iter`` to
+ ``pruned_transducer_stateless4/decode.py`` to use them.
+
+ We suggest that you try both types of checkpoints and choose the one
+ that produces the lowest WERs.
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./pruned_transducer_stateless4/decode.py --help
+
+shows the options for decoding.
+
+The following shows two examples (for two types of checkpoints):
+
+.. code-block:: bash
+
+ for m in greedy_search fast_beam_search modified_beam_search; do
+ for epoch in 25 20; do
+ for avg in 7 5 3 1; do
+ ./pruned_transducer_stateless4/decode.py \
+ --epoch $epoch \
+ --avg $avg \
+ --exp-dir pruned_transducer_stateless4/exp \
+ --max-duration 600 \
+ --decoding-method $m
+ done
+ done
+ done
+
+
+.. code-block:: bash
+
+ for m in greedy_search fast_beam_search modified_beam_search; do
+ for iter in 474000; do
+ for avg in 8 10 12 14 16 18; do
+ ./pruned_transducer_stateless4/decode.py \
+ --iter $iter \
+ --avg $avg \
+ --exp-dir pruned_transducer_stateless4/exp \
+ --max-duration 600 \
+ --decoding-method $m
+ done
+ done
+ done
+
+
+.. Note::
+
+ Supporting decoding methods are as follows:
+
+ - ``greedy_search`` : It takes the symbol with largest posterior probability
+ of each frame as the decoding result.
+
+ - ``beam_search`` : It implements Algorithm 1 in https://arxiv.org/pdf/1211.3711.pdf and
+ `espnet/nets/beam_search_transducer.py `_
+ is used as a reference. Basicly, it keeps topk states for each frame, and expands the kept states with their own contexts to
+ next frame.
+
+ - ``modified_beam_search`` : It implements the same algorithm as ``beam_search`` above, but it
+ runs in batch mode with ``--max-sym-per-frame=1`` being hardcoded.
+
+ - ``fast_beam_search`` : It implements graph composition between the output ``log_probs`` and
+ given ``FSAs``. It is hard to describe the details in several lines of texts, you can read
+ our paper in https://arxiv.org/pdf/2211.00484.pdf or our `rnnt decode code in k2 `_. ``fast_beam_search`` can decode with ``FSAs`` on GPU efficiently.
+
+ - ``fast_beam_search_LG`` : The same as ``fast_beam_search`` above, ``fast_beam_search`` uses
+ an trivial graph that has only one state, while ``fast_beam_search_LG`` uses an LG graph
+ (with N-gram LM).
+
+ - ``fast_beam_search_nbest`` : It produces the decoding results as follows:
+
+ - (1) Use ``fast_beam_search`` to get a lattice
+ - (2) Select ``num_paths`` paths from the lattice using ``k2.random_paths()``
+ - (3) Unique the selected paths
+ - (4) Intersect the selected paths with the lattice and compute the
+ shortest path from the intersection result
+ - (5) The path with the largest score is used as the decoding output.
+
+ - ``fast_beam_search_nbest_LG`` : It implements same logic as ``fast_beam_search_nbest``, the
+ only difference is that it uses ``fast_beam_search_LG`` to generate the lattice.
+
+
+Export Model
+------------
+
+`pruned_transducer_stateless4/export.py `_ supports exporting checkpoints from ``pruned_transducer_stateless4/exp`` in the following ways.
+
+Export ``model.state_dict()``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Checkpoints saved by ``pruned_transducer_stateless4/train.py`` also include
+``optimizer.state_dict()``. It is useful for resuming training. But after training,
+we are interested only in ``model.state_dict()``. You can use the following
+command to extract ``model.state_dict()``.
+
+.. code-block:: bash
+
+ # Assume that --epoch 25 --avg 3 produces the smallest WER
+ # (You can get such information after running ./pruned_transducer_stateless4/decode.py)
+
+ epoch=25
+ avg=3
+
+ ./pruned_transducer_stateless4/export.py \
+ --exp-dir ./pruned_transducer_stateless4/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch $epoch \
+ --avg $avg
+
+It will generate a file ``./pruned_transducer_stateless4/exp/pretrained.pt``.
+
+.. hint::
+
+ To use the generated ``pretrained.pt`` for ``pruned_transducer_stateless4/decode.py``,
+ you can run:
+
+ .. code-block:: bash
+
+ cd pruned_transducer_stateless4/exp
+ ln -s pretrained.pt epoch-999.pt
+
+ And then pass ``--epoch 999 --avg 1 --use-averaged-model 0`` to
+ ``./pruned_transducer_stateless4/decode.py``.
+
+To use the exported model with ``./pruned_transducer_stateless4/pretrained.py``, you
+can run:
+
+.. code-block:: bash
+
+ ./pruned_transducer_stateless4/pretrained.py \
+ --checkpoint ./pruned_transducer_stateless4/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method greedy_search \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+
+Export model using ``torch.jit.script()``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: bash
+
+ ./pruned_transducer_stateless4/export.py \
+ --exp-dir ./pruned_transducer_stateless4/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 25 \
+ --avg 3 \
+ --jit 1
+
+It will generate a file ``cpu_jit.pt`` in the given ``exp_dir``. You can later
+load it by ``torch.jit.load("cpu_jit.pt")``.
+
+Note ``cpu`` in the name ``cpu_jit.pt`` means the parameters when loaded into Python
+are on CPU. You can use ``to("cuda")`` to move them to a CUDA device.
+
+.. NOTE::
+
+ You will need this ``cpu_jit.pt`` when deploying with Sherpa framework.
+
+
+Download pretrained models
+--------------------------
+
+If you don't want to train from scratch, you can download the pretrained models
+by visiting the following links:
+
+ - `pruned_transducer_stateless `_
+
+ - `pruned_transducer_stateless2 `_
+
+ - `pruned_transducer_stateless4 `_
+
+ - `pruned_transducer_stateless5 `_
+
+ See ``_
+ for the details of the above pretrained models
+
+
+Deploy with Sherpa
+------------------
+
+Please see ``_
+for how to deploy the models in ``sherpa``.
diff --git a/docs/source/recipes/librispeech/tdnn_lstm_ctc.rst b/docs/source/recipes/Non-streaming-ASR/librispeech/tdnn_lstm_ctc.rst
similarity index 100%
rename from docs/source/recipes/librispeech/tdnn_lstm_ctc.rst
rename to docs/source/recipes/Non-streaming-ASR/librispeech/tdnn_lstm_ctc.rst
diff --git a/docs/source/recipes/Non-streaming-ASR/librispeech/zipformer_ctc_blankskip.rst b/docs/source/recipes/Non-streaming-ASR/librispeech/zipformer_ctc_blankskip.rst
new file mode 100644
index 000000000..56a420605
--- /dev/null
+++ b/docs/source/recipes/Non-streaming-ASR/librispeech/zipformer_ctc_blankskip.rst
@@ -0,0 +1,453 @@
+Zipformer CTC Blank Skip
+========================
+
+.. hint::
+
+ Please scroll down to the bottom of this page to find download links
+ for pretrained models if you don't want to train a model from scratch.
+
+
+This tutorial shows you how to train a Zipformer model based on the guidance from
+a co-trained CTC model using `blank skip method `_
+with the `LibriSpeech `_ dataset.
+
+.. note::
+
+ We use both CTC and RNN-T loss to train. During the forward pass, the encoder output
+ is first used to calculate the CTC posterior probability; then for each output frame,
+ if its blank posterior is bigger than some threshold, it will be simply discarded
+ from the encoder output. To prevent information loss, we also put a convolution module
+ similar to the one used in conformer (referred to as “LConv”) before the frame reduction.
+
+
+Data preparation
+----------------
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./prepare.sh
+
+The script ``./prepare.sh`` handles the data preparation for you, **automagically**.
+All you need to do is to run it.
+
+.. note::
+
+ We encourage you to read ``./prepare.sh``.
+
+The data preparation contains several stages. You can use the following two
+options:
+
+ - ``--stage``
+ - ``--stop-stage``
+
+to control which stage(s) should be run. By default, all stages are executed.
+
+
+For example,
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./prepare.sh --stage 0 --stop-stage 0
+
+means to run only stage 0.
+
+To run stage 2 to stage 5, use:
+
+.. code-block:: bash
+
+ $ ./prepare.sh --stage 2 --stop-stage 5
+
+.. hint::
+
+ If you have pre-downloaded the `LibriSpeech `_
+ dataset and the `musan `_ dataset, say,
+ they are saved in ``/tmp/LibriSpeech`` and ``/tmp/musan``, you can modify
+ the ``dl_dir`` variable in ``./prepare.sh`` to point to ``/tmp`` so that
+ ``./prepare.sh`` won't re-download them.
+
+.. note::
+
+ All generated files by ``./prepare.sh``, e.g., features, lexicon, etc,
+ are saved in ``./data`` directory.
+
+We provide the following YouTube video showing how to run ``./prepare.sh``.
+
+.. note::
+
+ To get the latest news of `next-gen Kaldi `_, please subscribe
+ the following YouTube channel by `Nadira Povey `_:
+
+ ``_
+
+.. youtube:: ofEIoJL-mGM
+
+Training
+--------
+
+For stability, it doesn`t use blank skip method until model warm-up.
+
+Configurable options
+~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./pruned_transducer_stateless7_ctc_bs/train.py --help
+
+shows you the training options that can be passed from the commandline.
+The following options are used quite often:
+
+ - ``--full-libri``
+
+ If it's True, the training part uses all the training data, i.e.,
+ 960 hours. Otherwise, the training part uses only the subset
+ ``train-clean-100``, which has 100 hours of training data.
+
+ .. CAUTION::
+
+ The training set is perturbed by speed with two factors: 0.9 and 1.1.
+ If ``--full-libri`` is True, each epoch actually processes
+ ``3x960 == 2880`` hours of data.
+
+ - ``--num-epochs``
+
+ It is the number of epochs to train. For instance,
+ ``./pruned_transducer_stateless7_ctc_bs/train.py --num-epochs 30`` trains for 30 epochs
+ and generates ``epoch-1.pt``, ``epoch-2.pt``, ..., ``epoch-30.pt``
+ in the folder ``./pruned_transducer_stateless7_ctc_bs/exp``.
+
+ - ``--start-epoch``
+
+ It's used to resume training.
+ ``./pruned_transducer_stateless7_ctc_bs/train.py --start-epoch 10`` loads the
+ checkpoint ``./pruned_transducer_stateless7_ctc_bs/exp/epoch-9.pt`` and starts
+ training from epoch 10, based on the state from epoch 9.
+
+ - ``--world-size``
+
+ It is used for multi-GPU single-machine DDP training.
+
+ - (a) If it is 1, then no DDP training is used.
+
+ - (b) If it is 2, then GPU 0 and GPU 1 are used for DDP training.
+
+ The following shows some use cases with it.
+
+ **Use case 1**: You have 4 GPUs, but you only want to use GPU 0 and
+ GPU 2 for training. You can do the following:
+
+ .. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ export CUDA_VISIBLE_DEVICES="0,2"
+ $ ./pruned_transducer_stateless7_ctc_bs/train.py --world-size 2
+
+ **Use case 2**: You have 4 GPUs and you want to use all of them
+ for training. You can do the following:
+
+ .. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./pruned_transducer_stateless7_ctc_bs/train.py --world-size 4
+
+ **Use case 3**: You have 4 GPUs but you only want to use GPU 3
+ for training. You can do the following:
+
+ .. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ export CUDA_VISIBLE_DEVICES="3"
+ $ ./pruned_transducer_stateless7_ctc_bs/train.py --world-size 1
+
+ .. caution::
+
+ Only multi-GPU single-machine DDP training is implemented at present.
+ Multi-GPU multi-machine DDP training will be added later.
+
+ - ``--max-duration``
+
+ It specifies the number of seconds over all utterances in a
+ batch, before **padding**.
+ If you encounter CUDA OOM, please reduce it.
+
+ .. HINT::
+
+ Due to padding, the number of seconds of all utterances in a
+ batch will usually be larger than ``--max-duration``.
+
+ A larger value for ``--max-duration`` may cause OOM during training,
+ while a smaller value may increase the training time. You have to
+ tune it.
+
+
+Pre-configured options
+~~~~~~~~~~~~~~~~~~~~~~
+
+There are some training options, e.g., weight decay,
+number of warmup steps, results dir, etc,
+that are not passed from the commandline.
+They are pre-configured by the function ``get_params()`` in
+`pruned_transducer_stateless7_ctc_bs/train.py `_
+
+You don't need to change these pre-configured parameters. If you really need to change
+them, please modify ``./pruned_transducer_stateless7_ctc_bs/train.py`` directly.
+
+Training logs
+~~~~~~~~~~~~~
+
+Training logs and checkpoints are saved in ``pruned_transducer_stateless7_ctc_bs/exp``.
+You will find the following files in that directory:
+
+ - ``epoch-1.pt``, ``epoch-2.pt``, ...
+
+ These are checkpoint files saved at the end of each epoch, containing model
+ ``state_dict`` and optimizer ``state_dict``.
+ To resume training from some checkpoint, say ``epoch-10.pt``, you can use:
+
+ .. code-block:: bash
+
+ $ ./pruned_transducer_stateless7_ctc_bs/train.py --start-epoch 11
+
+ - ``checkpoint-436000.pt``, ``checkpoint-438000.pt``, ...
+
+ These are checkpoint files saved every ``--save-every-n`` batches,
+ containing model ``state_dict`` and optimizer ``state_dict``.
+ To resume training from some checkpoint, say ``checkpoint-436000``, you can use:
+
+ .. code-block:: bash
+
+ $ ./pruned_transducer_stateless7_ctc_bs/train.py --start-batch 436000
+
+ - ``tensorboard/``
+
+ This folder contains tensorBoard logs. Training loss, validation loss, learning
+ rate, etc, are recorded in these logs. You can visualize them by:
+
+ .. code-block:: bash
+
+ $ cd pruned_transducer_stateless7_ctc_bs/exp/tensorboard
+ $ tensorboard dev upload --logdir . --description "Zipformer-CTC co-training using blank skip for LibriSpeech with icefall"
+
+ It will print something like below:
+
+ .. code-block::
+
+ TensorFlow installation not found - running with reduced feature set.
+ Upload started and will continue reading any new data as it's added to the logdir.
+
+ To stop uploading, press Ctrl-C.
+
+ New experiment created. View your TensorBoard at: https://tensorboard.dev/experiment/xyOZUKpEQm62HBIlUD4uPA/
+
+ Note there is a URL in the above output. Click it and you will see
+ tensorboard.
+
+ .. hint::
+
+ If you don't have access to google, you can use the following command
+ to view the tensorboard log locally:
+
+ .. code-block:: bash
+
+ cd pruned_transducer_stateless7_ctc_bs/exp/tensorboard
+ tensorboard --logdir . --port 6008
+
+ It will print the following message:
+
+ .. code-block::
+
+ Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
+ TensorBoard 2.8.0 at http://localhost:6008/ (Press CTRL+C to quit)
+
+ Now start your browser and go to ``_ to view the tensorboard
+ logs.
+
+
+ - ``log/log-train-xxxx``
+
+ It is the detailed training log in text format, same as the one
+ you saw printed to the console during training.
+
+Usage example
+~~~~~~~~~~~~~
+
+You can use the following command to start the training using 4 GPUs:
+
+.. code-block:: bash
+
+ export CUDA_VISIBLE_DEVICES="0,1,2,3"
+ ./pruned_transducer_stateless7_ctc_bs/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --full-libri 1 \
+ --exp-dir pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --use-fp16 1
+
+Decoding
+--------
+
+The decoding part uses checkpoints saved by the training part, so you have
+to run the training part first.
+
+.. hint::
+
+ There are two kinds of checkpoints:
+
+ - (1) ``epoch-1.pt``, ``epoch-2.pt``, ..., which are saved at the end
+ of each epoch. You can pass ``--epoch`` to
+ ``pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py`` to use them.
+
+ - (2) ``checkpoints-436000.pt``, ``epoch-438000.pt``, ..., which are saved
+ every ``--save-every-n`` batches. You can pass ``--iter`` to
+ ``pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py`` to use them.
+
+ We suggest that you try both types of checkpoints and choose the one
+ that produces the lowest WERs.
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py --help
+
+shows the options for decoding.
+
+The following shows the example using ``epoch-*.pt``:
+
+.. code-block:: bash
+
+ for m in greedy_search fast_beam_search modified_beam_search; do
+ ./pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py \
+ --epoch 30 \
+ --avg 13 \
+ --exp-dir pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method $m
+ done
+
+To test CTC branch, you can use the following command:
+
+.. code-block:: bash
+
+ for m in ctc-decoding 1best; do
+ ./pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py \
+ --epoch 30 \
+ --avg 13 \
+ --exp-dir pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method $m
+ done
+
+Export models
+-------------
+
+`pruned_transducer_stateless7_ctc_bs/export.py `_ supports exporting checkpoints from ``pruned_transducer_stateless7_ctc_bs/exp`` in the following ways.
+
+Export ``model.state_dict()``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Checkpoints saved by ``pruned_transducer_stateless7_ctc_bs/train.py`` also include
+``optimizer.state_dict()``. It is useful for resuming training. But after training,
+we are interested only in ``model.state_dict()``. You can use the following
+command to extract ``model.state_dict()``.
+
+.. code-block:: bash
+
+ ./pruned_transducer_stateless7_ctc_bs/export.py \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 13 \
+ --jit 0
+
+It will generate a file ``./pruned_transducer_stateless7_ctc_bs/exp/pretrained.pt``.
+
+.. hint::
+
+ To use the generated ``pretrained.pt`` for ``pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py``,
+ you can run:
+
+ .. code-block:: bash
+
+ cd pruned_transducer_stateless7_ctc_bs/exp
+ ln -s pretrained epoch-9999.pt
+
+ And then pass ``--epoch 9999 --avg 1 --use-averaged-model 0`` to
+ ``./pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py``.
+
+To use the exported model with ``./pruned_transducer_stateless7_ctc_bs/pretrained.py``, you
+can run:
+
+.. code-block:: bash
+
+ ./pruned_transducer_stateless7_ctc_bs/pretrained.py \
+ --checkpoint ./pruned_transducer_stateless7_ctc_bs/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method greedy_search \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+To test CTC branch using the exported model with ``./pruned_transducer_stateless7_ctc_bs/pretrained_ctc.py``:
+
+.. code-block:: bash
+
+ ./pruned_transducer_stateless7_ctc_bs/jit_pretrained_ctc.py \
+ --checkpoint ./pruned_transducer_stateless7_ctc_bs/exp/pretrained.pt \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --method ctc-decoding \
+ --sample-rate 16000 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+Export model using ``torch.jit.script()``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: bash
+
+ ./pruned_transducer_stateless7_ctc_bs/export.py \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 13 \
+ --jit 1
+
+It will generate a file ``cpu_jit.pt`` in the given ``exp_dir``. You can later
+load it by ``torch.jit.load("cpu_jit.pt")``.
+
+Note ``cpu`` in the name ``cpu_jit.pt`` means the parameters when loaded into Python
+are on CPU. You can use ``to("cuda")`` to move them to a CUDA device.
+
+To use the generated files with ``./pruned_transducer_stateless7_ctc_bs/jit_pretrained.py``:
+
+.. code-block:: bash
+
+ ./pruned_transducer_stateless7_ctc_bs/jit_pretrained.py \
+ --nn-model-filename ./pruned_transducer_stateless7_ctc_bs/exp/cpu_jit.pt \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+To test CTC branch using the generated files with ``./pruned_transducer_stateless7_ctc_bs/jit_pretrained_ctc.py``:
+
+.. code-block:: bash
+
+ ./pruned_transducer_stateless7_ctc_bs/jit_pretrained_ctc.py \
+ --model-filename ./pruned_transducer_stateless7_ctc_bs/exp/cpu_jit.pt \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --method ctc-decoding \
+ --sample-rate 16000 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+Download pretrained models
+--------------------------
+
+If you don't want to train from scratch, you can download the pretrained models
+by visiting the following links:
+
+ - ``_
+
+ See ``_
+ for the details of the above pretrained models
diff --git a/docs/source/recipes/Non-streaming-ASR/librispeech/zipformer_mmi.rst b/docs/source/recipes/Non-streaming-ASR/librispeech/zipformer_mmi.rst
new file mode 100644
index 000000000..a7b59a992
--- /dev/null
+++ b/docs/source/recipes/Non-streaming-ASR/librispeech/zipformer_mmi.rst
@@ -0,0 +1,422 @@
+Zipformer MMI
+===============
+
+.. hint::
+
+ Please scroll down to the bottom of this page to find download links
+ for pretrained models if you don't want to train a model from scratch.
+
+
+This tutorial shows you how to train an Zipformer MMI model
+with the `LibriSpeech `_ dataset.
+
+We use LF-MMI to compute the loss.
+
+.. note::
+
+ You can find the document about LF-MMI training at the following address:
+
+ ``_
+
+
+Data preparation
+----------------
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./prepare.sh
+
+The script ``./prepare.sh`` handles the data preparation for you, **automagically**.
+All you need to do is to run it.
+
+.. note::
+
+ We encourage you to read ``./prepare.sh``.
+
+The data preparation contains several stages. You can use the following two
+options:
+
+ - ``--stage``
+ - ``--stop-stage``
+
+to control which stage(s) should be run. By default, all stages are executed.
+
+
+For example,
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./prepare.sh --stage 0 --stop-stage 0
+
+means to run only stage 0.
+
+To run stage 2 to stage 5, use:
+
+.. code-block:: bash
+
+ $ ./prepare.sh --stage 2 --stop-stage 5
+
+.. hint::
+
+ If you have pre-downloaded the `LibriSpeech `_
+ dataset and the `musan `_ dataset, say,
+ they are saved in ``/tmp/LibriSpeech`` and ``/tmp/musan``, you can modify
+ the ``dl_dir`` variable in ``./prepare.sh`` to point to ``/tmp`` so that
+ ``./prepare.sh`` won't re-download them.
+
+.. note::
+
+ All generated files by ``./prepare.sh``, e.g., features, lexicon, etc,
+ are saved in ``./data`` directory.
+
+We provide the following YouTube video showing how to run ``./prepare.sh``.
+
+.. note::
+
+ To get the latest news of `next-gen Kaldi `_, please subscribe
+ the following YouTube channel by `Nadira Povey `_:
+
+ ``_
+
+.. youtube:: ofEIoJL-mGM
+
+Training
+--------
+
+For stability, it uses CTC loss for model warm-up and then switches to MMI loss.
+
+Configurable options
+~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./zipformer_mmi/train.py --help
+
+shows you the training options that can be passed from the commandline.
+The following options are used quite often:
+
+ - ``--full-libri``
+
+ If it's True, the training part uses all the training data, i.e.,
+ 960 hours. Otherwise, the training part uses only the subset
+ ``train-clean-100``, which has 100 hours of training data.
+
+ .. CAUTION::
+
+ The training set is perturbed by speed with two factors: 0.9 and 1.1.
+ If ``--full-libri`` is True, each epoch actually processes
+ ``3x960 == 2880`` hours of data.
+
+ - ``--num-epochs``
+
+ It is the number of epochs to train. For instance,
+ ``./zipformer_mmi/train.py --num-epochs 30`` trains for 30 epochs
+ and generates ``epoch-1.pt``, ``epoch-2.pt``, ..., ``epoch-30.pt``
+ in the folder ``./zipformer_mmi/exp``.
+
+ - ``--start-epoch``
+
+ It's used to resume training.
+ ``./zipformer_mmi/train.py --start-epoch 10`` loads the
+ checkpoint ``./zipformer_mmi/exp/epoch-9.pt`` and starts
+ training from epoch 10, based on the state from epoch 9.
+
+ - ``--world-size``
+
+ It is used for multi-GPU single-machine DDP training.
+
+ - (a) If it is 1, then no DDP training is used.
+
+ - (b) If it is 2, then GPU 0 and GPU 1 are used for DDP training.
+
+ The following shows some use cases with it.
+
+ **Use case 1**: You have 4 GPUs, but you only want to use GPU 0 and
+ GPU 2 for training. You can do the following:
+
+ .. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ export CUDA_VISIBLE_DEVICES="0,2"
+ $ ./zipformer_mmi/train.py --world-size 2
+
+ **Use case 2**: You have 4 GPUs and you want to use all of them
+ for training. You can do the following:
+
+ .. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./zipformer_mmi/train.py --world-size 4
+
+ **Use case 3**: You have 4 GPUs but you only want to use GPU 3
+ for training. You can do the following:
+
+ .. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ export CUDA_VISIBLE_DEVICES="3"
+ $ ./zipformer_mmi/train.py --world-size 1
+
+ .. caution::
+
+ Only multi-GPU single-machine DDP training is implemented at present.
+ Multi-GPU multi-machine DDP training will be added later.
+
+ - ``--max-duration``
+
+ It specifies the number of seconds over all utterances in a
+ batch, before **padding**.
+ If you encounter CUDA OOM, please reduce it.
+
+ .. HINT::
+
+ Due to padding, the number of seconds of all utterances in a
+ batch will usually be larger than ``--max-duration``.
+
+ A larger value for ``--max-duration`` may cause OOM during training,
+ while a smaller value may increase the training time. You have to
+ tune it.
+
+
+Pre-configured options
+~~~~~~~~~~~~~~~~~~~~~~
+
+There are some training options, e.g., weight decay,
+number of warmup steps, results dir, etc,
+that are not passed from the commandline.
+They are pre-configured by the function ``get_params()`` in
+`zipformer_mmi/train.py `_
+
+You don't need to change these pre-configured parameters. If you really need to change
+them, please modify ``./zipformer_mmi/train.py`` directly.
+
+Training logs
+~~~~~~~~~~~~~
+
+Training logs and checkpoints are saved in ``zipformer_mmi/exp``.
+You will find the following files in that directory:
+
+ - ``epoch-1.pt``, ``epoch-2.pt``, ...
+
+ These are checkpoint files saved at the end of each epoch, containing model
+ ``state_dict`` and optimizer ``state_dict``.
+ To resume training from some checkpoint, say ``epoch-10.pt``, you can use:
+
+ .. code-block:: bash
+
+ $ ./zipformer_mmi/train.py --start-epoch 11
+
+ - ``checkpoint-436000.pt``, ``checkpoint-438000.pt``, ...
+
+ These are checkpoint files saved every ``--save-every-n`` batches,
+ containing model ``state_dict`` and optimizer ``state_dict``.
+ To resume training from some checkpoint, say ``checkpoint-436000``, you can use:
+
+ .. code-block:: bash
+
+ $ ./zipformer_mmi/train.py --start-batch 436000
+
+ - ``tensorboard/``
+
+ This folder contains tensorBoard logs. Training loss, validation loss, learning
+ rate, etc, are recorded in these logs. You can visualize them by:
+
+ .. code-block:: bash
+
+ $ cd zipformer_mmi/exp/tensorboard
+ $ tensorboard dev upload --logdir . --description "Zipformer MMI training for LibriSpeech with icefall"
+
+ It will print something like below:
+
+ .. code-block::
+
+ TensorFlow installation not found - running with reduced feature set.
+ Upload started and will continue reading any new data as it's added to the logdir.
+
+ To stop uploading, press Ctrl-C.
+
+ New experiment created. View your TensorBoard at: https://tensorboard.dev/experiment/xyOZUKpEQm62HBIlUD4uPA/
+
+ Note there is a URL in the above output. Click it and you will see
+ tensorboard.
+
+ .. hint::
+
+ If you don't have access to google, you can use the following command
+ to view the tensorboard log locally:
+
+ .. code-block:: bash
+
+ cd zipformer_mmi/exp/tensorboard
+ tensorboard --logdir . --port 6008
+
+ It will print the following message:
+
+ .. code-block::
+
+ Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
+ TensorBoard 2.8.0 at http://localhost:6008/ (Press CTRL+C to quit)
+
+ Now start your browser and go to ``_ to view the tensorboard
+ logs.
+
+
+ - ``log/log-train-xxxx``
+
+ It is the detailed training log in text format, same as the one
+ you saw printed to the console during training.
+
+Usage example
+~~~~~~~~~~~~~
+
+You can use the following command to start the training using 4 GPUs:
+
+.. code-block:: bash
+
+ export CUDA_VISIBLE_DEVICES="0,1,2,3"
+ ./zipformer_mmi/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --full-libri 1 \
+ --exp-dir zipformer_mmi/exp \
+ --max-duration 500 \
+ --use-fp16 1 \
+ --num-workers 2
+
+Decoding
+--------
+
+The decoding part uses checkpoints saved by the training part, so you have
+to run the training part first.
+
+.. hint::
+
+ There are two kinds of checkpoints:
+
+ - (1) ``epoch-1.pt``, ``epoch-2.pt``, ..., which are saved at the end
+ of each epoch. You can pass ``--epoch`` to
+ ``zipformer_mmi/decode.py`` to use them.
+
+ - (2) ``checkpoints-436000.pt``, ``epoch-438000.pt``, ..., which are saved
+ every ``--save-every-n`` batches. You can pass ``--iter`` to
+ ``zipformer_mmi/decode.py`` to use them.
+
+ We suggest that you try both types of checkpoints and choose the one
+ that produces the lowest WERs.
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./zipformer_mmi/decode.py --help
+
+shows the options for decoding.
+
+The following shows the example using ``epoch-*.pt``:
+
+.. code-block:: bash
+
+ for m in nbest nbest-rescoring-LG nbest-rescoring-3-gram nbest-rescoring-4-gram; do
+ ./zipformer_mmi/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir ./zipformer_mmi/exp/ \
+ --max-duration 100 \
+ --lang-dir data/lang_bpe_500 \
+ --nbest-scale 1.2 \
+ --hp-scale 1.0 \
+ --decoding-method $m
+ done
+
+
+Export models
+-------------
+
+`zipformer_mmi/export.py `_ supports exporting checkpoints from ``zipformer_mmi/exp`` in the following ways.
+
+Export ``model.state_dict()``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Checkpoints saved by ``zipformer_mmi/train.py`` also include
+``optimizer.state_dict()``. It is useful for resuming training. But after training,
+we are interested only in ``model.state_dict()``. You can use the following
+command to extract ``model.state_dict()``.
+
+.. code-block:: bash
+
+ ./zipformer_mmi/export.py \
+ --exp-dir ./zipformer_mmi/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 9 \
+ --jit 0
+
+It will generate a file ``./zipformer_mmi/exp/pretrained.pt``.
+
+.. hint::
+
+ To use the generated ``pretrained.pt`` for ``zipformer_mmi/decode.py``,
+ you can run:
+
+ .. code-block:: bash
+
+ cd zipformer_mmi/exp
+ ln -s pretrained epoch-9999.pt
+
+ And then pass ``--epoch 9999 --avg 1 --use-averaged-model 0`` to
+ ``./zipformer_mmi/decode.py``.
+
+To use the exported model with ``./zipformer_mmi/pretrained.py``, you
+can run:
+
+.. code-block:: bash
+
+ ./zipformer_mmi/pretrained.py \
+ --checkpoint ./zipformer_mmi/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method 1best \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+Export model using ``torch.jit.script()``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: bash
+
+ ./zipformer_mmi/export.py \
+ --exp-dir ./zipformer_mmi/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 9 \
+ --jit 1
+
+It will generate a file ``cpu_jit.pt`` in the given ``exp_dir``. You can later
+load it by ``torch.jit.load("cpu_jit.pt")``.
+
+Note ``cpu`` in the name ``cpu_jit.pt`` means the parameters when loaded into Python
+are on CPU. You can use ``to("cuda")`` to move them to a CUDA device.
+
+To use the generated files with ``./zipformer_mmi/jit_pretrained.py``:
+
+.. code-block:: bash
+
+ ./zipformer_mmi/jit_pretrained.py \
+ --nn-model-filename ./zipformer_mmi/exp/cpu_jit.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method 1best \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+Download pretrained models
+--------------------------
+
+If you don't want to train from scratch, you can download the pretrained models
+by visiting the following links:
+
+ - ``_
+
+ See ``_
+ for the details of the above pretrained models
diff --git a/docs/source/recipes/timit/index.rst b/docs/source/recipes/Non-streaming-ASR/timit/index.rst
similarity index 100%
rename from docs/source/recipes/timit/index.rst
rename to docs/source/recipes/Non-streaming-ASR/timit/index.rst
diff --git a/docs/source/recipes/timit/tdnn_ligru_ctc.rst b/docs/source/recipes/Non-streaming-ASR/timit/tdnn_ligru_ctc.rst
similarity index 100%
rename from docs/source/recipes/timit/tdnn_ligru_ctc.rst
rename to docs/source/recipes/Non-streaming-ASR/timit/tdnn_ligru_ctc.rst
diff --git a/docs/source/recipes/timit/tdnn_lstm_ctc.rst b/docs/source/recipes/Non-streaming-ASR/timit/tdnn_lstm_ctc.rst
similarity index 100%
rename from docs/source/recipes/timit/tdnn_lstm_ctc.rst
rename to docs/source/recipes/Non-streaming-ASR/timit/tdnn_lstm_ctc.rst
diff --git a/docs/source/recipes/yesno/images/tdnn-tensorboard-log.png b/docs/source/recipes/Non-streaming-ASR/yesno/images/tdnn-tensorboard-log.png
similarity index 100%
rename from docs/source/recipes/yesno/images/tdnn-tensorboard-log.png
rename to docs/source/recipes/Non-streaming-ASR/yesno/images/tdnn-tensorboard-log.png
diff --git a/docs/source/recipes/yesno/index.rst b/docs/source/recipes/Non-streaming-ASR/yesno/index.rst
similarity index 100%
rename from docs/source/recipes/yesno/index.rst
rename to docs/source/recipes/Non-streaming-ASR/yesno/index.rst
diff --git a/docs/source/recipes/yesno/tdnn.rst b/docs/source/recipes/Non-streaming-ASR/yesno/tdnn.rst
similarity index 100%
rename from docs/source/recipes/yesno/tdnn.rst
rename to docs/source/recipes/Non-streaming-ASR/yesno/tdnn.rst
diff --git a/docs/source/recipes/Streaming-ASR/index.rst b/docs/source/recipes/Streaming-ASR/index.rst
new file mode 100644
index 000000000..8c0ffe447
--- /dev/null
+++ b/docs/source/recipes/Streaming-ASR/index.rst
@@ -0,0 +1,12 @@
+Streaming ASR
+=============
+
+.. toctree::
+ :maxdepth: 1
+
+ introduction
+
+.. toctree::
+ :maxdepth: 2
+
+ librispeech/index
diff --git a/docs/source/recipes/Streaming-ASR/introduction.rst b/docs/source/recipes/Streaming-ASR/introduction.rst
new file mode 100644
index 000000000..d81156659
--- /dev/null
+++ b/docs/source/recipes/Streaming-ASR/introduction.rst
@@ -0,0 +1,52 @@
+Introduction
+============
+
+This page shows you how we implement streaming **X-former transducer** models for ASR.
+
+.. HINT::
+ X-former transducer here means the encoder of the transducer model uses Multi-Head Attention,
+ like `Conformer `_, `EmFormer `_ etc.
+
+Currently we have implemented two types of streaming models, one uses Conformer as encoder, the other uses Emformer as encoder.
+
+Streaming Conformer
+-------------------
+
+The main idea of training a streaming model is to make the model see limited contexts
+in training time, we can achieve this by applying a mask to the output of self-attention.
+In icefall, we implement the streaming conformer the way just like what `WeNet `_ did.
+
+.. NOTE::
+ The conformer-transducer recipes in LibriSpeech datasets, like, `pruned_transducer_stateless `_,
+ `pruned_transducer_stateless2 `_,
+ `pruned_transducer_stateless3 `_,
+ `pruned_transducer_stateless4 `_,
+ `pruned_transducer_stateless5 `_
+ all support streaming.
+
+.. NOTE::
+ Training a streaming conformer model in ``icefall`` is almost the same as training a
+ non-streaming model, all you need to do is passing several extra arguments.
+ See :doc:`Pruned transducer statelessX ` for more details.
+
+.. HINT::
+ If you want to adapt a non-streaming conformer model to be streaming, please refer
+ to `this pull request `_.
+
+
+Streaming Emformer
+------------------
+
+The Emformer model proposed `here `_ uses more
+complicated techniques. It has a memory bank component to memorize history information,
+what' more, it also introduces right context in training time by hard-copying part of
+the input features.
+
+We have three variants of Emformer models in ``icefall``.
+
+ - ``pruned_stateless_emformer_rnnt2`` using Emformer from torchaudio, see `LibriSpeech recipe `_.
+ - ``conv_emformer_transducer_stateless`` using ConvEmformer implemented by ourself. Different from the Emformer in torchaudio,
+ ConvEmformer has a convolution in each layer and uses the mechanisms in our reworked conformer model.
+ See `LibriSpeech recipe `_.
+ - ``conv_emformer_transducer_stateless2`` using ConvEmformer implemented by ourself. The only difference from the above one is that
+ it uses a simplified memory bank. See `LibriSpeech recipe `_.
diff --git a/docs/source/recipes/librispeech/images/librispeech-lstm-transducer-tensorboard-log.png b/docs/source/recipes/Streaming-ASR/librispeech/images/librispeech-lstm-transducer-tensorboard-log.png
similarity index 100%
rename from docs/source/recipes/librispeech/images/librispeech-lstm-transducer-tensorboard-log.png
rename to docs/source/recipes/Streaming-ASR/librispeech/images/librispeech-lstm-transducer-tensorboard-log.png
diff --git a/docs/source/recipes/Streaming-ASR/librispeech/images/streaming-librispeech-pruned-transducer-tensorboard-log.jpg b/docs/source/recipes/Streaming-ASR/librispeech/images/streaming-librispeech-pruned-transducer-tensorboard-log.jpg
new file mode 100644
index 000000000..9c77b8bae
Binary files /dev/null and b/docs/source/recipes/Streaming-ASR/librispeech/images/streaming-librispeech-pruned-transducer-tensorboard-log.jpg differ
diff --git a/docs/source/recipes/librispeech/index.rst b/docs/source/recipes/Streaming-ASR/librispeech/index.rst
similarity index 61%
rename from docs/source/recipes/librispeech/index.rst
rename to docs/source/recipes/Streaming-ASR/librispeech/index.rst
index 6c91b6750..d52e08058 100644
--- a/docs/source/recipes/librispeech/index.rst
+++ b/docs/source/recipes/Streaming-ASR/librispeech/index.rst
@@ -4,6 +4,8 @@ LibriSpeech
.. toctree::
:maxdepth: 1
- tdnn_lstm_ctc
- conformer_ctc
+ pruned_transducer_stateless
+
lstm_pruned_stateless_transducer
+
+ zipformer_transducer
diff --git a/docs/source/recipes/librispeech/lstm_pruned_stateless_transducer.rst b/docs/source/recipes/Streaming-ASR/librispeech/lstm_pruned_stateless_transducer.rst
similarity index 95%
rename from docs/source/recipes/librispeech/lstm_pruned_stateless_transducer.rst
rename to docs/source/recipes/Streaming-ASR/librispeech/lstm_pruned_stateless_transducer.rst
index 643855cc2..ce8ba1453 100644
--- a/docs/source/recipes/librispeech/lstm_pruned_stateless_transducer.rst
+++ b/docs/source/recipes/Streaming-ASR/librispeech/lstm_pruned_stateless_transducer.rst
@@ -515,10 +515,10 @@ To use the generated files with ``./lstm_transducer_stateless2/jit_pretrained``:
Please see ``_
for how to use the exported models in ``sherpa``.
-.. _export-model-for-ncnn:
+.. _export-lstm-transducer-model-for-ncnn:
-Export model for ncnn
-~~~~~~~~~~~~~~~~~~~~~
+Export LSTM transducer models for ncnn
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We support exporting pretrained LSTM transducer models to
`ncnn `_ using
@@ -531,16 +531,36 @@ First, let us install a modified version of ``ncnn``:
git clone https://github.com/csukuangfj/ncnn
cd ncnn
git submodule update --recursive --init
- python3 setup.py bdist_wheel
- ls -lh dist/
- pip install ./dist/*.whl
+
+ # Note: We don't use "python setup.py install" or "pip install ." here
+
+ mkdir -p build-wheel
+ cd build-wheel
+
+ cmake \
+ -DCMAKE_BUILD_TYPE=Release \
+ -DNCNN_PYTHON=ON \
+ -DNCNN_BUILD_BENCHMARK=OFF \
+ -DNCNN_BUILD_EXAMPLES=OFF \
+ -DNCNN_BUILD_TOOLS=ON \
+ ..
+
+ make -j4
+
+ cd ..
+
+ # Note: $PWD here is /path/to/ncnn
+
+ export PYTHONPATH=$PWD/python:$PYTHONPATH
+ export PATH=$PWD/tools/pnnx/build/src:$PATH
+ export PATH=$PWD/build-wheel/tools/quantize:$PATH
# now build pnnx
cd tools/pnnx
mkdir build
cd build
+ cmake ..
make -j4
- export PATH=$PWD/src:$PATH
./src/pnnx
@@ -549,6 +569,9 @@ First, let us install a modified version of ``ncnn``:
We assume that you have added the path to the binary ``pnnx`` to the
environment variable ``PATH``.
+ We also assume that you have added ``build/tools/quantize`` to the environment
+ variable ``PATH`` so that you are able to use ``ncnn2int8`` later.
+
Second, let us export the model using ``torch.jit.trace()`` that is suitable
for ``pnnx``:
@@ -634,3 +657,6 @@ by visiting the following links:
You can find more usages of the pretrained models in
``_
+
+Export ConvEmformer transducer models for ncnn
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/docs/source/recipes/Streaming-ASR/librispeech/pruned_transducer_stateless.rst b/docs/source/recipes/Streaming-ASR/librispeech/pruned_transducer_stateless.rst
new file mode 100644
index 000000000..de7102ba8
--- /dev/null
+++ b/docs/source/recipes/Streaming-ASR/librispeech/pruned_transducer_stateless.rst
@@ -0,0 +1,735 @@
+Pruned transducer statelessX
+============================
+
+This tutorial shows you how to run a **streaming** conformer transducer model
+with the `LibriSpeech `_ dataset.
+
+.. Note::
+
+ The tutorial is suitable for `pruned_transducer_stateless `_,
+ `pruned_transducer_stateless2 `_,
+ `pruned_transducer_stateless4 `_,
+ `pruned_transducer_stateless5 `_,
+ We will take pruned_transducer_stateless4 as an example in this tutorial.
+
+.. HINT::
+
+ We assume you have read the page :ref:`install icefall` and have setup
+ the environment for ``icefall``.
+
+.. HINT::
+
+ We recommend you to use a GPU or several GPUs to run this recipe.
+
+.. hint::
+
+ Please scroll down to the bottom of this page to find download links
+ for pretrained models if you don't want to train a model from scratch.
+
+
+We use pruned RNN-T to compute the loss.
+
+.. note::
+
+ You can find the paper about pruned RNN-T at the following address:
+
+ ``_
+
+The transducer model consists of 3 parts:
+
+ - Encoder, a.k.a, the transcription network. We use a Conformer model (the reworked version by Daniel Povey)
+ - Decoder, a.k.a, the prediction network. We use a stateless model consisting of
+ ``nn.Embedding`` and ``nn.Conv1d``
+ - Joiner, a.k.a, the joint network.
+
+.. caution::
+
+ Contrary to the conventional RNN-T models, we use a stateless decoder.
+ That is, it has no recurrent connections.
+
+
+Data preparation
+----------------
+
+.. hint::
+
+ The data preparation is the same as other recipes on LibriSpeech dataset,
+ if you have finished this step, you can skip to ``Training`` directly.
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./prepare.sh
+
+The script ``./prepare.sh`` handles the data preparation for you, **automagically**.
+All you need to do is to run it.
+
+The data preparation contains several stages, you can use the following two
+options:
+
+ - ``--stage``
+ - ``--stop-stage``
+
+to control which stage(s) should be run. By default, all stages are executed.
+
+
+For example,
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./prepare.sh --stage 0 --stop-stage 0
+
+means to run only stage 0.
+
+To run stage 2 to stage 5, use:
+
+.. code-block:: bash
+
+ $ ./prepare.sh --stage 2 --stop-stage 5
+
+.. HINT::
+
+ If you have pre-downloaded the `LibriSpeech `_
+ dataset and the `musan `_ dataset, say,
+ they are saved in ``/tmp/LibriSpeech`` and ``/tmp/musan``, you can modify
+ the ``dl_dir`` variable in ``./prepare.sh`` to point to ``/tmp`` so that
+ ``./prepare.sh`` won't re-download them.
+
+.. NOTE::
+
+ All generated files by ``./prepare.sh``, e.g., features, lexicon, etc,
+ are saved in ``./data`` directory.
+
+We provide the following YouTube video showing how to run ``./prepare.sh``.
+
+.. note::
+
+ To get the latest news of `next-gen Kaldi `_, please subscribe
+ the following YouTube channel by `Nadira Povey `_:
+
+ ``_
+
+.. youtube:: ofEIoJL-mGM
+
+
+Training
+--------
+
+.. NOTE::
+
+ We put the streaming and non-streaming model in one recipe, to train a streaming model you only
+ need to add **4** extra options comparing with training a non-streaming model. These options are
+ ``--dynamic-chunk-training``, ``--num-left-chunks``, ``--causal-convolution``, ``--short-chunk-size``.
+ You can see the configurable options below for their meanings or read https://arxiv.org/pdf/2012.05481.pdf for more details.
+
+Configurable options
+~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./pruned_transducer_stateless4/train.py --help
+
+
+shows you the training options that can be passed from the commandline.
+The following options are used quite often:
+
+ - ``--exp-dir``
+
+ The directory to save checkpoints, training logs and tensorboard.
+
+ - ``--full-libri``
+
+ If it's True, the training part uses all the training data, i.e.,
+ 960 hours. Otherwise, the training part uses only the subset
+ ``train-clean-100``, which has 100 hours of training data.
+
+ .. CAUTION::
+ The training set is perturbed by speed with two factors: 0.9 and 1.1.
+ If ``--full-libri`` is True, each epoch actually processes
+ ``3x960 == 2880`` hours of data.
+
+ - ``--num-epochs``
+
+ It is the number of epochs to train. For instance,
+ ``./pruned_transducer_stateless4/train.py --num-epochs 30`` trains for 30 epochs
+ and generates ``epoch-1.pt``, ``epoch-2.pt``, ..., ``epoch-30.pt``
+ in the folder ``./pruned_transducer_stateless4/exp``.
+
+ - ``--start-epoch``
+
+ It's used to resume training.
+ ``./pruned_transducer_stateless4/train.py --start-epoch 10`` loads the
+ checkpoint ``./pruned_transducer_stateless4/exp/epoch-9.pt`` and starts
+ training from epoch 10, based on the state from epoch 9.
+
+ - ``--world-size``
+
+ It is used for multi-GPU single-machine DDP training.
+
+ - (a) If it is 1, then no DDP training is used.
+
+ - (b) If it is 2, then GPU 0 and GPU 1 are used for DDP training.
+
+ The following shows some use cases with it.
+
+ **Use case 1**: You have 4 GPUs, but you only want to use GPU 0 and
+ GPU 2 for training. You can do the following:
+
+ .. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ export CUDA_VISIBLE_DEVICES="0,2"
+ $ ./pruned_transducer_stateless4/train.py --world-size 2
+
+ **Use case 2**: You have 4 GPUs and you want to use all of them
+ for training. You can do the following:
+
+ .. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./pruned_transducer_stateless4/train.py --world-size 4
+
+ **Use case 3**: You have 4 GPUs but you only want to use GPU 3
+ for training. You can do the following:
+
+ .. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ export CUDA_VISIBLE_DEVICES="3"
+ $ ./pruned_transducer_stateless4/train.py --world-size 1
+
+ .. caution::
+
+ Only multi-GPU single-machine DDP training is implemented at present.
+ Multi-GPU multi-machine DDP training will be added later.
+
+ - ``--max-duration``
+
+ It specifies the number of seconds over all utterances in a
+ batch, before **padding**.
+ If you encounter CUDA OOM, please reduce it.
+
+ .. HINT::
+
+ Due to padding, the number of seconds of all utterances in a
+ batch will usually be larger than ``--max-duration``.
+
+ A larger value for ``--max-duration`` may cause OOM during training,
+ while a smaller value may increase the training time. You have to
+ tune it.
+
+ - ``--use-fp16``
+
+ If it is True, the model will train with half precision, from our experiment
+ results, by using half precision you can train with two times larger ``--max-duration``
+ so as to get almost 2X speed up.
+
+ - ``--dynamic-chunk-training``
+
+ The flag that indicates whether to train a streaming model or not, it
+ **MUST** be True if you want to train a streaming model.
+
+ - ``--short-chunk-size``
+
+ When training a streaming attention model with chunk masking, the chunk size
+ would be either max sequence length of current batch or uniformly sampled from
+ (1, short_chunk_size). The default value is 25, you don't have to change it most of the time.
+
+ - ``--num-left-chunks``
+
+ It indicates how many left context (in chunks) that can be seen when calculating attention.
+ The default value is 4, you don't have to change it most of the time.
+
+
+ - ``--causal-convolution``
+
+ Whether to use causal convolution in conformer encoder layer, this requires
+ to be True when training a streaming model.
+
+
+Pre-configured options
+~~~~~~~~~~~~~~~~~~~~~~
+
+There are some training options, e.g., number of encoder layers,
+encoder dimension, decoder dimension, number of warmup steps etc,
+that are not passed from the commandline.
+They are pre-configured by the function ``get_params()`` in
+`pruned_transducer_stateless4/train.py `_
+
+You don't need to change these pre-configured parameters. If you really need to change
+them, please modify ``./pruned_transducer_stateless4/train.py`` directly.
+
+
+.. NOTE::
+
+ The options for `pruned_transducer_stateless5 `_ are a little different from
+ other recipes. It allows you to configure ``--num-encoder-layers``, ``--dim-feedforward``, ``--nhead``, ``--encoder-dim``, ``--decoder-dim``, ``--joiner-dim`` from commandline, so that you can train models with different size with pruned_transducer_stateless5.
+
+
+Training logs
+~~~~~~~~~~~~~
+
+Training logs and checkpoints are saved in ``--exp-dir`` (e.g. ``pruned_transducer_stateless4/exp``.
+You will find the following files in that directory:
+
+ - ``epoch-1.pt``, ``epoch-2.pt``, ...
+
+ These are checkpoint files saved at the end of each epoch, containing model
+ ``state_dict`` and optimizer ``state_dict``.
+ To resume training from some checkpoint, say ``epoch-10.pt``, you can use:
+
+ .. code-block:: bash
+
+ $ ./pruned_transducer_stateless4/train.py --start-epoch 11
+
+ - ``checkpoint-436000.pt``, ``checkpoint-438000.pt``, ...
+
+ These are checkpoint files saved every ``--save-every-n`` batches,
+ containing model ``state_dict`` and optimizer ``state_dict``.
+ To resume training from some checkpoint, say ``checkpoint-436000``, you can use:
+
+ .. code-block:: bash
+
+ $ ./pruned_transducer_stateless4/train.py --start-batch 436000
+
+ - ``tensorboard/``
+
+ This folder contains tensorBoard logs. Training loss, validation loss, learning
+ rate, etc, are recorded in these logs. You can visualize them by:
+
+ .. code-block:: bash
+
+ $ cd pruned_transducer_stateless4/exp/tensorboard
+ $ tensorboard dev upload --logdir . --description "pruned transducer training for LibriSpeech with icefall"
+
+ It will print something like below:
+
+ .. code-block::
+
+ TensorFlow installation not found - running with reduced feature set.
+ Upload started and will continue reading any new data as it's added to the logdir.
+
+ To stop uploading, press Ctrl-C.
+
+ New experiment created. View your TensorBoard at: https://tensorboard.dev/experiment/97VKXf80Ru61CnP2ALWZZg/
+
+ [2022-11-20T15:50:50] Started scanning logdir.
+ Uploading 4468 scalars...
+ [2022-11-20T15:53:02] Total uploaded: 210171 scalars, 0 tensors, 0 binary objects
+ Listening for new data in logdir...
+
+ Note there is a URL in the above output. Click it and you will see
+ the following screenshot:
+
+ .. figure:: images/streaming-librispeech-pruned-transducer-tensorboard-log.jpg
+ :width: 600
+ :alt: TensorBoard screenshot
+ :align: center
+ :target: https://tensorboard.dev/experiment/97VKXf80Ru61CnP2ALWZZg/
+
+ TensorBoard screenshot.
+
+ .. hint::
+
+ If you don't have access to google, you can use the following command
+ to view the tensorboard log locally:
+
+ .. code-block:: bash
+
+ cd pruned_transducer_stateless4/exp/tensorboard
+ tensorboard --logdir . --port 6008
+
+ It will print the following message:
+
+ .. code-block::
+
+ Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
+ TensorBoard 2.8.0 at http://localhost:6008/ (Press CTRL+C to quit)
+
+ Now start your browser and go to ``_ to view the tensorboard
+ logs.
+
+
+ - ``log/log-train-xxxx``
+
+ It is the detailed training log in text format, same as the one
+ you saw printed to the console during training.
+
+Usage example
+~~~~~~~~~~~~~
+
+You can use the following command to start the training using 4 GPUs:
+
+.. code-block:: bash
+
+ export CUDA_VISIBLE_DEVICES="0,1,2,3"
+ ./pruned_transducer_stateless4/train.py \
+ --world-size 4 \
+ --dynamic-chunk-training 1 \
+ --causal-convolution 1 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --exp-dir pruned_transducer_stateless4/exp \
+ --full-libri 1 \
+ --max-duration 300
+
+.. NOTE::
+
+ Comparing with training a non-streaming model, you only need to add two extra options,
+ ``--dynamic-chunk-training 1`` and ``--causal-convolution 1`` .
+
+
+Decoding
+--------
+
+The decoding part uses checkpoints saved by the training part, so you have
+to run the training part first.
+
+.. hint::
+
+ There are two kinds of checkpoints:
+
+ - (1) ``epoch-1.pt``, ``epoch-2.pt``, ..., which are saved at the end
+ of each epoch. You can pass ``--epoch`` to
+ ``pruned_transducer_stateless4/decode.py`` to use them.
+
+ - (2) ``checkpoints-436000.pt``, ``epoch-438000.pt``, ..., which are saved
+ every ``--save-every-n`` batches. You can pass ``--iter`` to
+ ``pruned_transducer_stateless4/decode.py`` to use them.
+
+ We suggest that you try both types of checkpoints and choose the one
+ that produces the lowest WERs.
+
+.. tip::
+
+ To decode a streaming model, you can use either ``simulate streaming decoding`` in ``decode.py`` or
+ ``real streaming decoding`` in ``streaming_decode.py``, the difference between ``decode.py`` and
+ ``streaming_decode.py`` is that, ``decode.py`` processes the whole acoustic frames at one time with masking (i.e. same as training),
+ but ``streaming_decode.py`` processes the acoustic frames chunk by chunk (so it can only see limited context).
+
+.. NOTE::
+
+ ``simulate streaming decoding`` in ``decode.py`` and ``real streaming decoding`` in ``streaming_decode.py`` should
+ produce almost the same results given the same ``--decode-chunk-size`` and ``--left-context``.
+
+
+Simulate streaming decoding
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./pruned_transducer_stateless4/decode.py --help
+
+shows the options for decoding.
+The following options are important for streaming models:
+
+ ``--simulate-streaming``
+
+ If you want to decode a streaming model with ``decode.py``, you **MUST** set
+ ``--simulate-streaming`` to ``True``. ``simulate`` here means the acoustic frames
+ are not processed frame by frame (or chunk by chunk), instead, the whole sequence
+ is processed at one time with masking (the same as training).
+
+ ``--causal-convolution``
+
+ If True, the convolution module in encoder layers will be causal convolution.
+ This is **MUST** be True when decoding with a streaming model.
+
+ ``--decode-chunk-size``
+
+ For streaming models, we will calculate the chunk-wise attention, ``--decode-chunk-size``
+ indicates the chunk length (in frames after subsampling) for chunk-wise attention.
+ For ``simulate streaming decoding`` the ``decode-chunk-size`` is used to generate
+ the attention mask.
+
+ ``--left-context``
+
+ ``--left-context`` indicates how many left context frames (after subsampling) can be seen
+ for current chunk when calculating chunk-wise attention. Normally, ``left-context`` should equal
+ to ``decode-chunk-size * num-left-chunks``, where ``num-left-chunks`` is the option used
+ to train this model. For ``simulate streaming decoding`` the ``left-context`` is used to generate
+ the attention mask.
+
+
+The following shows two examples (for the two types of checkpoints):
+
+.. code-block:: bash
+
+ for m in greedy_search fast_beam_search modified_beam_search; do
+ for epoch in 25 20; do
+ for avg in 7 5 3 1; do
+ ./pruned_transducer_stateless4/decode.py \
+ --epoch $epoch \
+ --avg $avg \
+ --simulate-streaming 1 \
+ --causal-convolution 1 \
+ --decode-chunk-size 16 \
+ --left-context 64 \
+ --exp-dir pruned_transducer_stateless4/exp \
+ --max-duration 600 \
+ --decoding-method $m
+ done
+ done
+ done
+
+
+.. code-block:: bash
+
+ for m in greedy_search fast_beam_search modified_beam_search; do
+ for iter in 474000; do
+ for avg in 8 10 12 14 16 18; do
+ ./pruned_transducer_stateless4/decode.py \
+ --iter $iter \
+ --avg $avg \
+ --simulate-streaming 1 \
+ --causal-convolution 1 \
+ --decode-chunk-size 16 \
+ --left-context 64 \
+ --exp-dir pruned_transducer_stateless4/exp \
+ --max-duration 600 \
+ --decoding-method $m
+ done
+ done
+ done
+
+
+Real streaming decoding
+~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./pruned_transducer_stateless4/streaming_decode.py --help
+
+shows the options for decoding.
+The following options are important for streaming models:
+
+ ``--decode-chunk-size``
+
+ For streaming models, we will calculate the chunk-wise attention, ``--decode-chunk-size``
+ indicates the chunk length (in frames after subsampling) for chunk-wise attention.
+ For ``real streaming decoding``, we will process ``decode-chunk-size`` acoustic frames at each time.
+
+ ``--left-context``
+
+ ``--left-context`` indicates how many left context frames (after subsampling) can be seen
+ for current chunk when calculating chunk-wise attention. Normally, ``left-context`` should equal
+ to ``decode-chunk-size * num-left-chunks``, where ``num-left-chunks`` is the option used
+ to train this model.
+
+ ``--num-decode-streams``
+
+ The number of decoding streams that can be run in parallel (very similar to the ``bath size``).
+ For ``real streaming decoding``, the batches will be packed dynamically, for example, if the
+ ``num-decode-streams`` equals to 10, then, sequence 1 to 10 will be decoded at first, after a while,
+ suppose sequence 1 and 2 are done, so, sequence 3 to 12 will be processed parallelly in a batch.
+
+
+.. NOTE::
+
+ We also try adding ``--right-context`` in the real streaming decoding, but it seems not to benefit
+ the performance for all the models, the reasons might be the training and decoding mismatch. You
+ can try decoding with ``--right-context`` to see if it helps. The default value is 0.
+
+
+The following shows two examples (for the two types of checkpoints):
+
+.. code-block:: bash
+
+ for m in greedy_search fast_beam_search modified_beam_search; do
+ for epoch in 25 20; do
+ for avg in 7 5 3 1; do
+ ./pruned_transducer_stateless4/decode.py \
+ --epoch $epoch \
+ --avg $avg \
+ --decode-chunk-size 16 \
+ --left-context 64 \
+ --num-decode-streams 100 \
+ --exp-dir pruned_transducer_stateless4/exp \
+ --max-duration 600 \
+ --decoding-method $m
+ done
+ done
+ done
+
+
+.. code-block:: bash
+
+ for m in greedy_search fast_beam_search modified_beam_search; do
+ for iter in 474000; do
+ for avg in 8 10 12 14 16 18; do
+ ./pruned_transducer_stateless4/decode.py \
+ --iter $iter \
+ --avg $avg \
+ --decode-chunk-size 16 \
+ --left-context 64 \
+ --num-decode-streams 100 \
+ --exp-dir pruned_transducer_stateless4/exp \
+ --max-duration 600 \
+ --decoding-method $m
+ done
+ done
+ done
+
+
+.. tip::
+
+ Supporting decoding methods are as follows:
+
+ - ``greedy_search`` : It takes the symbol with largest posterior probability
+ of each frame as the decoding result.
+
+ - ``beam_search`` : It implements Algorithm 1 in https://arxiv.org/pdf/1211.3711.pdf and
+ `espnet/nets/beam_search_transducer.py `_
+ is used as a reference. Basicly, it keeps topk states for each frame, and expands the kept states with their own contexts to
+ next frame.
+
+ - ``modified_beam_search`` : It implements the same algorithm as ``beam_search`` above, but it
+ runs in batch mode with ``--max-sym-per-frame=1`` being hardcoded.
+
+ - ``fast_beam_search`` : It implements graph composition between the output ``log_probs`` and
+ given ``FSAs``. It is hard to describe the details in several lines of texts, you can read
+ our paper in https://arxiv.org/pdf/2211.00484.pdf or our `rnnt decode code in k2 `_. ``fast_beam_search`` can decode with ``FSAs`` on GPU efficiently.
+
+ - ``fast_beam_search_LG`` : The same as ``fast_beam_search`` above, ``fast_beam_search`` uses
+ an trivial graph that has only one state, while ``fast_beam_search_LG`` uses an LG graph
+ (with N-gram LM).
+
+ - ``fast_beam_search_nbest`` : It produces the decoding results as follows:
+
+ - (1) Use ``fast_beam_search`` to get a lattice
+ - (2) Select ``num_paths`` paths from the lattice using ``k2.random_paths()``
+ - (3) Unique the selected paths
+ - (4) Intersect the selected paths with the lattice and compute the
+ shortest path from the intersection result
+ - (5) The path with the largest score is used as the decoding output.
+
+ - ``fast_beam_search_nbest_LG`` : It implements same logic as ``fast_beam_search_nbest``, the
+ only difference is that it uses ``fast_beam_search_LG`` to generate the lattice.
+
+.. NOTE::
+
+ The supporting decoding methods in ``streaming_decode.py`` might be less than that in ``decode.py``, if needed,
+ you can implement them by yourself or file a issue in `icefall `_ .
+
+
+Export Model
+------------
+
+`pruned_transducer_stateless4/export.py `_ supports exporting checkpoints from ``pruned_transducer_stateless4/exp`` in the following ways.
+
+Export ``model.state_dict()``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Checkpoints saved by ``pruned_transducer_stateless4/train.py`` also include
+``optimizer.state_dict()``. It is useful for resuming training. But after training,
+we are interested only in ``model.state_dict()``. You can use the following
+command to extract ``model.state_dict()``.
+
+.. code-block:: bash
+
+ # Assume that --epoch 25 --avg 3 produces the smallest WER
+ # (You can get such information after running ./pruned_transducer_stateless4/decode.py)
+
+ epoch=25
+ avg=3
+
+ ./pruned_transducer_stateless4/export.py \
+ --exp-dir ./pruned_transducer_stateless4/exp \
+ --streaming-model 1 \
+ --causal-convolution 1 \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch $epoch \
+ --avg $avg
+
+.. caution::
+
+ ``--streaming-model`` and ``--causal-convolution`` require to be True to export
+ a streaming mdoel.
+
+It will generate a file ``./pruned_transducer_stateless4/exp/pretrained.pt``.
+
+.. hint::
+
+ To use the generated ``pretrained.pt`` for ``pruned_transducer_stateless4/decode.py``,
+ you can run:
+
+ .. code-block:: bash
+
+ cd pruned_transducer_stateless4/exp
+ ln -s pretrained.pt epoch-999.pt
+
+ And then pass ``--epoch 999 --avg 1 --use-averaged-model 0`` to
+ ``./pruned_transducer_stateless4/decode.py``.
+
+To use the exported model with ``./pruned_transducer_stateless4/pretrained.py``, you
+can run:
+
+.. code-block:: bash
+
+ ./pruned_transducer_stateless4/pretrained.py \
+ --checkpoint ./pruned_transducer_stateless4/exp/pretrained.pt \
+ --simulate-streaming 1 \
+ --causal-convolution 1 \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method greedy_search \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+
+Export model using ``torch.jit.script()``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: bash
+
+ ./pruned_transducer_stateless4/export.py \
+ --exp-dir ./pruned_transducer_stateless4/exp \
+ --streaming-model 1 \
+ --causal-convolution 1 \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 25 \
+ --avg 3 \
+ --jit 1
+
+.. caution::
+
+ ``--streaming-model`` and ``--causal-convolution`` require to be True to export
+ a streaming mdoel.
+
+It will generate a file ``cpu_jit.pt`` in the given ``exp_dir``. You can later
+load it by ``torch.jit.load("cpu_jit.pt")``.
+
+Note ``cpu`` in the name ``cpu_jit.pt`` means the parameters when loaded into Python
+are on CPU. You can use ``to("cuda")`` to move them to a CUDA device.
+
+.. NOTE::
+
+ You will need this ``cpu_jit.pt`` when deploying with Sherpa framework.
+
+
+Download pretrained models
+--------------------------
+
+If you don't want to train from scratch, you can download the pretrained models
+by visiting the following links:
+
+ - `pruned_transducer_stateless `_
+
+ - `pruned_transducer_stateless2 `_
+
+ - `pruned_transducer_stateless4 `_
+
+ - `pruned_transducer_stateless5 `_
+
+ See ``_
+ for the details of the above pretrained models
+
+
+Deploy with Sherpa
+------------------
+
+Please see ``_
+for how to deploy the models in ``sherpa``.
diff --git a/docs/source/recipes/Streaming-ASR/librispeech/zipformer_transducer.rst b/docs/source/recipes/Streaming-ASR/librispeech/zipformer_transducer.rst
new file mode 100644
index 000000000..f0e8961d7
--- /dev/null
+++ b/docs/source/recipes/Streaming-ASR/librispeech/zipformer_transducer.rst
@@ -0,0 +1,654 @@
+Zipformer Transducer
+====================
+
+This tutorial shows you how to run a **streaming** zipformer transducer model
+with the `LibriSpeech `_ dataset.
+
+.. Note::
+
+ The tutorial is suitable for `pruned_transducer_stateless7_streaming `_,
+
+.. HINT::
+
+ We assume you have read the page :ref:`install icefall` and have setup
+ the environment for ``icefall``.
+
+.. HINT::
+
+ We recommend you to use a GPU or several GPUs to run this recipe.
+
+.. hint::
+
+ Please scroll down to the bottom of this page to find download links
+ for pretrained models if you don't want to train a model from scratch.
+
+
+We use pruned RNN-T to compute the loss.
+
+.. note::
+
+ You can find the paper about pruned RNN-T at the following address:
+
+ ``_
+
+The transducer model consists of 3 parts:
+
+ - Encoder, a.k.a, the transcription network. We use a Zipformer model (proposed by Daniel Povey)
+ - Decoder, a.k.a, the prediction network. We use a stateless model consisting of
+ ``nn.Embedding`` and ``nn.Conv1d``
+ - Joiner, a.k.a, the joint network.
+
+.. caution::
+
+ Contrary to the conventional RNN-T models, we use a stateless decoder.
+ That is, it has no recurrent connections.
+
+
+Data preparation
+----------------
+
+.. hint::
+
+ The data preparation is the same as other recipes on LibriSpeech dataset,
+ if you have finished this step, you can skip to ``Training`` directly.
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./prepare.sh
+
+The script ``./prepare.sh`` handles the data preparation for you, **automagically**.
+All you need to do is to run it.
+
+The data preparation contains several stages, you can use the following two
+options:
+
+ - ``--stage``
+ - ``--stop-stage``
+
+to control which stage(s) should be run. By default, all stages are executed.
+
+
+For example,
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./prepare.sh --stage 0 --stop-stage 0
+
+means to run only stage 0.
+
+To run stage 2 to stage 5, use:
+
+.. code-block:: bash
+
+ $ ./prepare.sh --stage 2 --stop-stage 5
+
+.. HINT::
+
+ If you have pre-downloaded the `LibriSpeech `_
+ dataset and the `musan `_ dataset, say,
+ they are saved in ``/tmp/LibriSpeech`` and ``/tmp/musan``, you can modify
+ the ``dl_dir`` variable in ``./prepare.sh`` to point to ``/tmp`` so that
+ ``./prepare.sh`` won't re-download them.
+
+.. NOTE::
+
+ All generated files by ``./prepare.sh``, e.g., features, lexicon, etc,
+ are saved in ``./data`` directory.
+
+We provide the following YouTube video showing how to run ``./prepare.sh``.
+
+.. note::
+
+ To get the latest news of `next-gen Kaldi `_, please subscribe
+ the following YouTube channel by `Nadira Povey `_:
+
+ ``_
+
+.. youtube:: ofEIoJL-mGM
+
+
+Training
+--------
+
+Configurable options
+~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./pruned_transducer_stateless7_streaming/train.py --help
+
+
+shows you the training options that can be passed from the commandline.
+The following options are used quite often:
+
+ - ``--exp-dir``
+
+ The directory to save checkpoints, training logs and tensorboard.
+
+ - ``--full-libri``
+
+ If it's True, the training part uses all the training data, i.e.,
+ 960 hours. Otherwise, the training part uses only the subset
+ ``train-clean-100``, which has 100 hours of training data.
+
+ .. CAUTION::
+ The training set is perturbed by speed with two factors: 0.9 and 1.1.
+ If ``--full-libri`` is True, each epoch actually processes
+ ``3x960 == 2880`` hours of data.
+
+ - ``--num-epochs``
+
+ It is the number of epochs to train. For instance,
+ ``./pruned_transducer_stateless7_streaming/train.py --num-epochs 30`` trains for 30 epochs
+ and generates ``epoch-1.pt``, ``epoch-2.pt``, ..., ``epoch-30.pt``
+ in the folder ``./pruned_transducer_stateless7_streaming/exp``.
+
+ - ``--start-epoch``
+
+ It's used to resume training.
+ ``./pruned_transducer_stateless7_streaming/train.py --start-epoch 10`` loads the
+ checkpoint ``./pruned_transducer_stateless7_streaming/exp/epoch-9.pt`` and starts
+ training from epoch 10, based on the state from epoch 9.
+
+ - ``--world-size``
+
+ It is used for multi-GPU single-machine DDP training.
+
+ - (a) If it is 1, then no DDP training is used.
+
+ - (b) If it is 2, then GPU 0 and GPU 1 are used for DDP training.
+
+ The following shows some use cases with it.
+
+ **Use case 1**: You have 4 GPUs, but you only want to use GPU 0 and
+ GPU 2 for training. You can do the following:
+
+ .. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ export CUDA_VISIBLE_DEVICES="0,2"
+ $ ./pruned_transducer_stateless7_streaming/train.py --world-size 2
+
+ **Use case 2**: You have 4 GPUs and you want to use all of them
+ for training. You can do the following:
+
+ .. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./pruned_transducer_stateless7_streaming/train.py --world-size 4
+
+ **Use case 3**: You have 4 GPUs but you only want to use GPU 3
+ for training. You can do the following:
+
+ .. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ export CUDA_VISIBLE_DEVICES="3"
+ $ ./pruned_transducer_stateless7_streaming/train.py --world-size 1
+
+ .. caution::
+
+ Only multi-GPU single-machine DDP training is implemented at present.
+ Multi-GPU multi-machine DDP training will be added later.
+
+ - ``--max-duration``
+
+ It specifies the number of seconds over all utterances in a
+ batch, before **padding**.
+ If you encounter CUDA OOM, please reduce it.
+
+ .. HINT::
+
+ Due to padding, the number of seconds of all utterances in a
+ batch will usually be larger than ``--max-duration``.
+
+ A larger value for ``--max-duration`` may cause OOM during training,
+ while a smaller value may increase the training time. You have to
+ tune it.
+
+ - ``--use-fp16``
+
+ If it is True, the model will train with half precision, from our experiment
+ results, by using half precision you can train with two times larger ``--max-duration``
+ so as to get almost 2X speed up.
+
+ We recommend using ``--use-fp16 True``.
+
+ - ``--short-chunk-size``
+
+ When training a streaming attention model with chunk masking, the chunk size
+ would be either max sequence length of current batch or uniformly sampled from
+ (1, short_chunk_size). The default value is 50, you don't have to change it most of the time.
+
+ - ``--num-left-chunks``
+
+ It indicates how many left context (in chunks) that can be seen when calculating attention.
+ The default value is 4, you don't have to change it most of the time.
+
+
+ - ``--decode-chunk-len``
+
+ The chunk size for decoding (in frames before subsampling). It is used for validation.
+ The default value is 32 (i.e., 320ms).
+
+
+Pre-configured options
+~~~~~~~~~~~~~~~~~~~~~~
+
+There are some training options, e.g., number of encoder layers,
+encoder dimension, decoder dimension, number of warmup steps etc,
+that are not passed from the commandline.
+They are pre-configured by the function ``get_params()`` in
+`pruned_transducer_stateless7_streaming/train.py `_
+
+You don't need to change these pre-configured parameters. If you really need to change
+them, please modify ``./pruned_transducer_stateless7_streaming/train.py`` directly.
+
+
+Training logs
+~~~~~~~~~~~~~
+
+Training logs and checkpoints are saved in ``--exp-dir`` (e.g. ``pruned_transducer_stateless7_streaming/exp``.
+You will find the following files in that directory:
+
+ - ``epoch-1.pt``, ``epoch-2.pt``, ...
+
+ These are checkpoint files saved at the end of each epoch, containing model
+ ``state_dict`` and optimizer ``state_dict``.
+ To resume training from some checkpoint, say ``epoch-10.pt``, you can use:
+
+ .. code-block:: bash
+
+ $ ./pruned_transducer_stateless7_streaming/train.py --start-epoch 11
+
+ - ``checkpoint-436000.pt``, ``checkpoint-438000.pt``, ...
+
+ These are checkpoint files saved every ``--save-every-n`` batches,
+ containing model ``state_dict`` and optimizer ``state_dict``.
+ To resume training from some checkpoint, say ``checkpoint-436000``, you can use:
+
+ .. code-block:: bash
+
+ $ ./pruned_transducer_stateless7_streaming/train.py --start-batch 436000
+
+ - ``tensorboard/``
+
+ This folder contains tensorBoard logs. Training loss, validation loss, learning
+ rate, etc, are recorded in these logs. You can visualize them by:
+
+ .. code-block:: bash
+
+ $ cd pruned_transducer_stateless7_streaming/exp/tensorboard
+ $ tensorboard dev upload --logdir . --description "pruned transducer training for LibriSpeech with icefall"
+
+ .. hint::
+
+ If you don't have access to google, you can use the following command
+ to view the tensorboard log locally:
+
+ .. code-block:: bash
+
+ cd pruned_transducer_stateless7_streaming/exp/tensorboard
+ tensorboard --logdir . --port 6008
+
+ It will print the following message:
+
+ .. code-block::
+
+ Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
+ TensorBoard 2.8.0 at http://localhost:6008/ (Press CTRL+C to quit)
+
+ Now start your browser and go to ``_ to view the tensorboard
+ logs.
+
+
+ - ``log/log-train-xxxx``
+
+ It is the detailed training log in text format, same as the one
+ you saw printed to the console during training.
+
+Usage example
+~~~~~~~~~~~~~
+
+You can use the following command to start the training using 4 GPUs:
+
+.. code-block:: bash
+
+ export CUDA_VISIBLE_DEVICES="0,1,2,3"
+ ./pruned_transducer_stateless7_streaming/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --use-fp16 1 \
+ --exp-dir pruned_transducer_stateless7_streaming/exp \
+ --full-libri 1 \
+ --max-duration 550
+
+Decoding
+--------
+
+The decoding part uses checkpoints saved by the training part, so you have
+to run the training part first.
+
+.. hint::
+
+ There are two kinds of checkpoints:
+
+ - (1) ``epoch-1.pt``, ``epoch-2.pt``, ..., which are saved at the end
+ of each epoch. You can pass ``--epoch`` to
+ ``pruned_transducer_stateless7_streaming/decode.py`` to use them.
+
+ - (2) ``checkpoints-436000.pt``, ``epoch-438000.pt``, ..., which are saved
+ every ``--save-every-n`` batches. You can pass ``--iter`` to
+ ``pruned_transducer_stateless7_streaming/decode.py`` to use them.
+
+ We suggest that you try both types of checkpoints and choose the one
+ that produces the lowest WERs.
+
+.. tip::
+
+ To decode a streaming model, you can use either ``simulate streaming decoding`` in ``decode.py`` or
+ ``real chunk-wise streaming decoding`` in ``streaming_decode.py``. The difference between ``decode.py`` and
+ ``streaming_decode.py`` is that, ``decode.py`` processes the whole acoustic frames at one time with masking (i.e. same as training),
+ but ``streaming_decode.py`` processes the acoustic frames chunk by chunk.
+
+.. NOTE::
+
+ ``simulate streaming decoding`` in ``decode.py`` and ``real chunk-size streaming decoding`` in ``streaming_decode.py`` should
+ produce almost the same results given the same ``--decode-chunk-len``.
+
+
+Simulate streaming decoding
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./pruned_transducer_stateless7_streaming/decode.py --help
+
+shows the options for decoding.
+The following options are important for streaming models:
+
+ ``--decode-chunk-len``
+
+ It is same as in ``train.py``, which specifies the chunk size for decoding (in frames before subsampling).
+ The default value is 32 (i.e., 320ms).
+
+
+The following shows two examples (for the two types of checkpoints):
+
+.. code-block:: bash
+
+ for m in greedy_search fast_beam_search modified_beam_search; do
+ for epoch in 30; do
+ for avg in 12 11 10 9 8; do
+ ./pruned_transducer_stateless7_streaming/decode.py \
+ --epoch $epoch \
+ --avg $avg \
+ --decode-chunk-len 32 \
+ --exp-dir pruned_transducer_stateless7_streaming/exp \
+ --max-duration 600 \
+ --decoding-method $m
+ done
+ done
+ done
+
+
+.. code-block:: bash
+
+ for m in greedy_search fast_beam_search modified_beam_search; do
+ for iter in 474000; do
+ for avg in 8 10 12 14 16 18; do
+ ./pruned_transducer_stateless7_streaming/decode.py \
+ --iter $iter \
+ --avg $avg \
+ --decode-chunk-len 32 \
+ --exp-dir pruned_transducer_stateless7_streaming/exp \
+ --max-duration 600 \
+ --decoding-method $m
+ done
+ done
+ done
+
+
+Real streaming decoding
+~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: bash
+
+ $ cd egs/librispeech/ASR
+ $ ./pruned_transducer_stateless7_streaming/streaming_decode.py --help
+
+shows the options for decoding.
+The following options are important for streaming models:
+
+ ``--decode-chunk-len``
+
+ It is same as in ``train.py``, which specifies the chunk size for decoding (in frames before subsampling).
+ The default value is 32 (i.e., 320ms).
+ For ``real streaming decoding``, we will process ``decode-chunk-len`` acoustic frames at each time.
+
+ ``--num-decode-streams``
+
+ The number of decoding streams that can be run in parallel (very similar to the ``bath size``).
+ For ``real streaming decoding``, the batches will be packed dynamically, for example, if the
+ ``num-decode-streams`` equals to 10, then, sequence 1 to 10 will be decoded at first, after a while,
+ suppose sequence 1 and 2 are done, so, sequence 3 to 12 will be processed parallelly in a batch.
+
+
+The following shows two examples (for the two types of checkpoints):
+
+.. code-block:: bash
+
+ for m in greedy_search fast_beam_search modified_beam_search; do
+ for epoch in 30; do
+ for avg in 12 11 10 9 8; do
+ ./pruned_transducer_stateless7_streaming/decode.py \
+ --epoch $epoch \
+ --avg $avg \
+ --decode-chunk-len 32 \
+ --num-decode-streams 100 \
+ --exp-dir pruned_transducer_stateless7_streaming/exp \
+ --decoding-method $m
+ done
+ done
+ done
+
+
+.. code-block:: bash
+
+ for m in greedy_search fast_beam_search modified_beam_search; do
+ for iter in 474000; do
+ for avg in 8 10 12 14 16 18; do
+ ./pruned_transducer_stateless7_streaming/decode.py \
+ --iter $iter \
+ --avg $avg \
+ --decode-chunk-len 16 \
+ --num-decode-streams 100 \
+ --exp-dir pruned_transducer_stateless7_streaming/exp \
+ --decoding-method $m
+ done
+ done
+ done
+
+
+.. tip::
+
+ Supporting decoding methods are as follows:
+
+ - ``greedy_search`` : It takes the symbol with largest posterior probability
+ of each frame as the decoding result.
+
+ - ``beam_search`` : It implements Algorithm 1 in https://arxiv.org/pdf/1211.3711.pdf and
+ `espnet/nets/beam_search_transducer.py `_
+ is used as a reference. Basicly, it keeps topk states for each frame, and expands the kept states with their own contexts to
+ next frame.
+
+ - ``modified_beam_search`` : It implements the same algorithm as ``beam_search`` above, but it
+ runs in batch mode with ``--max-sym-per-frame=1`` being hardcoded.
+
+ - ``fast_beam_search`` : It implements graph composition between the output ``log_probs`` and
+ given ``FSAs``. It is hard to describe the details in several lines of texts, you can read
+ our paper in https://arxiv.org/pdf/2211.00484.pdf or our `rnnt decode code in k2 `_. ``fast_beam_search`` can decode with ``FSAs`` on GPU efficiently.
+
+ - ``fast_beam_search_LG`` : The same as ``fast_beam_search`` above, ``fast_beam_search`` uses
+ an trivial graph that has only one state, while ``fast_beam_search_LG`` uses an LG graph
+ (with N-gram LM).
+
+ - ``fast_beam_search_nbest`` : It produces the decoding results as follows:
+
+ - (1) Use ``fast_beam_search`` to get a lattice
+ - (2) Select ``num_paths`` paths from the lattice using ``k2.random_paths()``
+ - (3) Unique the selected paths
+ - (4) Intersect the selected paths with the lattice and compute the
+ shortest path from the intersection result
+ - (5) The path with the largest score is used as the decoding output.
+
+ - ``fast_beam_search_nbest_LG`` : It implements same logic as ``fast_beam_search_nbest``, the
+ only difference is that it uses ``fast_beam_search_LG`` to generate the lattice.
+
+.. NOTE::
+
+ The supporting decoding methods in ``streaming_decode.py`` might be less than that in ``decode.py``, if needed,
+ you can implement them by yourself or file a issue in `icefall `_ .
+
+
+Export Model
+------------
+
+Currently it supports exporting checkpoints from ``pruned_transducer_stateless7_streaming/exp`` in the following ways.
+
+Export ``model.state_dict()``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Checkpoints saved by ``pruned_transducer_stateless7_streaming/train.py`` also include
+``optimizer.state_dict()``. It is useful for resuming training. But after training,
+we are interested only in ``model.state_dict()``. You can use the following
+command to extract ``model.state_dict()``.
+
+.. code-block:: bash
+
+ # Assume that --epoch 30 --avg 9 produces the smallest WER
+ # (You can get such information after running ./pruned_transducer_stateless7_streaming/decode.py)
+
+ epoch=30
+ avg=9
+
+ ./pruned_transducer_stateless7_streaming/export.py \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch $epoch \
+ --avg $avg \
+ --use-averaged-model=True \
+ --decode-chunk-len 32
+
+It will generate a file ``./pruned_transducer_stateless7_streaming/exp/pretrained.pt``.
+
+.. hint::
+
+ To use the generated ``pretrained.pt`` for ``pruned_transducer_stateless7_streaming/decode.py``,
+ you can run:
+
+ .. code-block:: bash
+
+ cd pruned_transducer_stateless7_streaming/exp
+ ln -s pretrained.pt epoch-999.pt
+
+ And then pass ``--epoch 999 --avg 1 --use-averaged-model 0`` to
+ ``./pruned_transducer_stateless7_streaming/decode.py``.
+
+To use the exported model with ``./pruned_transducer_stateless7_streaming/pretrained.py``, you
+can run:
+
+.. code-block:: bash
+
+ ./pruned_transducer_stateless7_streaming/pretrained.py \
+ --checkpoint ./pruned_transducer_stateless7_streaming/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method greedy_search \
+ --decode-chunk-len 32 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+
+Export model using ``torch.jit.script()``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: bash
+
+ ./pruned_transducer_stateless7_streaming/export.py \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 9 \
+ --decode-chunk-len 32 \
+ --jit 1
+
+.. caution::
+
+ ``--decode-chunk-len`` is required to export a ScriptModule.
+
+It will generate a file ``cpu_jit.pt`` in the given ``exp_dir``. You can later
+load it by ``torch.jit.load("cpu_jit.pt")``.
+
+Note ``cpu`` in the name ``cpu_jit.pt`` means the parameters when loaded into Python
+are on CPU. You can use ``to("cuda")`` to move them to a CUDA device.
+
+Export model using ``torch.jit.trace()``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: bash
+
+ epoch=30
+ avg=9
+
+ ./pruned_transducer_stateless7_streaming/jit_trace_export.py \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --use-averaged-model=True \
+ --decode-chunk-len 32 \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --epoch $epoch \
+ --avg $avg
+
+.. caution::
+
+ ``--decode-chunk-len`` is required to export a ScriptModule.
+
+It will generate 3 files:
+
+ - ``./pruned_transducer_stateless7_streaming/exp/encoder_jit_trace.pt``
+ - ``./pruned_transducer_stateless7_streaming/exp/decoder_jit_trace.pt``
+ - ``./pruned_transducer_stateless7_streaming/exp/joiner_jit_trace.pt``
+
+To use the generated files with ``./pruned_transducer_stateless7_streaming/jit_trace_pretrained.py``:
+
+.. code-block:: bash
+
+ ./pruned_transducer_stateless7_streaming/jit_trace_pretrained.py \
+ --encoder-model-filename ./pruned_transducer_stateless7_streaming/exp/encoder_jit_trace.pt \
+ --decoder-model-filename ./pruned_transducer_stateless7_streaming/exp/decoder_jit_trace.pt \
+ --joiner-model-filename ./pruned_transducer_stateless7_streaming/exp/joiner_jit_trace.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --decode-chunk-len 32 \
+ /path/to/foo.wav
+
+
+Download pretrained models
+--------------------------
+
+If you don't want to train from scratch, you can download the pretrained models
+by visiting the following links:
+
+ - `pruned_transducer_stateless7_streaming `_
+
+ See ``_
+ for the details of the above pretrained models
+
+Deploy with Sherpa
+------------------
+
+Please see ``_
+for how to deploy the models in ``sherpa``.
diff --git a/docs/source/recipes/index.rst b/docs/source/recipes/index.rst
index 9d1d83d29..63793275c 100644
--- a/docs/source/recipes/index.rst
+++ b/docs/source/recipes/index.rst
@@ -13,7 +13,5 @@ We may add recipes for other tasks as well in the future.
:maxdepth: 2
:caption: Table of Contents
- aishell/index
- librispeech/index
- timit/index
- yesno/index
+ Non-streaming-ASR/index
+ Streaming-ASR/index
diff --git a/egs/alimeeting/ASR_v2/README.md b/egs/alimeeting/ASR_v2/README.md
new file mode 100644
index 000000000..f70327501
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/README.md
@@ -0,0 +1,38 @@
+
+# Introduction
+
+This recipe trains multi-domain ASR models for AliMeeting. By multi-domain, we mean that
+we train a single model on close-talk and far-field conditions. This recipe optionally
+uses [GSS]-based enhancement for far-field array microphone.
+We pool data in the following 4 ways and train a single model on the pooled data:
+
+(i) individual headset microphone (IHM)
+(ii) IHM with simulated reverb
+(iii) Single distant microphone (SDM)
+(iv) GSS-enhanced array microphones
+
+This is different from `alimeeting/ASR` since that recipe trains a model only on the
+far-field audio. Additionally, we use text normalization here similar to the original
+M2MeT challenge, so the results should be more comparable to those from Table 4 of
+the [paper](https://arxiv.org/abs/2110.07393).
+
+The following additional packages need to be installed to run this recipe:
+* `pip install jieba`
+* `pip install paddlepaddle`
+* `pip install git+https://github.com/desh2608/gss.git`
+
+[./RESULTS.md](./RESULTS.md) contains the latest results.
+
+## Performance Record
+
+### pruned_transducer_stateless7
+
+The following are decoded using `modified_beam_search`:
+
+| Evaluation set | eval WER | test WER |
+|--------------------------|------------|---------|
+| IHM | 9.58 | 11.53 |
+| SDM | 23.37 | 25.85 |
+| MDM (GSS-enhanced) | 11.82 | 14.22 |
+
+See [RESULTS](/egs/alimeeting/ASR_v2/RESULTS.md) for details.
diff --git a/egs/alimeeting/ASR_v2/RESULTS.md b/egs/alimeeting/ASR_v2/RESULTS.md
new file mode 100644
index 000000000..15b24250d
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/RESULTS.md
@@ -0,0 +1,90 @@
+## Results (CER)
+
+#### 2022-12-09
+
+#### Zipformer (pruned_transducer_stateless7)
+
+Zipformer encoder + non-current decoder. The decoder
+contains only an embedding layer, a Conv1d (with kernel size 2) and a linear
+layer (to transform tensor dim).
+
+All the results below are using a single model that is trained by combining the following
+data: IHM, IHM+reverb, SDM, and GSS-enhanced MDM. Speed perturbation and MUSAN noise
+augmentation are applied on top of the pooled data.
+
+**WERs for IHM:**
+
+| | eval | test | comment |
+|---------------------------|------------|------------|------------------------------------------|
+| greedy search | 10.13 | 12.21 | --epoch 15 --avg 8 --max-duration 500 |
+| modified beam search | 9.58 | 11.53 | --epoch 15 --avg 8 --max-duration 500 --beam-size 4 |
+| fast beam search | 9.92 | 12.07 | --epoch 15 --avg 8 --max-duration 500 --beam-size 4 --max-contexts 4 --max-states 8 |
+
+**WERs for SDM:**
+
+| | eval | test | comment |
+|---------------------------|------------|------------|------------------------------------------|
+| greedy search | 23.70 | 26.41 | --epoch 15 --avg 8 --max-duration 500 |
+| modified beam search | 23.37 | 25.85 | --epoch 15 --avg 8 --max-duration 500 --beam-size 4 |
+| fast beam search | 23.60 | 26.38 | --epoch 15 --avg 8 --max-duration 500 --beam-size 4 --max-contexts 4 --max-states 8 |
+
+**WERs for GSS-enhanced MDM:**
+
+| | eval | test | comment |
+|---------------------------|------------|------------|------------------------------------------|
+| greedy search | 12.24 | 14.99 | --epoch 15 --avg 8 --max-duration 500 |
+| modified beam search | 11.82 | 14.22 | --epoch 15 --avg 8 --max-duration 500 --beam-size 4 |
+| fast beam search | 12.30 | 14.98 | --epoch 15 --avg 8 --max-duration 500 --beam-size 4 --max-contexts 4 --max-states 8 |
+
+The training command for reproducing is given below:
+
+```
+export CUDA_VISIBLE_DEVICES="0,1,2,3"
+
+./pruned_transducer_stateless7/train.py \
+ --world-size 4 \
+ --num-epochs 15 \
+ --exp-dir pruned_transducer_stateless7/exp \
+ --max-duration 300 \
+ --max-cuts 100 \
+ --prune-range 5 \
+ --lr-factor 5 \
+ --lm-scale 0.25 \
+ --use-fp16 True
+```
+
+The decoding command is:
+```
+# greedy search
+./pruned_transducer_stateless7/decode.py \
+ --epoch 15 \
+ --avg 8 \
+ --exp-dir ./pruned_transducer_stateless7/exp \
+ --max-duration 500 \
+ --decoding-method greedy_search
+
+# modified beam search
+./pruned_transducer_stateless7/decode.py \
+ --epoch 15 \
+ --avg 8 \
+ --exp-dir ./pruned_transducer_stateless7/exp \
+ --max-duration 500 \
+ --decoding-method modified_beam_search \
+ --beam-size 4
+
+# fast beam search
+./pruned_transducer_stateless7/decode.py \
+ --epoch 15 \
+ --avg 8 \
+ --exp-dir ./pruned_transducer_stateless5/exp \
+ --max-duration 500 \
+ --decoding-method fast_beam_search \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8
+```
+
+Pretrained model is available at
+
+The tensorboard training log can be found at
+
diff --git a/egs/alimeeting/ASR_v2/local/__init__.py b/egs/alimeeting/ASR_v2/local/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/egs/alimeeting/ASR_v2/local/compute_fbank_alimeeting.py b/egs/alimeeting/ASR_v2/local/compute_fbank_alimeeting.py
new file mode 100755
index 000000000..c6aa2ab36
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/local/compute_fbank_alimeeting.py
@@ -0,0 +1,193 @@
+#!/usr/bin/env python3
+# Copyright 2022 Johns Hopkins University (authors: Desh Raj)
+#
+# 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.
+
+
+"""
+This file computes fbank features of the AliMeeting dataset.
+For the training data, we prepare IHM, reverberated IHM, SDM, and GSS-enhanced
+audios. For the test data, we separately prepare IHM, SDM, and GSS-enhanced
+parts (which are the 3 evaluation settings).
+It looks for manifests in the directory data/manifests.
+
+The generated fbank features are saved in data/fbank.
+"""
+import logging
+from pathlib import Path
+
+import torch
+import torch.multiprocessing
+from lhotse import CutSet, LilcomChunkyWriter
+from lhotse.features.kaldifeat import (
+ KaldifeatFbank,
+ KaldifeatFbankConfig,
+ KaldifeatFrameOptions,
+ KaldifeatMelOptions,
+)
+from lhotse.recipes.utils import read_manifests_if_cached
+
+# Torch's multithreaded behavior needs to be disabled or
+# it wastes a lot of CPU and slow things down.
+# Do this outside of main() in case it needs to take effect
+# even when we are not invoking the main (e.g. when spawning subprocesses).
+torch.set_num_threads(1)
+torch.set_num_interop_threads(1)
+torch.multiprocessing.set_sharing_strategy("file_system")
+
+
+def compute_fbank_ami():
+ src_dir = Path("data/manifests")
+ output_dir = Path("data/fbank")
+
+ sampling_rate = 16000
+ num_mel_bins = 80
+
+ extractor = KaldifeatFbank(
+ KaldifeatFbankConfig(
+ frame_opts=KaldifeatFrameOptions(sampling_rate=sampling_rate),
+ mel_opts=KaldifeatMelOptions(num_bins=num_mel_bins),
+ device="cuda",
+ )
+ )
+
+ logging.info("Reading manifests")
+ manifests_ihm = read_manifests_if_cached(
+ dataset_parts=["train", "eval", "test"],
+ output_dir=src_dir,
+ prefix="alimeeting-ihm",
+ suffix="jsonl.gz",
+ )
+ manifests_sdm = read_manifests_if_cached(
+ dataset_parts=["train", "eval", "test"],
+ output_dir=src_dir,
+ prefix="alimeeting-sdm",
+ suffix="jsonl.gz",
+ )
+ # For GSS we already have cuts so we read them directly.
+ manifests_gss = read_manifests_if_cached(
+ dataset_parts=["train", "eval", "test"],
+ output_dir=src_dir,
+ prefix="alimeeting-gss",
+ suffix="jsonl.gz",
+ )
+
+ def _extract_feats(cuts: CutSet, storage_path: Path, manifest_path: Path) -> None:
+ cuts = cuts + cuts.perturb_speed(0.9) + cuts.perturb_speed(1.1)
+ _ = cuts.compute_and_store_features_batch(
+ extractor=extractor,
+ storage_path=storage_path,
+ manifest_path=manifest_path,
+ batch_duration=5000,
+ num_workers=8,
+ storage_type=LilcomChunkyWriter,
+ )
+
+ logging.info(
+ "Preparing training cuts: IHM + reverberated IHM + SDM + GSS (optional)"
+ )
+
+ logging.info("Processing train split IHM")
+ cuts_ihm = (
+ CutSet.from_manifests(**manifests_ihm["train"])
+ .trim_to_supervisions(keep_overlapping=False, keep_all_channels=False)
+ .modify_ids(lambda x: x + "-ihm")
+ )
+ _extract_feats(
+ cuts_ihm,
+ output_dir / "feats_train_ihm",
+ src_dir / "cuts_train_ihm.jsonl.gz",
+ )
+
+ logging.info("Processing train split IHM + reverberated IHM")
+ cuts_ihm_rvb = cuts_ihm.reverb_rir()
+ _extract_feats(
+ cuts_ihm_rvb,
+ output_dir / "feats_train_ihm_rvb",
+ src_dir / "cuts_train_ihm_rvb.jsonl.gz",
+ )
+
+ logging.info("Processing train split SDM")
+ cuts_sdm = (
+ CutSet.from_manifests(**manifests_sdm["train"])
+ .trim_to_supervisions(keep_overlapping=False)
+ .modify_ids(lambda x: x + "-sdm")
+ )
+ _extract_feats(
+ cuts_sdm,
+ output_dir / "feats_train_sdm",
+ src_dir / "cuts_train_sdm.jsonl.gz",
+ )
+
+ logging.info("Processing train split GSS")
+ cuts_gss = (
+ CutSet.from_manifests(**manifests_gss["train"])
+ .trim_to_supervisions(keep_overlapping=False)
+ .modify_ids(lambda x: x + "-gss")
+ )
+ _extract_feats(
+ cuts_gss,
+ output_dir / "feats_train_gss",
+ src_dir / "cuts_train_gss.jsonl.gz",
+ )
+
+ logging.info("Preparing test cuts: IHM, SDM, GSS (optional)")
+ for split in ["eval", "test"]:
+ logging.info(f"Processing {split} IHM")
+ cuts_ihm = (
+ CutSet.from_manifests(**manifests_ihm[split])
+ .trim_to_supervisions(keep_overlapping=False, keep_all_channels=False)
+ .compute_and_store_features_batch(
+ extractor=extractor,
+ storage_path=output_dir / f"feats_{split}_ihm",
+ manifest_path=src_dir / f"cuts_{split}_ihm.jsonl.gz",
+ batch_duration=500,
+ num_workers=4,
+ storage_type=LilcomChunkyWriter,
+ )
+ )
+ logging.info(f"Processing {split} SDM")
+ cuts_sdm = (
+ CutSet.from_manifests(**manifests_sdm[split])
+ .trim_to_supervisions(keep_overlapping=False)
+ .compute_and_store_features_batch(
+ extractor=extractor,
+ storage_path=output_dir / f"feats_{split}_sdm",
+ manifest_path=src_dir / f"cuts_{split}_sdm.jsonl.gz",
+ batch_duration=500,
+ num_workers=4,
+ storage_type=LilcomChunkyWriter,
+ )
+ )
+ logging.info(f"Processing {split} GSS")
+ cuts_gss = (
+ CutSet.from_manifests(**manifests_gss[split])
+ .trim_to_supervisions(keep_overlapping=False)
+ .compute_and_store_features_batch(
+ extractor=extractor,
+ storage_path=output_dir / f"feats_{split}_gss",
+ manifest_path=src_dir / f"cuts_{split}_gss.jsonl.gz",
+ batch_duration=500,
+ num_workers=4,
+ storage_type=LilcomChunkyWriter,
+ )
+ )
+
+
+if __name__ == "__main__":
+ formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
+ logging.basicConfig(format=formatter, level=logging.INFO)
+
+ compute_fbank_ami()
diff --git a/egs/alimeeting/ASR_v2/local/compute_fbank_musan.py b/egs/alimeeting/ASR_v2/local/compute_fbank_musan.py
new file mode 120000
index 000000000..5833f2484
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/local/compute_fbank_musan.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/local/compute_fbank_musan.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/local/prepare_alimeeting_enhanced.py b/egs/alimeeting/ASR_v2/local/prepare_alimeeting_enhanced.py
new file mode 100644
index 000000000..f1512efa5
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/local/prepare_alimeeting_enhanced.py
@@ -0,0 +1,158 @@
+#!/usr/local/bin/python
+# -*- coding: utf-8 -*-
+# Data preparation for AliMeeting GSS-enhanced dataset.
+
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+from lhotse import Recording, RecordingSet, SupervisionSet
+from lhotse.qa import fix_manifests
+from lhotse.recipes.utils import read_manifests_if_cached
+from lhotse.utils import fastcopy
+from tqdm import tqdm
+
+logging.basicConfig(
+ format="%(asctime)s %(levelname)-8s %(message)s",
+ level=logging.INFO,
+ datefmt="%Y-%m-%d %H:%M:%S",
+)
+
+
+def get_args():
+ import argparse
+
+ parser = argparse.ArgumentParser(description="AMI enhanced dataset preparation.")
+ parser.add_argument(
+ "manifests_dir",
+ type=Path,
+ help="Path to directory containing AliMeeting manifests.",
+ )
+ parser.add_argument(
+ "enhanced_dir",
+ type=Path,
+ help="Path to enhanced data directory.",
+ )
+ parser.add_argument(
+ "--num-jobs",
+ "-j",
+ type=int,
+ default=1,
+ help="Number of parallel jobs to run.",
+ )
+ parser.add_argument(
+ "--min-segment-duration",
+ "-d",
+ type=float,
+ default=0.0,
+ help="Minimum duration of a segment in seconds.",
+ )
+ return parser.parse_args()
+
+
+def find_recording_and_create_new_supervision(enhanced_dir, supervision):
+ """
+ Given a supervision (corresponding to original AMI recording), this function finds the
+ enhanced recording correspoding to the supervision, and returns this recording and
+ a new supervision whose start and end times are adjusted to match the enhanced recording.
+ """
+ file_name = Path(
+ f"{supervision.recording_id}-{supervision.speaker}-{int(100*supervision.start):06d}_{int(100*supervision.end):06d}.flac"
+ )
+ save_path = enhanced_dir / f"{supervision.recording_id}" / file_name
+ if save_path.exists():
+ recording = Recording.from_file(save_path)
+ if recording.duration == 0:
+ logging.warning(f"Skipping {save_path} which has duration 0 seconds.")
+ return None
+
+ # Old supervision is wrt to the original recording, we create new supervision
+ # wrt to the enhanced segment
+ new_supervision = fastcopy(
+ supervision,
+ recording_id=recording.id,
+ start=0,
+ duration=recording.duration,
+ )
+ return recording, new_supervision
+ else:
+ logging.warning(f"{save_path} does not exist.")
+ return None
+
+
+def main(args):
+ # Get arguments
+ manifests_dir = args.manifests_dir
+ enhanced_dir = args.enhanced_dir
+
+ # Load manifests from cache if they exist (saves time)
+ manifests = read_manifests_if_cached(
+ dataset_parts=["train", "eval", "test"],
+ output_dir=manifests_dir,
+ prefix="alimeeting-sdm",
+ suffix="jsonl.gz",
+ )
+ if not manifests:
+ raise ValueError(
+ "AliMeeting SDM manifests not found in {}".format(manifests_dir)
+ )
+
+ with ThreadPoolExecutor(args.num_jobs) as ex:
+ for part in ["train", "eval", "test"]:
+ logging.info(f"Processing {part}...")
+ supervisions_orig = manifests[part]["supervisions"].filter(
+ lambda s: s.duration >= args.min_segment_duration
+ )
+ futures = []
+
+ for supervision in tqdm(
+ supervisions_orig,
+ desc="Distributing tasks",
+ ):
+ futures.append(
+ ex.submit(
+ find_recording_and_create_new_supervision,
+ enhanced_dir,
+ supervision,
+ )
+ )
+
+ recordings = []
+ supervisions = []
+ for future in tqdm(
+ futures,
+ total=len(futures),
+ desc="Processing tasks",
+ ):
+ result = future.result()
+ if result is not None:
+ recording, new_supervision = result
+ recordings.append(recording)
+ supervisions.append(new_supervision)
+
+ # Remove duplicates from the recordings
+ recordings_nodup = {}
+ for recording in recordings:
+ if recording.id not in recordings_nodup:
+ recordings_nodup[recording.id] = recording
+ else:
+ logging.warning("Recording {} is duplicated.".format(recording.id))
+ recordings = RecordingSet.from_recordings(recordings_nodup.values())
+ supervisions = SupervisionSet.from_segments(supervisions)
+
+ recordings, supervisions = fix_manifests(
+ recordings=recordings, supervisions=supervisions
+ )
+
+ logging.info(f"Writing {part} enhanced manifests")
+ recordings.to_file(
+ manifests_dir / f"alimeeting-gss_recordings_{part}.jsonl.gz"
+ )
+ supervisions.to_file(
+ manifests_dir / f"alimeeting-gss_supervisions_{part}.jsonl.gz"
+ )
+
+
+if __name__ == "__main__":
+ args = get_args()
+ main(args)
diff --git a/egs/alimeeting/ASR_v2/local/prepare_alimeeting_gss.sh b/egs/alimeeting/ASR_v2/local/prepare_alimeeting_gss.sh
new file mode 100755
index 000000000..76db19832
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/local/prepare_alimeeting_gss.sh
@@ -0,0 +1,98 @@
+#!/bin/bash
+# This script is used to run GSS-based enhancement on AMI data.
+set -euo pipefail
+nj=4
+stage=0
+
+. shared/parse_options.sh || exit 1
+
+if [ $# != 2 ]; then
+ echo "Wrong #arguments ($#, expected 2)"
+ echo "Usage: local/prepare_alimeeting_gss.sh [options] "
+ echo "e.g. local/prepare_alimeeting_gss.sh data/manifests exp/ami_gss"
+ echo "main options (for others, see top of script file)"
+ echo " --nj # number of parallel jobs"
+ echo " --stage # stage to start running from"
+ exit 1;
+fi
+
+DATA_DIR=$1
+EXP_DIR=$2
+
+mkdir -p $EXP_DIR
+
+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 ]; then
+ log "Stage 1: Prepare cut sets"
+ for part in train eval test; do
+ lhotse cut simple \
+ -r $DATA_DIR/alimeeting-mdm_recordings_${part}.jsonl.gz \
+ -s $DATA_DIR/alimeeting-mdm_supervisions_${part}.jsonl.gz \
+ $EXP_DIR/cuts_${part}.jsonl.gz
+ done
+fi
+
+if [ $stage -le 2 ]; then
+ log "Stage 2: Trim cuts to supervisions (1 cut per supervision segment)"
+ for part in train eval test; do
+ lhotse cut trim-to-supervisions --discard-overlapping \
+ $EXP_DIR/cuts_${part}.jsonl.gz $EXP_DIR/cuts_per_segment_${part}.jsonl.gz
+ done
+fi
+
+if [ $stage -le 3 ]; then
+ log "Stage 3: Split manifests for multi-GPU processing (optional)"
+ for part in train eval test; do
+ gss utils split $nj $EXP_DIR/cuts_per_segment_${part}.jsonl.gz \
+ $EXP_DIR/cuts_per_segment_${part}_split$nj
+ done
+fi
+
+if [ $stage -le 4 ]; then
+ log "Stage 4: Enhance train segments using GSS (requires GPU)"
+ # for train, we use smaller context and larger batches to speed-up processing
+ for JOB in $(seq $nj); do
+ gss enhance cuts $EXP_DIR/cuts_train.jsonl.gz \
+ $EXP_DIR/cuts_per_segment_train_split$nj/cuts_per_segment_train.JOB.jsonl.gz $EXP_DIR/enhanced \
+ --bss-iterations 10 \
+ --context-duration 5.0 \
+ --use-garbage-class \
+ --channels 0,1,2,3,4,5,6,7 \
+ --min-segment-length 0.05 \
+ --max-segment-length 25.0 \
+ --max-batch-duration 60.0 \
+ --num-buckets 4 \
+ --num-workers 4
+ done
+fi
+
+if [ $stage -le 5 ]; then
+ log "Stage 5: Enhance eval/test segments using GSS (using GPU)"
+ # for eval/test, we use larger context and smaller batches to get better quality
+ for part in eval test; do
+ for JOB in $(seq $nj); do
+ gss enhance cuts $EXP_DIR/cuts_${part}.jsonl.gz \
+ $EXP_DIR/cuts_per_segment_${part}_split$nj/cuts_per_segment_${part}.JOB.jsonl.gz \
+ $EXP_DIR/enhanced \
+ --bss-iterations 10 \
+ --context-duration 15.0 \
+ --use-garbage-class \
+ --channels 0,1,2,3,4,5,6,7 \
+ --min-segment-length 0.05 \
+ --max-segment-length 16.0 \
+ --max-batch-duration 45.0 \
+ --num-buckets 4 \
+ --num-workers 4
+ done
+ done
+fi
+
+if [ $stage -le 6 ]; then
+ log "Stage 6: Prepare manifests for GSS-enhanced data"
+ python local/prepare_alimeeting_enhanced.py $DATA_DIR $EXP_DIR/enhanced -j $nj --min-segment-duration 0.05
+fi
diff --git a/egs/alimeeting/ASR_v2/local/prepare_char.py b/egs/alimeeting/ASR_v2/local/prepare_char.py
new file mode 120000
index 000000000..ee5dd34f1
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/local/prepare_char.py
@@ -0,0 +1 @@
+../../ASR/local/prepare_char.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/local/prepare_words.py b/egs/alimeeting/ASR_v2/local/prepare_words.py
new file mode 120000
index 000000000..970bfd60c
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/local/prepare_words.py
@@ -0,0 +1 @@
+../../ASR/local/prepare_words.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/local/text2segments.py b/egs/alimeeting/ASR_v2/local/text2segments.py
new file mode 120000
index 000000000..bf4547794
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/local/text2segments.py
@@ -0,0 +1 @@
+../../ASR/local/text2segments.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/local/text2token.py b/egs/alimeeting/ASR_v2/local/text2token.py
new file mode 120000
index 000000000..f6b8531b6
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/local/text2token.py
@@ -0,0 +1 @@
+../../ASR/local/text2token.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/prepare.sh b/egs/alimeeting/ASR_v2/prepare.sh
new file mode 100755
index 000000000..76a108771
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/prepare.sh
@@ -0,0 +1,125 @@
+#!/usr/bin/env bash
+
+set -eou pipefail
+
+stage=-1
+stop_stage=100
+use_gss=true # Use GSS-based enhancement with MDM setting
+
+# We assume dl_dir (download dir) contains the following
+# directories and files. If not, they will be downloaded
+# by this script automatically.
+#
+# - $dl_dir/alimeeting
+# This directory contains the following files downloaded from
+# https://openslr.org/62/
+#
+# - Train_Ali_far.tar.gz
+# - Train_Ali_near.tar.gz
+# - Test_Ali.tar.gz
+# - Eval_Ali.tar.gz
+#
+# - $dl_dir/musan
+# This directory contains the following directories downloaded from
+# http://www.openslr.org/17/
+#
+# - music
+# - noise
+# - speech
+
+dl_dir=$PWD/download
+
+. shared/parse_options.sh || exit 1
+
+# All files generated by this script are saved in "data".
+# You can safely remove "data" and rerun this script to regenerate it.
+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]}) $*"
+}
+
+log "dl_dir: $dl_dir"
+
+if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then
+ log "Stage 0: Download data"
+
+ if [ ! -f $dl_dir/alimeeting/Train_Ali_far.tar.gz ]; then
+ lhotse download ali-meeting $dl_dir/alimeeting
+ fi
+fi
+
+if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
+ log "Stage 1: Prepare alimeeting manifest"
+ # We assume that you have downloaded the alimeeting corpus
+ # to $dl_dir/alimeeting
+ for part in ihm sdm mdm; do
+ mkdir -p data/manifests/alimeeting
+ lhotse prepare ali-meeting --mic $part --save-mono --normalize-text m2met \
+ $dl_dir/alimeeting data/manifests
+ done
+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 $dl_dir/musan data/manifests
+fi
+
+if [ $stage -le 3 ] && [ $stop_stage -ge 3 ] && [ $use_gss = true ]; then
+ log "Stage 3: Apply GSS enhancement on MDM data (this stage requires a GPU)"
+ # We assume that you have installed the GSS package: https://github.com/desh2608/gss
+ local/prepare_alimeeting_gss.sh data/manifests exp/alimeeting_gss
+fi
+
+if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
+ log "Stage 4: Compute fbank for musan"
+ mkdir -p data/fbank
+ python local/compute_fbank_musan.py
+fi
+
+if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then
+ log "Stage 5: Compute fbank for alimeeting"
+ mkdir -p data/fbank
+ python local/compute_fbank_alimeeting.py
+ log "Combine features from train splits"
+ lhotse combine data/manifests/cuts_train_{ihm,ihm_rvb,sdm,gss}.jsonl.gz - | shuf |\
+ gzip -c > data/manifests/cuts_train_all.jsonl.gz
+fi
+
+if [ $stage -le 6 ] && [ $stop_stage -ge 6 ]; then
+ log "Stage 6: Prepare char based lang"
+ lang_char_dir=data/lang_char
+ mkdir -p $lang_char_dir
+
+ # Prepare text.
+ # Note: in Linux, you can install jq with the following command:
+ # wget -O jq https://github.com/stedolan/jq/releases/download/jq-1.6/jq-linux64
+ gunzip -c data/manifests/alimeeting-sdm_supervisions_train.jsonl.gz \
+ | jq ".text" | sed 's/"//g' \
+ | ./local/text2token.py -t "char" > $lang_char_dir/text
+
+ # Prepare words segments
+ python ./local/text2segments.py \
+ --input $lang_char_dir/text \
+ --output $lang_char_dir/text_words_segmentation
+
+ cat $lang_char_dir/text_words_segmentation | sed "s/ /\n/g" \
+ | sort -u | sed "/^$/d" \
+ | uniq > $lang_char_dir/words_no_ids.txt
+
+ # Prepare words.txt
+ if [ ! -f $lang_char_dir/words.txt ]; then
+ ./local/prepare_words.py \
+ --input-file $lang_char_dir/words_no_ids.txt \
+ --output-file $lang_char_dir/words.txt
+ fi
+
+ if [ ! -f $lang_char_dir/L_disambig.pt ]; then
+ ./local/prepare_char.py
+ fi
+fi
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/__init__.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/asr_datamodule.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/asr_datamodule.py
new file mode 100644
index 000000000..1cfd053c7
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/asr_datamodule.py
@@ -0,0 +1,419 @@
+# Copyright 2021 Piotr Żelasko
+#
+# 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.
+
+
+import argparse
+import logging
+import re
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import torch
+from lhotse import CutSet, Fbank, FbankConfig, load_manifest, load_manifest_lazy
+from lhotse.cut import Cut
+from lhotse.dataset import (
+ CutConcatenate,
+ CutMix,
+ DynamicBucketingSampler,
+ K2SpeechRecognitionDataset,
+ PrecomputedFeatures,
+ SpecAugment,
+)
+from lhotse.dataset.input_strategies import OnTheFlyFeatures
+from lhotse.utils import fix_random_seed
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from icefall.utils import str2bool
+
+
+class _SeedWorkers:
+ def __init__(self, seed: int):
+ self.seed = seed
+
+ def __call__(self, worker_id: int):
+ fix_random_seed(self.seed + worker_id)
+
+
+class AlimeetingAsrDataModule:
+ """
+ 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.
+ """
+
+ def __init__(self, args: argparse.Namespace):
+ self.args = args
+
+ @classmethod
+ def add_arguments(cls, parser: argparse.ArgumentParser):
+ 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(
+ "--manifest-dir",
+ type=Path,
+ default=Path("data/manifests"),
+ help="Path to directory with train/valid/test cuts.",
+ )
+ group.add_argument(
+ "--enable-musan",
+ type=str2bool,
+ default=True,
+ help=(
+ "When enabled, select noise from MUSAN and mix it "
+ "with training dataset. "
+ ),
+ )
+ group.add_argument(
+ "--concatenate-cuts",
+ type=str2bool,
+ default=False,
+ 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(
+ "--max-duration",
+ type=int,
+ default=100.0,
+ help=(
+ "Maximum pooled recordings duration (seconds) in a "
+ "single batch. You can reduce it if it causes CUDA OOM."
+ ),
+ )
+ group.add_argument(
+ "--max-cuts", type=int, default=None, help="Maximum cuts in a single batch."
+ )
+ group.add_argument(
+ "--num-buckets",
+ type=int,
+ default=50,
+ help=(
+ "The number of buckets for the BucketingSampler"
+ "(you might want to increase it for larger datasets)."
+ ),
+ )
+ 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."
+ ),
+ )
+ group.add_argument(
+ "--shuffle",
+ type=str2bool,
+ default=True,
+ help=(
+ "When enabled (=default), the examples will be "
+ "shuffled for each epoch."
+ ),
+ )
+
+ group.add_argument(
+ "--num-workers",
+ type=int,
+ default=8,
+ help=(
+ "The number of training dataloader workers that " "collect the batches."
+ ),
+ )
+ group.add_argument(
+ "--enable-spec-aug",
+ type=str2bool,
+ default=True,
+ help="When enabled, use SpecAugment for training dataset.",
+ )
+ group.add_argument(
+ "--spec-aug-time-warp-factor",
+ type=int,
+ default=80,
+ help=(
+ "Used only when --enable-spec-aug is True. "
+ "It specifies the factor for time warping in SpecAugment. "
+ "Larger values mean more warping. "
+ "A value less than 1 means to disable time warp."
+ ),
+ )
+
+ def train_dataloaders(
+ self,
+ cuts_train: CutSet,
+ sampler_state_dict: Optional[Dict[str, Any]] = None,
+ ) -> DataLoader:
+ """
+ Args:
+ cuts_train:
+ CutSet for training.
+ sampler_state_dict:
+ The state dict for the training sampler.
+ """
+ logging.info("About to get Musan cuts")
+
+ transforms = []
+ if self.args.enable_musan:
+ logging.info("Enable MUSAN")
+ cuts_musan = load_manifest(self.args.manifest_dir / "musan_cuts.jsonl.gz")
+ transforms.append(
+ CutMix(cuts=cuts_musan, prob=0.5, snr=(10, 20), preserve_id=True)
+ )
+ else:
+ logging.info("Disable MUSAN")
+
+ if self.args.concatenate_cuts:
+ logging.info(
+ "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 = []
+ if self.args.enable_spec_aug:
+ logging.info("Enable SpecAugment")
+ logging.info(f"Time warp factor: {self.args.spec_aug_time_warp_factor}")
+ input_transforms.append(
+ SpecAugment(
+ time_warp_factor=self.args.spec_aug_time_warp_factor,
+ num_frame_masks=2,
+ features_mask_size=27,
+ num_feature_masks=2,
+ frames_mask_size=100,
+ )
+ )
+ else:
+ logging.info("Disable SpecAugment")
+
+ logging.info("About to create train dataset")
+ if self.args.on_the_fly_feats:
+ train = K2SpeechRecognitionDataset(
+ cut_transforms=transforms,
+ input_strategy=OnTheFlyFeatures(Fbank(FbankConfig(num_mel_bins=80))),
+ input_transforms=input_transforms,
+ )
+ else:
+ train = K2SpeechRecognitionDataset(
+ cut_transforms=transforms,
+ input_transforms=input_transforms,
+ )
+
+ logging.info("Using DynamicBucketingSampler.")
+ train_sampler = DynamicBucketingSampler(
+ cuts_train,
+ max_duration=self.args.max_duration,
+ max_cuts=self.args.max_cuts,
+ shuffle=False,
+ num_buckets=self.args.num_buckets,
+ drop_last=True,
+ )
+ logging.info("About to create train dataloader")
+
+ if sampler_state_dict is not None:
+ logging.info("Loading sampler state dict")
+ train_sampler.load_state_dict(sampler_state_dict)
+
+ # 'seed' is derived from the current random state, which will have
+ # previously been set in the main process.
+ seed = torch.randint(0, 100000, ()).item()
+ worker_init_fn = _SeedWorkers(seed)
+
+ train_dl = DataLoader(
+ train,
+ sampler=train_sampler,
+ batch_size=None,
+ num_workers=self.args.num_workers,
+ persistent_workers=False,
+ worker_init_fn=worker_init_fn,
+ )
+
+ return train_dl
+
+ def valid_dataloaders(self, cuts_valid: CutSet) -> DataLoader:
+
+ transforms = []
+ if self.args.concatenate_cuts:
+ transforms = [
+ CutConcatenate(
+ duration_factor=self.args.duration_factor, gap=self.args.gap
+ )
+ ] + transforms
+
+ logging.info("About to create dev dataset")
+ if self.args.on_the_fly_feats:
+ validate = K2SpeechRecognitionDataset(
+ cut_transforms=transforms,
+ input_strategy=OnTheFlyFeatures(Fbank(FbankConfig(num_mel_bins=80))),
+ )
+ else:
+ validate = K2SpeechRecognitionDataset(
+ cut_transforms=transforms,
+ )
+ valid_sampler = DynamicBucketingSampler(
+ cuts_valid,
+ max_duration=self.args.max_duration,
+ shuffle=False,
+ )
+ logging.info("About to create dev dataloader")
+ valid_dl = DataLoader(
+ validate,
+ sampler=valid_sampler,
+ batch_size=None,
+ num_workers=2,
+ persistent_workers=False,
+ )
+
+ return valid_dl
+
+ def test_dataloaders(self, cuts: CutSet) -> DataLoader:
+ logging.debug("About to create test dataset")
+ test = K2SpeechRecognitionDataset(
+ input_strategy=OnTheFlyFeatures(Fbank(FbankConfig(num_mel_bins=80)))
+ if self.args.on_the_fly_feats
+ else PrecomputedFeatures(),
+ return_cuts=True,
+ )
+ sampler = DynamicBucketingSampler(
+ cuts, max_duration=self.args.max_duration, shuffle=False
+ )
+ logging.debug("About to create test dataloader")
+ test_dl = DataLoader(
+ test,
+ batch_size=None,
+ sampler=sampler,
+ num_workers=self.args.num_workers,
+ )
+ return test_dl
+
+ def remove_short_cuts(self, cut: Cut) -> bool:
+ """
+ See: https://github.com/k2-fsa/icefall/issues/500
+ Basically, the zipformer model subsamples the input using the following formula:
+ num_out_frames = ((num_in_frames - 7)//2 + 1)//2
+ For num_out_frames to be at least 1, num_in_frames must be at least 9.
+ """
+ return cut.duration >= 0.09
+
+ @lru_cache()
+ def train_cuts(self, sp: Optional[Any] = None) -> CutSet:
+ logging.info("About to get AMI train cuts")
+
+ def _remove_short_and_long_utt(c: Cut):
+ if c.duration < 0.1 or c.duration > 25.0:
+ return False
+
+ # In pruned RNN-T, we require that T >= S
+ # where T is the number of feature frames after subsampling
+ # and S is the number of tokens in the utterance
+
+ # In ./zipformer.py, the conv module uses the following expression
+ # for subsampling
+ T = ((c.num_frames - 7) // 2 + 1) // 2
+ tokens = c.supervisions[0].text
+ return T >= len(tokens)
+
+ cuts_train = load_manifest_lazy(
+ self.args.manifest_dir / "cuts_train_all.jsonl.gz"
+ )
+
+ return cuts_train.filter(_remove_short_and_long_utt)
+
+ @lru_cache()
+ def eval_ihm_cuts(self) -> CutSet:
+ logging.info("About to get AliMeeting IHM eval cuts")
+ cs = load_manifest_lazy(self.args.manifest_dir / "cuts_eval_ihm.jsonl.gz")
+ return cs.filter(self.remove_short_cuts)
+
+ @lru_cache()
+ def eval_sdm_cuts(self) -> CutSet:
+ logging.info("About to get AliMeeting SDM eval cuts")
+ cs = load_manifest_lazy(self.args.manifest_dir / "cuts_eval_sdm.jsonl.gz")
+ return cs.filter(self.remove_short_cuts)
+
+ @lru_cache()
+ def eval_gss_cuts(self) -> CutSet:
+ if not (self.args.manifest_dir / "cuts_eval_gss.jsonl.gz").exists():
+ logging.info("No GSS dev cuts found")
+ return None
+ logging.info("About to get AliMeeting GSS-enhanced eval cuts")
+ cs = load_manifest_lazy(self.args.manifest_dir / "cuts_eval_gss.jsonl.gz")
+ return cs.filter(self.remove_short_cuts)
+
+ @lru_cache()
+ def test_ihm_cuts(self) -> CutSet:
+ logging.info("About to get AliMeeting IHM test cuts")
+ cs = load_manifest_lazy(self.args.manifest_dir / "cuts_test_ihm.jsonl.gz")
+ return cs.filter(self.remove_short_cuts)
+
+ @lru_cache()
+ def test_sdm_cuts(self) -> CutSet:
+ logging.info("About to get AliMeeting SDM test cuts")
+ cs = load_manifest_lazy(self.args.manifest_dir / "cuts_test_sdm.jsonl.gz")
+ return cs.filter(self.remove_short_cuts)
+
+ @lru_cache()
+ def test_gss_cuts(self) -> CutSet:
+ if not (self.args.manifest_dir / "cuts_test_gss.jsonl.gz").exists():
+ logging.info("No GSS test cuts found")
+ return None
+ logging.info("About to get AliMeeting GSS-enhanced test cuts")
+ cs = load_manifest_lazy(self.args.manifest_dir / "cuts_test_gss.jsonl.gz")
+ return cs.filter(self.remove_short_cuts)
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/beam_search.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/beam_search.py
new file mode 120000
index 000000000..37516affc
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/beam_search.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7/beam_search.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/decode.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/decode.py
new file mode 100755
index 000000000..53381c1f4
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/decode.py
@@ -0,0 +1,698 @@
+#!/usr/bin/env python3
+#
+# Copyright 2021 Xiaomi Corporation (Author: Fangjun Kuang)
+#
+# 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.
+"""
+Usage:
+(1) greedy search
+./pruned_transducer_stateless7/decode.py \
+ --epoch 15 \
+ --avg 8 \
+ --exp-dir ./pruned_transducer_stateless7/exp \
+ --max-duration 500 \
+ --decoding-method greedy_search
+
+(2) modified beam search
+./pruned_transducer_stateless7/decode.py \
+ --epoch 15 \
+ --avg 8 \
+ --exp-dir ./pruned_transducer_stateless7/exp \
+ --max-duration 500 \
+ --decoding-method modified_beam_search \
+ --beam-size 4
+
+(3) fast beam search
+./pruned_transducer_stateless7/decode.py \
+ --epoch 15 \
+ --avg 8 \
+ --exp-dir ./pruned_transducer_stateless7/exp \
+ --max-duration 500 \
+ --decoding-method fast_beam_search \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8
+"""
+
+
+import argparse
+import logging
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import k2
+import sentencepiece as spm
+import torch
+import torch.nn as nn
+from asr_datamodule import AlimeetingAsrDataModule
+from beam_search import (
+ beam_search,
+ fast_beam_search_nbest_LG,
+ fast_beam_search_one_best,
+ greedy_search,
+ greedy_search_batch,
+ modified_beam_search,
+)
+from train import add_model_arguments, get_params, get_transducer_model
+
+from icefall import NgramLm
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from icefall.lexicon import Lexicon
+from icefall.utils import (
+ AttributeDict,
+ setup_logger,
+ store_transcripts,
+ str2bool,
+ write_error_stats,
+)
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=30,
+ help="""It specifies the checkpoint to use for decoding.
+ Note: Epoch counts from 0.
+ You can specify --avg to use more checkpoints for model averaging.""",
+ )
+
+ parser.add_argument(
+ "--iter",
+ type=int,
+ default=0,
+ help="""If positive, --epoch is ignored and it
+ will use the checkpoint exp_dir/checkpoint-iter.pt.
+ You can specify --avg to use more checkpoints for model averaging.
+ """,
+ )
+
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=10,
+ help="Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch' and '--iter'",
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=True,
+ help="Whether to load averaged model. Currently it only supports "
+ "using --epoch. If True, it would decode with the averaged model "
+ "over the epoch range from `epoch-avg` (excluded) to `epoch`."
+ "Actually only the models with epoch number of `epoch-avg` and "
+ "`epoch` are loaded for averaging. ",
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="pruned_transducer_stateless2/exp",
+ help="The experiment dir",
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=str,
+ default="data/lang_char",
+ help="""The lang dir
+ It contains language related input files such as
+ "lexicon.txt"
+ """,
+ )
+
+ parser.add_argument(
+ "--decoding-method",
+ type=str,
+ default="greedy_search",
+ help="""Possible values are:
+ - greedy_search
+ - beam_search
+ - modified_beam_search
+ - fast_beam_search
+ - fast_beam_search_nbest
+ - fast_beam_search_nbest_oracle
+ - fast_beam_search_nbest_LG
+ If you use fast_beam_search_nbest_LG, you have to specify
+ `--lang-dir`, which should contain `LG.pt`.
+ """,
+ )
+
+ parser.add_argument(
+ "--beam-size",
+ type=int,
+ default=4,
+ help="""An interger indicating how many candidates we will keep for each
+ frame. Used only when --decoding-method is beam_search or
+ modified_beam_search.""",
+ )
+
+ parser.add_argument(
+ "--beam",
+ type=float,
+ default=4,
+ help="""A floating point value to calculate the cutoff score during beam
+ search (i.e., `cutoff = max-score - beam`), which is the same as the
+ `beam` in Kaldi.
+ Used only when --decoding-method is fast_beam_search""",
+ )
+
+ parser.add_argument(
+ "--ngram-lm-scale",
+ type=float,
+ default=0.01,
+ help="""
+ Used only when --decoding_method is fast_beam_search_nbest_LG.
+ It specifies the scale for n-gram LM scores.
+ """,
+ )
+
+ parser.add_argument(
+ "--max-contexts",
+ type=int,
+ default=8,
+ help="""Used only when --decoding-method is
+ fast_beam_search, fast_beam_search_nbest, fast_beam_search_nbest_LG,
+ and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--max-states",
+ type=int,
+ default=64,
+ help="""Used only when --decoding-method is
+ fast_beam_search, fast_beam_search_nbest, fast_beam_search_nbest_LG,
+ and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram; " "2 means tri-gram",
+ )
+ parser.add_argument(
+ "--max-sym-per-frame",
+ type=int,
+ default=1,
+ help="""Maximum number of symbols per frame.
+ Used only when --decoding_method is greedy_search""",
+ )
+
+ parser.add_argument(
+ "--num-paths",
+ type=int,
+ default=200,
+ help="""Number of paths for nbest decoding.
+ Used only when the decoding method is fast_beam_search_nbest,
+ fast_beam_search_nbest_LG, and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--nbest-scale",
+ type=float,
+ default=0.5,
+ help="""Scale applied to lattice scores when computing nbest paths.
+ Used only when the decoding method is fast_beam_search_nbest,
+ fast_beam_search_nbest_LG, and fast_beam_search_nbest_oracle""",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def decode_one_batch(
+ params: AttributeDict,
+ model: nn.Module,
+ lexicon: Lexicon,
+ batch: dict,
+ decoding_graph: Optional[k2.Fsa] = None,
+) -> Dict[str, List[List[str]]]:
+ """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 greedy_search is used, it would be "greedy_search"
+ If beam search with a beam size of 7 is used, it would be
+ "beam_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`.
+ model:
+ The neural model.
+ batch:
+ It is the return value from iterating
+ `lhotse.dataset.K2SpeechRecognitionDataset`. See its documentation
+ for the format of the `batch`.
+ decoding_graph:
+ The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+ only when --decoding_method is fast_beam_search.
+ Returns:
+ Return the decoding result. See above description for the format of
+ the returned dict.
+ """
+ device = model.device
+ feature = batch["inputs"]
+ assert feature.ndim == 3
+
+ feature = feature.to(device)
+ # at entry, feature is (N, T, C)
+
+ supervisions = batch["supervisions"]
+ feature_lens = supervisions["num_frames"].to(device)
+
+ encoder_out, encoder_out_lens = model.encoder(x=feature, x_lens=feature_lens)
+ hyps = []
+
+ if params.decoding_method == "fast_beam_search":
+ hyp_tokens = fast_beam_search_one_best(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ )
+ for i in range(encoder_out.size(0)):
+ hyps.append([lexicon.token_table[idx] for idx in hyp_tokens[i]])
+ elif params.decoding_method == "fast_beam_search_nbest_LG":
+ hyp_tokens = fast_beam_search_nbest_LG(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ num_paths=params.num_paths,
+ nbest_scale=params.nbest_scale,
+ )
+ for i in range(encoder_out.size(0)):
+ hyps.append([lexicon.token_table[idx] for idx in hyp_tokens[i]])
+ elif params.decoding_method == "greedy_search" and params.max_sym_per_frame == 1:
+ hyp_tokens = greedy_search_batch(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ )
+ for i in range(encoder_out.size(0)):
+ hyps.append([lexicon.token_table[idx] for idx in hyp_tokens[i]])
+ elif params.decoding_method == "modified_beam_search":
+ hyp_tokens = modified_beam_search(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam_size,
+ )
+ for i in range(encoder_out.size(0)):
+ hyps.append([lexicon.token_table[idx] for idx in hyp_tokens[i]])
+ else:
+ batch_size = encoder_out.size(0)
+
+ for i in range(batch_size):
+ # fmt: off
+ encoder_out_i = encoder_out[i:i+1, :encoder_out_lens[i]]
+ # fmt: on
+ if params.decoding_method == "greedy_search":
+ hyp = greedy_search(
+ model=model,
+ encoder_out=encoder_out_i,
+ max_sym_per_frame=params.max_sym_per_frame,
+ )
+ elif params.decoding_method == "beam_search":
+ hyp = beam_search(
+ model=model,
+ encoder_out=encoder_out_i,
+ beam=params.beam_size,
+ )
+ else:
+ raise ValueError(
+ f"Unsupported decoding method: {params.decoding_method}"
+ )
+ hyps.append([lexicon.token_table[idx] for idx in hyp])
+
+ if params.decoding_method == "greedy_search":
+ return {"greedy_search": hyps}
+ elif params.decoding_method == "fast_beam_search":
+ return {
+ (
+ f"beam_{params.beam}_"
+ f"max_contexts_{params.max_contexts}_"
+ f"max_states_{params.max_states}"
+ ): hyps
+ }
+ elif "fast_beam_search" in params.decoding_method:
+ key = f"beam_{params.beam}_"
+ key += f"max_contexts_{params.max_contexts}_"
+ key += f"max_states_{params.max_states}"
+ if "nbest" in params.decoding_method:
+ key += f"_num_paths_{params.num_paths}_"
+ key += f"nbest_scale_{params.nbest_scale}"
+ if "LG" in params.decoding_method:
+ key += f"_ngram_lm_scale_{params.ngram_lm_scale}"
+
+ return {key: hyps}
+ else:
+ return {f"beam_size_{params.beam_size}": hyps}
+
+
+def decode_dataset(
+ dl: torch.utils.data.DataLoader,
+ params: AttributeDict,
+ model: nn.Module,
+ lexicon: Lexicon,
+ decoding_graph: Optional[k2.Fsa] = None,
+) -> Dict[str, List[Tuple[List[str], List[str]]]]:
+ """Decode dataset.
+
+ Args:
+ dl:
+ PyTorch's dataloader containing the dataset to decode.
+ params:
+ It is returned by :func:`get_params`.
+ model:
+ The neural model.
+ decoding_graph:
+ The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+ only when --decoding_method is fast_beam_search.
+ Returns:
+ Return a dict, whose key may be "greedy_search" if greedy search
+ is used, or it may be "beam_7" if beam size of 7 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.
+ """
+ num_cuts = 0
+
+ try:
+ num_batches = len(dl)
+ except TypeError:
+ num_batches = "?"
+
+ if params.decoding_method == "greedy_search":
+ log_interval = 100
+ else:
+ log_interval = 2
+
+ results = defaultdict(list)
+ for batch_idx, batch in enumerate(dl):
+ texts = batch["supervisions"]["text"]
+ texts = [list(str(text).replace(" ", "")) for text in texts]
+ cut_ids = [cut.id for cut in batch["supervisions"]["cut"]]
+
+ hyps_dict = decode_one_batch(
+ params=params,
+ model=model,
+ lexicon=lexicon,
+ decoding_graph=decoding_graph,
+ batch=batch,
+ )
+
+ for name, hyps in hyps_dict.items():
+ this_batch = []
+ assert len(hyps) == len(texts)
+ for cut_id, hyp_words, ref_text in zip(cut_ids, hyps, texts):
+ this_batch.append((cut_id, ref_text, hyp_words))
+
+ results[name].extend(this_batch)
+
+ num_cuts += len(texts)
+
+ if batch_idx % log_interval == 0:
+ batch_str = f"{batch_idx}/{num_batches}"
+
+ logging.info(f"batch {batch_str}, cuts processed until now is {num_cuts}")
+ return results
+
+
+def save_results(
+ params: AttributeDict,
+ test_set_name: str,
+ results_dict: Dict[str, List[Tuple[str, List[str], List[str]]]],
+):
+ test_set_wers = dict()
+ for key, results in results_dict.items():
+ recog_path = (
+ params.res_dir / f"recogs-{test_set_name}-{key}-{params.suffix}.txt"
+ )
+ results = sorted(results)
+ 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.res_dir / f"errs-{test_set_name}-{key}-{params.suffix}.txt"
+ )
+ with open(errs_filename, "w") as f:
+ wer = write_error_stats(
+ f, f"{test_set_name}-{key}", results, enable_log=True
+ )
+ 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.res_dir / f"wer-summary-{test_set_name}-{key}-{params.suffix}.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()
+ AlimeetingAsrDataModule.add_arguments(parser)
+ args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ params = get_params()
+ params.update(vars(args))
+
+ assert params.decoding_method in (
+ "greedy_search",
+ "beam_search",
+ "fast_beam_search",
+ "fast_beam_search_nbest_LG",
+ "modified_beam_search",
+ )
+ params.res_dir = params.exp_dir / params.decoding_method
+
+ if params.iter > 0:
+ params.suffix = f"iter-{params.iter}-avg-{params.avg}"
+ else:
+ params.suffix = f"epoch-{params.epoch}-avg-{params.avg}"
+
+ if "fast_beam_search" in params.decoding_method:
+ params.suffix += f"-beam-{params.beam}"
+ params.suffix += f"-max-contexts-{params.max_contexts}"
+ params.suffix += f"-max-states-{params.max_states}"
+ if "nbest" in params.decoding_method:
+ params.suffix += f"-nbest-scale-{params.nbest_scale}"
+ params.suffix += f"-num-paths-{params.num_paths}"
+ if "LG" in params.decoding_method:
+ params.suffix += f"-ngram-lm-scale-{params.ngram_lm_scale}"
+ elif "beam_search" in params.decoding_method:
+ params.suffix += f"-{params.decoding_method}-beam-size-{params.beam_size}"
+ else:
+ params.suffix += f"-context-{params.context_size}"
+ params.suffix += f"-max-sym-per-frame-{params.max_sym_per_frame}"
+
+ setup_logger(f"{params.res_dir}/log-decode-{params.suffix}")
+ logging.info("Decoding started")
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"Device: {device}")
+
+ lexicon = Lexicon(params.lang_dir)
+ params.blank_id = lexicon.token_table[""]
+ params.vocab_size = max(lexicon.tokens) + 1
+
+ logging.info(params)
+
+ logging.info("About to create model")
+ model = get_transducer_model(params)
+
+ if not params.use_averaged_model:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ elif 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 i >= 1:
+ filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg + 1
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg + 1:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ filename_start = filenames[-1]
+ filename_end = filenames[0]
+ logging.info(
+ "Calculating the averaged model over iteration checkpoints"
+ f" from {filename_start} (excluded) to {filename_end}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+ else:
+ assert params.avg > 0, params.avg
+ start = params.epoch - params.avg
+ assert start >= 1, start
+ filename_start = f"{params.exp_dir}/epoch-{start}.pt"
+ filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
+ logging.info(
+ f"Calculating the averaged model over epoch range from "
+ f"{start} (excluded) to {params.epoch}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+
+ model.to(device)
+ model.eval()
+ model.device = device
+
+ if "fast_beam_search" in params.decoding_method:
+ decoding_graph = k2.trivial_graph(params.vocab_size - 1, device=device)
+ else:
+ decoding_graph = None
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ alimeeting = AlimeetingAsrDataModule(args)
+
+ eval_ihm_cuts = alimeeting.eval_ihm_cuts()
+ test_ihm_cuts = alimeeting.test_ihm_cuts()
+ eval_sdm_cuts = alimeeting.eval_sdm_cuts()
+ test_sdm_cuts = alimeeting.test_sdm_cuts()
+ eval_gss_cuts = alimeeting.eval_gss_cuts()
+ test_gss_cuts = alimeeting.test_gss_cuts()
+
+ eval_ihm_dl = alimeeting.test_dataloaders(eval_ihm_cuts)
+ test_ihm_dl = alimeeting.test_dataloaders(test_ihm_cuts)
+ eval_sdm_dl = alimeeting.test_dataloaders(eval_sdm_cuts)
+ test_sdm_dl = alimeeting.test_dataloaders(test_sdm_cuts)
+ if eval_gss_cuts is not None:
+ eval_gss_dl = alimeeting.test_dataloaders(eval_gss_cuts)
+ if test_gss_cuts is not None:
+ test_gss_dl = alimeeting.test_dataloaders(test_gss_cuts)
+
+ test_sets = {
+ "eval_ihm": (eval_ihm_dl, eval_ihm_cuts),
+ "test_ihm": (test_ihm_dl, test_ihm_cuts),
+ "eval_sdm": (eval_sdm_dl, eval_sdm_cuts),
+ "test_sdm": (test_sdm_dl, test_sdm_cuts),
+ }
+ if eval_gss_cuts is not None:
+ test_sets["eval_gss"] = (eval_gss_dl, eval_gss_cuts)
+ if test_gss_cuts is not None:
+ test_sets["test_gss"] = (test_gss_dl, test_gss_cuts)
+
+ for test_set in test_sets:
+ logging.info(f"Decoding {test_set}")
+ dl, cuts = test_sets[test_set]
+ results_dict = decode_dataset(
+ dl=dl,
+ params=params,
+ model=model,
+ lexicon=lexicon,
+ decoding_graph=decoding_graph,
+ )
+
+ save_results(
+ params=params,
+ test_set_name=test_set,
+ results_dict=results_dict,
+ )
+
+ logging.info("Done!")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/decoder.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/decoder.py
new file mode 120000
index 000000000..8283d8c5a
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/decoder.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7/decoder.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/encoder_interface.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/encoder_interface.py
new file mode 120000
index 000000000..0c2673d46
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/encoder_interface.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7/encoder_interface.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/export.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/export.py
new file mode 100755
index 000000000..23a88dd29
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/export.py
@@ -0,0 +1,320 @@
+#!/usr/bin/env python3
+#
+# Copyright 2021 Xiaomi Corporation (Author: Fangjun Kuang)
+#
+# 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.
+
+# This script converts several saved checkpoints
+# to a single one using model averaging.
+"""
+
+Usage:
+
+(1) Export to torchscript model using torch.jit.script()
+
+./pruned_transducer_stateless7/export.py \
+ --exp-dir ./pruned_transducer_stateless7/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 9 \
+ --jit 1
+
+It will generate a file `cpu_jit.pt` in the given `exp_dir`. You can later
+load it by `torch.jit.load("cpu_jit.pt")`.
+
+Note `cpu` in the name `cpu_jit.pt` means the parameters when loaded into Python
+are on CPU. You can use `to("cuda")` to move them to a CUDA device.
+
+Check
+https://github.com/k2-fsa/sherpa
+for how to use the exported models outside of icefall.
+
+(2) Export `model.state_dict()`
+
+./pruned_transducer_stateless7/export.py \
+ --exp-dir ./pruned_transducer_stateless7/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 20 \
+ --avg 10
+
+It will generate a file `pretrained.pt` in the given `exp_dir`. You can later
+load it by `icefall.checkpoint.load_checkpoint()`.
+
+To use the generated file with `pruned_transducer_stateless7/decode.py`,
+you can do:
+
+ cd /path/to/exp_dir
+ ln -s pretrained.pt epoch-9999.pt
+
+ cd /path/to/egs/librispeech/ASR
+ ./pruned_transducer_stateless7/decode.py \
+ --exp-dir ./pruned_transducer_stateless7/exp \
+ --epoch 9999 \
+ --avg 1 \
+ --max-duration 600 \
+ --decoding-method greedy_search \
+ --bpe-model data/lang_bpe_500/bpe.model
+
+Check ./pretrained.py for its usage.
+
+Note: If you don't want to train a model from scratch, we have
+provided one for you. You can get it at
+
+https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless7-2022-11-11
+
+with the following commands:
+
+ sudo apt-get install git-lfs
+ git lfs install
+ git clone https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless7-2022-11-11
+ # You will find the pre-trained model in icefall-asr-librispeech-pruned-transducer-stateless7-2022-11-11/exp
+"""
+
+import argparse
+import logging
+from pathlib import Path
+
+import sentencepiece as spm
+import torch
+import torch.nn as nn
+from scaling_converter import convert_scaled_to_non_scaled
+from train import add_model_arguments, get_params, get_transducer_model
+
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from icefall.lexicon import Lexicon
+from icefall.utils import str2bool
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=15,
+ help="""It specifies the checkpoint to use for decoding.
+ Note: Epoch counts from 1.
+ You can specify --avg to use more checkpoints for model averaging.""",
+ )
+
+ parser.add_argument(
+ "--iter",
+ type=int,
+ default=0,
+ help="""If positive, --epoch is ignored and it
+ will use the checkpoint exp_dir/checkpoint-iter.pt.
+ You can specify --avg to use more checkpoints for model averaging.
+ """,
+ )
+
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=8,
+ help="Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch' and '--iter'",
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=True,
+ help="Whether to load averaged model. Currently it only supports "
+ "using --epoch. If True, it would decode with the averaged model "
+ "over the epoch range from `epoch-avg` (excluded) to `epoch`."
+ "Actually only the models with epoch number of `epoch-avg` and "
+ "`epoch` are loaded for averaging. ",
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="pruned_transducer_stateless7/exp",
+ help="""It specifies the directory where all training related
+ files, e.g., checkpoints, log, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=str,
+ default="data/lang_char",
+ help="The lang dir",
+ )
+
+ parser.add_argument(
+ "--jit",
+ type=str2bool,
+ default=False,
+ help="""True to save a model after applying torch.jit.script.
+ It will generate a file named cpu_jit.pt
+
+ Check ./jit_pretrained.py for how to use it.
+ """,
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram; 2 means tri-gram",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+@torch.no_grad()
+def main():
+ args = get_parser().parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ params = get_params()
+ params.update(vars(args))
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"device: {device}")
+
+ lexicon = Lexicon(params.lang_dir)
+
+ params.blank_id = 0
+ params.vocab_size = max(lexicon.tokens) + 1
+
+ logging.info(params)
+
+ logging.info("About to create model")
+ model = get_transducer_model(params)
+
+ model.to(device)
+
+ if not params.use_averaged_model:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ elif 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 i >= 1:
+ filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg + 1
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg + 1:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ filename_start = filenames[-1]
+ filename_end = filenames[0]
+ logging.info(
+ "Calculating the averaged model over iteration checkpoints"
+ f" from {filename_start} (excluded) to {filename_end}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+ else:
+ assert params.avg > 0, params.avg
+ start = params.epoch - params.avg
+ assert start >= 1, start
+ filename_start = f"{params.exp_dir}/epoch-{start}.pt"
+ filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
+ logging.info(
+ f"Calculating the averaged model over epoch range from "
+ f"{start} (excluded) to {params.epoch}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+
+ model.to("cpu")
+ model.eval()
+
+ if params.jit is True:
+ convert_scaled_to_non_scaled(model, inplace=True)
+ logging.info("Using torch.jit.script()")
+ # We won't use the forward() method of the model in C++, so just ignore
+ # it here.
+ # Otherwise, one of its arguments is a ragged tensor and is not
+ # torch scriptabe.
+ model.__class__.forward = torch.jit.ignore(model.__class__.forward)
+ logging.info("Using torch.jit.script")
+ model = torch.jit.script(model)
+ filename = params.exp_dir / "cpu_jit.pt"
+ model.save(str(filename))
+ logging.info(f"Saved to {filename}")
+ else:
+ logging.info("Not using torchscript. Export model.state_dict()")
+ # Save it using a format so that it can be loaded
+ # by :func:`load_checkpoint`
+ filename = params.exp_dir / "pretrained.pt"
+ torch.save({"model": model.state_dict()}, str(filename))
+ logging.info(f"Saved to {filename}")
+
+
+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/alimeeting/ASR_v2/pruned_transducer_stateless7/jit_pretrained.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/jit_pretrained.py
new file mode 120000
index 000000000..a44034e34
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/jit_pretrained.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7/jit_pretrained.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/joiner.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/joiner.py
new file mode 120000
index 000000000..0f0c3c90a
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/joiner.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7/joiner.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/model.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/model.py
new file mode 120000
index 000000000..0d8bc665b
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/model.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7/model.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/optim.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/optim.py
new file mode 120000
index 000000000..8a05abb5f
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/optim.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7/optim.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/pretrained.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/pretrained.py
new file mode 120000
index 000000000..068f0f57f
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/pretrained.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7/pretrained.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/scaling.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/scaling.py
new file mode 120000
index 000000000..5f9be9fe0
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/scaling.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7/scaling.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/scaling_converter.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/scaling_converter.py
new file mode 120000
index 000000000..f9960e5c6
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/scaling_converter.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7/scaling_converter.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/test_model.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/test_model.py
new file mode 120000
index 000000000..7ceac5d10
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/test_model.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7/test_model.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/train.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/train.py
new file mode 100755
index 000000000..757d6535e
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/train.py
@@ -0,0 +1,1186 @@
+#!/usr/bin/env python3
+# Copyright 2021-2022 Xiaomi Corp. (authors: Fangjun Kuang,
+# Wei Kang,
+# Mingshuang Luo,)
+# Zengwei Yao)
+#
+# 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.
+"""
+Usage:
+
+export CUDA_VISIBLE_DEVICES="0,1,2,3"
+
+./pruned_transducer_stateless7/train.py \
+ --world-size 4 \
+ --num-epochs 15 \
+ --start-epoch 1 \
+ --exp-dir pruned_transducer_stateless7/exp \
+ --max-duration 150 \
+ --use-fp16 True
+
+"""
+
+
+import argparse
+import copy
+import logging
+import warnings
+from pathlib import Path
+from shutil import copyfile
+from typing import Any, Dict, Optional, Tuple, Union
+
+import k2
+import optim
+import sentencepiece as spm
+import torch
+import torch.multiprocessing as mp
+import torch.nn as nn
+from asr_datamodule import AlimeetingAsrDataModule
+from decoder import Decoder
+from joiner import Joiner
+from lhotse.dataset.sampling.base import CutSampler
+from lhotse.utils import fix_random_seed
+from model import Transducer
+from optim import Eden, ScaledAdam
+from torch import Tensor
+from torch.cuda.amp import GradScaler
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.tensorboard import SummaryWriter
+from zipformer import Zipformer
+
+from icefall import diagnostics
+from icefall.char_graph_compiler import CharCtcTrainingGraphCompiler
+from icefall.checkpoint import load_checkpoint, remove_checkpoints
+from icefall.checkpoint import save_checkpoint as save_checkpoint_impl
+from icefall.checkpoint import (
+ save_checkpoint_with_global_batch_idx,
+ update_averaged_model,
+)
+from icefall.dist import cleanup_dist, setup_dist
+from icefall.env import get_env_info
+from icefall.hooks import register_inf_check_hooks
+from icefall.lexicon import Lexicon
+from icefall.utils import AttributeDict, MetricsTracker, setup_logger, str2bool
+
+LRSchedulerType = Union[torch.optim.lr_scheduler._LRScheduler, optim.LRScheduler]
+
+
+def set_batch_count(model: Union[nn.Module, DDP], batch_count: float) -> None:
+ if isinstance(model, DDP):
+ # get underlying nn.Module
+ model = model.module
+ for module in model.modules():
+ if hasattr(module, "batch_count"):
+ module.batch_count = batch_count
+
+
+def add_model_arguments(parser: argparse.ArgumentParser):
+ parser.add_argument(
+ "--num-encoder-layers",
+ type=str,
+ default="2,4,3,2,4",
+ help="Number of zipformer encoder layers, comma separated.",
+ )
+
+ parser.add_argument(
+ "--feedforward-dims",
+ type=str,
+ default="1024,1024,2048,2048,1024",
+ help="Feedforward dimension of the zipformer encoder layers, comma separated.",
+ )
+
+ parser.add_argument(
+ "--nhead",
+ type=str,
+ default="8,8,8,8,8",
+ help="Number of attention heads in the zipformer encoder layers.",
+ )
+
+ parser.add_argument(
+ "--encoder-dims",
+ type=str,
+ default="384,384,384,384,384",
+ help="Embedding dimension in the 2 blocks of zipformer encoder layers, comma separated",
+ )
+
+ parser.add_argument(
+ "--attention-dims",
+ type=str,
+ default="192,192,192,192,192",
+ help="""Attention dimension in the 2 blocks of zipformer encoder layers, comma separated;
+ not the same as embedding dimension.""",
+ )
+
+ parser.add_argument(
+ "--encoder-unmasked-dims",
+ type=str,
+ default="256,256,256,256,256",
+ help="Unmasked dimensions in the encoders, relates to augmentation during training. "
+ "Must be <= each of encoder_dims. Empirically, less than 256 seems to make performance "
+ " worse.",
+ )
+
+ parser.add_argument(
+ "--zipformer-downsampling-factors",
+ type=str,
+ default="1,2,4,8,2",
+ help="Downsampling factor for each stack of encoder layers.",
+ )
+
+ parser.add_argument(
+ "--cnn-module-kernels",
+ type=str,
+ default="31,31,31,31,31",
+ help="Sizes of kernels in convolution modules",
+ )
+
+ parser.add_argument(
+ "--decoder-dim",
+ type=int,
+ default=512,
+ help="Embedding dimension in the decoder model.",
+ )
+
+ parser.add_argument(
+ "--joiner-dim",
+ type=int,
+ default=512,
+ help="""Dimension used in the joiner model.
+ Outputs from the encoder and decoder model are projected
+ to this dimension before adding.
+ """,
+ )
+
+
+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.",
+ )
+
+ parser.add_argument(
+ "--num-epochs",
+ type=int,
+ default=15,
+ help="Number of epochs to train.",
+ )
+
+ parser.add_argument(
+ "--start-epoch",
+ type=int,
+ default=1,
+ help="""Resume training from this epoch. It should be positive.
+ If larger than 1, it will load checkpoint from
+ exp-dir/epoch-{start_epoch-1}.pt
+ """,
+ )
+
+ parser.add_argument(
+ "--start-batch",
+ type=int,
+ default=0,
+ help="""If positive, --start-epoch is ignored and
+ it loads the checkpoint from exp-dir/checkpoint-{start_batch}.pt
+ """,
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="pruned_transducer_stateless7/exp",
+ help="""The experiment dir.
+ It specifies the directory where all training related
+ files, e.g., checkpoints, log, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=str,
+ default="data/lang_char",
+ help="""The lang dir
+ It contains language related input files such as
+ "lexicon.txt"
+ """,
+ )
+
+ parser.add_argument(
+ "--base-lr", type=float, default=0.05, help="The base learning rate."
+ )
+
+ parser.add_argument(
+ "--lr-batches",
+ type=float,
+ default=5000,
+ help="""Number of steps that affects how rapidly the learning rate
+ decreases. We suggest not to change this.""",
+ )
+
+ parser.add_argument(
+ "--lr-epochs",
+ type=float,
+ default=3.5,
+ help="""Number of epochs that affects how rapidly the learning rate decreases.
+ """,
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram; " "2 means tri-gram",
+ )
+
+ parser.add_argument(
+ "--prune-range",
+ type=int,
+ default=5,
+ help="The prune range for rnnt loss, it means how many symbols(context)"
+ "we are using to compute the loss",
+ )
+
+ parser.add_argument(
+ "--lm-scale",
+ type=float,
+ default=0.25,
+ help="The scale to smooth the loss with lm "
+ "(output of prediction network) part.",
+ )
+
+ parser.add_argument(
+ "--am-scale",
+ type=float,
+ default=0.0,
+ help="The scale to smooth the loss with am (output of encoder network)" "part.",
+ )
+
+ parser.add_argument(
+ "--simple-loss-scale",
+ type=float,
+ default=0.5,
+ help="To get pruning ranges, we will calculate a simple version"
+ "loss(joiner is just addition), this simple loss also uses for"
+ "training (as a regularization item). We will scale the simple loss"
+ "with this parameter before adding to the final loss.",
+ )
+
+ parser.add_argument(
+ "--seed",
+ type=int,
+ default=42,
+ help="The seed for random generators intended for reproducibility",
+ )
+
+ parser.add_argument(
+ "--print-diagnostics",
+ type=str2bool,
+ default=False,
+ help="Accumulate stats on activations, print them and exit.",
+ )
+
+ parser.add_argument(
+ "--inf-check",
+ type=str2bool,
+ default=False,
+ help="Add hooks to check for infinite module outputs and gradients.",
+ )
+
+ parser.add_argument(
+ "--save-every-n",
+ type=int,
+ default=5000,
+ help="""Save checkpoint after processing this number of batches"
+ periodically. We save checkpoint to exp-dir/ whenever
+ params.batch_idx_train % save_every_n == 0. The checkpoint filename
+ has the form: f'exp-dir/checkpoint-{params.batch_idx_train}.pt'
+ Note: It also saves checkpoint to `exp-dir/epoch-xxx.pt` at the
+ end of each epoch where `xxx` is the epoch number counting from 0.
+ """,
+ )
+
+ parser.add_argument(
+ "--keep-last-k",
+ type=int,
+ default=10,
+ help="""Only keep this number of checkpoints on disk.
+ For instance, if it is 3, there are only 3 checkpoints
+ in the exp-dir with filenames `checkpoint-xxx.pt`.
+ It does not affect checkpoints with name `epoch-xxx.pt`.
+ """,
+ )
+
+ parser.add_argument(
+ "--average-period",
+ type=int,
+ default=200,
+ help="""Update the averaged model, namely `model_avg`, after processing
+ this number of batches. `model_avg` is a separate version of model,
+ in which each floating-point parameter is the average of all the
+ parameters from the start of training. Each time we take the average,
+ we do: `model_avg = model * (average_period / batch_idx_train) +
+ model_avg * ((batch_idx_train - average_period) / batch_idx_train)`.
+ """,
+ )
+
+ parser.add_argument(
+ "--use-fp16",
+ type=str2bool,
+ default=False,
+ help="Whether to use half precision training.",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def get_params() -> AttributeDict:
+ """Return a dict containing training parameters.
+
+ All training related parameters that are not passed from the commandline
+ are 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`:
+
+ - 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
+
+ - reset_interval: Reset statistics if batch_idx % reset_interval is 0
+
+ - valid_interval: Run validation if batch_idx % valid_interval is 0
+
+ - feature_dim: The model input dim. It has to match the one used
+ in computing features.
+
+ - subsampling_factor: The subsampling factor for the model.
+
+ - encoder_dim: Hidden dim for multi-head attention model.
+
+ - num_decoder_layers: Number of decoder layer of transformer decoder.
+
+ - warm_step: The warmup period that dictates the decay of the
+ scale on "simple" (un-pruned) loss.
+ """
+ params = AttributeDict(
+ {
+ "best_train_loss": float("inf"),
+ "best_valid_loss": float("inf"),
+ "best_train_epoch": -1,
+ "best_valid_epoch": -1,
+ "batch_idx_train": 0,
+ "log_interval": 100,
+ "reset_interval": 200,
+ "valid_interval": 3000, # For the 100h subset, use 800
+ # parameters for zipformer
+ "feature_dim": 80,
+ "subsampling_factor": 4, # not passed in, this is fixed.
+ "warm_step": 2000,
+ "env_info": get_env_info(),
+ }
+ )
+
+ return params
+
+
+def get_encoder_model(params: AttributeDict) -> nn.Module:
+ # TODO: We can add an option to switch between Zipformer and Transformer
+ def to_int_tuple(s: str):
+ return tuple(map(int, s.split(",")))
+
+ encoder = Zipformer(
+ num_features=params.feature_dim,
+ output_downsampling_factor=2,
+ zipformer_downsampling_factors=to_int_tuple(
+ params.zipformer_downsampling_factors
+ ),
+ encoder_dims=to_int_tuple(params.encoder_dims),
+ attention_dim=to_int_tuple(params.attention_dims),
+ encoder_unmasked_dims=to_int_tuple(params.encoder_unmasked_dims),
+ nhead=to_int_tuple(params.nhead),
+ feedforward_dim=to_int_tuple(params.feedforward_dims),
+ cnn_module_kernels=to_int_tuple(params.cnn_module_kernels),
+ num_encoder_layers=to_int_tuple(params.num_encoder_layers),
+ )
+ return encoder
+
+
+def get_decoder_model(params: AttributeDict) -> nn.Module:
+ decoder = Decoder(
+ vocab_size=params.vocab_size,
+ decoder_dim=params.decoder_dim,
+ blank_id=params.blank_id,
+ context_size=params.context_size,
+ )
+ return decoder
+
+
+def get_joiner_model(params: AttributeDict) -> nn.Module:
+ joiner = Joiner(
+ encoder_dim=int(params.encoder_dims.split(",")[-1]),
+ decoder_dim=params.decoder_dim,
+ joiner_dim=params.joiner_dim,
+ vocab_size=params.vocab_size,
+ )
+ return joiner
+
+
+def get_transducer_model(params: AttributeDict) -> nn.Module:
+ encoder = get_encoder_model(params)
+ decoder = get_decoder_model(params)
+ joiner = get_joiner_model(params)
+
+ model = Transducer(
+ encoder=encoder,
+ decoder=decoder,
+ joiner=joiner,
+ encoder_dim=int(params.encoder_dims.split(",")[-1]),
+ decoder_dim=params.decoder_dim,
+ joiner_dim=params.joiner_dim,
+ vocab_size=params.vocab_size,
+ )
+ return model
+
+
+def load_checkpoint_if_available(
+ params: AttributeDict,
+ model: nn.Module,
+ model_avg: nn.Module = None,
+ optimizer: Optional[torch.optim.Optimizer] = None,
+ scheduler: Optional[LRSchedulerType] = None,
+) -> Optional[Dict[str, Any]]:
+ """Load checkpoint from file.
+
+ If params.start_batch is positive, it will load the checkpoint from
+ `params.exp_dir/checkpoint-{params.start_batch}.pt`. Otherwise, if
+ params.start_epoch is larger than 1, it will load the checkpoint from
+ `params.start_epoch - 1`.
+
+ Apart from loading state dict for `model` and `optimizer` 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.
+ model_avg:
+ The stored model averaged from the start of training.
+ optimizer:
+ The optimizer that we are using.
+ scheduler:
+ The scheduler that we are using.
+ Returns:
+ Return a dict containing previously saved training info.
+ """
+ if params.start_batch > 0:
+ filename = params.exp_dir / f"checkpoint-{params.start_batch}.pt"
+ elif params.start_epoch > 1:
+ filename = params.exp_dir / f"epoch-{params.start_epoch-1}.pt"
+ else:
+ return None
+
+ assert filename.is_file(), f"{filename} does not exist!"
+
+ saved_params = load_checkpoint(
+ filename,
+ model=model,
+ model_avg=model_avg,
+ 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]
+
+ if params.start_batch > 0:
+ if "cur_epoch" in saved_params:
+ params["start_epoch"] = saved_params["cur_epoch"]
+
+ if "cur_batch_idx" in saved_params:
+ params["cur_batch_idx"] = saved_params["cur_batch_idx"]
+
+ return saved_params
+
+
+def save_checkpoint(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ model_avg: Optional[nn.Module] = None,
+ optimizer: Optional[torch.optim.Optimizer] = None,
+ scheduler: Optional[LRSchedulerType] = None,
+ sampler: Optional[CutSampler] = None,
+ scaler: Optional[GradScaler] = 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.
+ model_avg:
+ The stored model averaged from the start of training.
+ optimizer:
+ The optimizer used in the training.
+ sampler:
+ The sampler for the training dataset.
+ scaler:
+ The scaler used for mix precision training.
+ """
+ if rank != 0:
+ return
+ filename = params.exp_dir / f"epoch-{params.cur_epoch}.pt"
+ save_checkpoint_impl(
+ filename=filename,
+ model=model,
+ model_avg=model_avg,
+ params=params,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=sampler,
+ scaler=scaler,
+ 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: Union[nn.Module, DDP],
+ graph_compiler: CharCtcTrainingGraphCompiler,
+ batch: dict,
+ is_training: bool,
+) -> Tuple[Tensor, MetricsTracker]:
+ """
+ Compute transducer 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 Zipformer in our case.
+ batch:
+ A batch of data. See `lhotse.dataset.K2SpeechRecognitionDataset()`
+ for the content in it.
+ 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.
+ warmup: a floating point value which increases throughout training;
+ values >= 1.0 are fully warmed up and have all modules present.
+ """
+ device = model.device if isinstance(model, DDP) else next(model.parameters()).device
+ feature = batch["inputs"]
+ # at entry, feature is (N, T, C)
+ assert feature.ndim == 3
+ feature = feature.to(device)
+
+ supervisions = batch["supervisions"]
+ feature_lens = supervisions["num_frames"].to(device)
+
+ batch_idx_train = params.batch_idx_train
+ warm_step = params.warm_step
+
+ texts = batch["supervisions"]["text"]
+
+ y = graph_compiler.texts_to_ids(texts)
+ if type(y) == list:
+ y = k2.RaggedTensor(y).to(device)
+ else:
+ y = y.to(device)
+
+ with torch.set_grad_enabled(is_training):
+ simple_loss, pruned_loss = model(
+ x=feature,
+ x_lens=feature_lens,
+ y=y,
+ prune_range=params.prune_range,
+ am_scale=params.am_scale,
+ lm_scale=params.lm_scale,
+ )
+
+ s = params.simple_loss_scale
+ # take down the scale on the simple loss from 1.0 at the start
+ # to params.simple_loss scale by warm_step.
+ simple_loss_scale = (
+ s
+ if batch_idx_train >= warm_step
+ else 1.0 - (batch_idx_train / warm_step) * (1.0 - s)
+ )
+ pruned_loss_scale = (
+ 1.0
+ if batch_idx_train >= warm_step
+ else 0.1 + 0.9 * (batch_idx_train / warm_step)
+ )
+
+ loss = simple_loss_scale * simple_loss + pruned_loss_scale * pruned_loss
+
+ assert loss.requires_grad == is_training
+
+ info = MetricsTracker()
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ info["frames"] = ((feature_lens - 7) // 2).sum().item()
+
+ # Note: We use reduction=sum while computing the loss.
+ info["loss"] = loss.detach().cpu().item()
+ info["simple_loss"] = simple_loss.detach().cpu().item()
+ info["pruned_loss"] = pruned_loss.detach().cpu().item()
+
+ return loss, info
+
+
+def compute_validation_loss(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ graph_compiler: CharCtcTrainingGraphCompiler,
+ valid_dl: torch.utils.data.DataLoader,
+ world_size: int = 1,
+) -> MetricsTracker:
+ """Run the validation process."""
+ model.eval()
+
+ tot_loss = MetricsTracker()
+
+ for batch_idx, batch in enumerate(valid_dl):
+ loss, loss_info = compute_loss(
+ params=params,
+ model=model,
+ graph_compiler=graph_compiler,
+ batch=batch,
+ is_training=False,
+ )
+ assert loss.requires_grad is False
+ tot_loss = tot_loss + loss_info
+
+ if world_size > 1:
+ tot_loss.reduce(loss.device)
+
+ loss_value = tot_loss["loss"] / tot_loss["frames"]
+ if loss_value < params.best_valid_loss:
+ params.best_valid_epoch = params.cur_epoch
+ params.best_valid_loss = loss_value
+
+ return tot_loss
+
+
+def train_one_epoch(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ optimizer: torch.optim.Optimizer,
+ scheduler: LRSchedulerType,
+ graph_compiler: CharCtcTrainingGraphCompiler,
+ train_dl: torch.utils.data.DataLoader,
+ valid_dl: torch.utils.data.DataLoader,
+ scaler: GradScaler,
+ model_avg: Optional[nn.Module] = None,
+ tb_writer: Optional[SummaryWriter] = None,
+ world_size: int = 1,
+ rank: int = 0,
+) -> 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.
+ scheduler:
+ The learning rate scheduler, we call step() every step.
+ train_dl:
+ Dataloader for the training dataset.
+ valid_dl:
+ Dataloader for the validation dataset.
+ scaler:
+ The scaler used for mix precision training.
+ model_avg:
+ The stored model averaged from the start of training.
+ tb_writer:
+ Writer to write log messages to tensorboard.
+ world_size:
+ Number of nodes in DDP training. If it is 1, DDP is disabled.
+ rank:
+ The rank of the node in DDP training. If no DDP is used, it should
+ be set to 0.
+ """
+ model.train()
+
+ tot_loss = MetricsTracker()
+
+ cur_batch_idx = params.get("cur_batch_idx", 0)
+
+ for batch_idx, batch in enumerate(train_dl):
+ if batch_idx < cur_batch_idx:
+ continue
+ cur_batch_idx = batch_idx
+
+ params.batch_idx_train += 1
+ batch_size = len(batch["supervisions"]["text"])
+
+ try:
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, loss_info = compute_loss(
+ params=params,
+ model=model,
+ graph_compiler=graph_compiler,
+ batch=batch,
+ is_training=True,
+ )
+ # summary stats
+ tot_loss = (tot_loss * (1 - 1 / params.reset_interval)) + loss_info
+
+ # NOTE: We use reduction==sum and loss is computed over utterances
+ # in the batch and there is no normalization to it so far.
+ scaler.scale(loss).backward()
+ set_batch_count(model, params.batch_idx_train)
+ scheduler.step_batch(params.batch_idx_train)
+
+ scaler.step(optimizer)
+ scaler.update()
+ optimizer.zero_grad()
+ except: # noqa
+ display_and_save_batch(batch, params=params, graph_compiler=graph_compiler)
+ raise
+
+ if params.print_diagnostics and batch_idx == 5:
+ return
+
+ if (
+ rank == 0
+ and params.batch_idx_train > 0
+ and params.batch_idx_train % params.average_period == 0
+ ):
+ update_averaged_model(
+ params=params,
+ model_cur=model,
+ model_avg=model_avg,
+ )
+
+ if (
+ params.batch_idx_train > 0
+ and params.batch_idx_train % params.save_every_n == 0
+ ):
+ params.cur_batch_idx = batch_idx
+ save_checkpoint_with_global_batch_idx(
+ out_dir=params.exp_dir,
+ global_batch_idx=params.batch_idx_train,
+ model=model,
+ model_avg=model_avg,
+ params=params,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=train_dl.sampler,
+ scaler=scaler,
+ rank=rank,
+ )
+ del params.cur_batch_idx
+ remove_checkpoints(
+ out_dir=params.exp_dir,
+ topk=params.keep_last_k,
+ rank=rank,
+ )
+
+ if batch_idx % 100 == 0 and params.use_fp16:
+ # If the grad scale was less than 1, try increasing it. The _growth_interval
+ # of the grad scaler is configurable, but we can't configure it to have different
+ # behavior depending on the current grad scale.
+ cur_grad_scale = scaler._scale.item()
+ if cur_grad_scale < 1.0 or (cur_grad_scale < 8.0 and batch_idx % 400 == 0):
+ scaler.update(cur_grad_scale * 2.0)
+ if cur_grad_scale < 0.01:
+ logging.warning(f"Grad scale is small: {cur_grad_scale}")
+ if cur_grad_scale < 1.0e-05:
+ raise RuntimeError(
+ f"grad_scale is too small, exiting: {cur_grad_scale}"
+ )
+
+ if batch_idx % params.log_interval == 0:
+ cur_lr = scheduler.get_last_lr()[0]
+ cur_grad_scale = scaler._scale.item() if params.use_fp16 else 1.0
+
+ logging.info(
+ f"Epoch {params.cur_epoch}, "
+ f"batch {batch_idx}, loss[{loss_info}], "
+ f"tot_loss[{tot_loss}], batch size: {batch_size}, "
+ f"lr: {cur_lr:.2e}, "
+ + (f"grad_scale: {scaler._scale.item()}" if params.use_fp16 else "")
+ )
+
+ if tb_writer is not None:
+ tb_writer.add_scalar(
+ "train/learning_rate", cur_lr, params.batch_idx_train
+ )
+
+ loss_info.write_summary(
+ tb_writer, "train/current_", params.batch_idx_train
+ )
+ tot_loss.write_summary(tb_writer, "train/tot_", params.batch_idx_train)
+ if params.use_fp16:
+ tb_writer.add_scalar(
+ "train/grad_scale", cur_grad_scale, params.batch_idx_train
+ )
+
+ if batch_idx % params.valid_interval == 0 and not params.print_diagnostics:
+ logging.info("Computing validation loss")
+ valid_info = 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}, validation: {valid_info}")
+ logging.info(
+ f"Maximum memory allocated so far is {torch.cuda.max_memory_allocated()//1000000}MB"
+ )
+ if tb_writer is not None:
+ valid_info.write_summary(
+ tb_writer, "train/valid_", params.batch_idx_train
+ )
+
+ loss_value = tot_loss["loss"] / tot_loss["frames"]
+ params.train_loss = loss_value
+ 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(params.seed)
+ 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")
+
+ if args.tensorboard and rank == 0:
+ tb_writer = SummaryWriter(log_dir=f"{params.exp_dir}/tensorboard")
+ else:
+ tb_writer = None
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", rank)
+ logging.info(f"Device: {device}")
+
+ lexicon = Lexicon(params.lang_dir)
+ graph_compiler = CharCtcTrainingGraphCompiler(
+ lexicon=lexicon,
+ device=device,
+ )
+
+ params.blank_id = lexicon.token_table[""]
+ params.vocab_size = max(lexicon.tokens) + 1
+
+ logging.info(params)
+
+ logging.info("About to create model")
+ model = get_transducer_model(params)
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ assert params.save_every_n >= params.average_period
+ model_avg: Optional[nn.Module] = None
+ if rank == 0:
+ # model_avg is only used with rank 0
+ model_avg = copy.deepcopy(model).to(torch.float64)
+
+ assert params.start_epoch > 0, params.start_epoch
+ checkpoints = load_checkpoint_if_available(
+ params=params, model=model, model_avg=model_avg
+ )
+
+ model.to(device)
+ if world_size > 1:
+ logging.info("Using DDP")
+ model = DDP(model, device_ids=[rank], find_unused_parameters=True)
+
+ parameters_names = []
+ parameters_names.append(
+ [name_param_pair[0] for name_param_pair in model.named_parameters()]
+ )
+ optimizer = ScaledAdam(
+ model.parameters(),
+ lr=params.base_lr,
+ clipping_scale=2.0,
+ parameters_names=parameters_names,
+ )
+
+ scheduler = Eden(optimizer, params.lr_batches, params.lr_epochs)
+
+ if checkpoints and "optimizer" in checkpoints:
+ logging.info("Loading optimizer state dict")
+ optimizer.load_state_dict(checkpoints["optimizer"])
+
+ if (
+ checkpoints
+ and "scheduler" in checkpoints
+ and checkpoints["scheduler"] is not None
+ ):
+ logging.info("Loading scheduler state dict")
+ scheduler.load_state_dict(checkpoints["scheduler"])
+
+ if params.print_diagnostics:
+ opts = diagnostics.TensorDiagnosticOptions(
+ 2**22
+ ) # allow 4 megabytes per sub-module
+ diagnostic = diagnostics.attach_diagnostics(model, opts)
+
+ if params.start_batch > 0 and checkpoints and "sampler" in checkpoints:
+ # We only load the sampler's state dict when it loads a checkpoint
+ # saved in the middle of an epoch
+ sampler_state_dict = checkpoints["sampler"]
+ else:
+ sampler_state_dict = None
+
+ if params.inf_check:
+ register_inf_check_hooks(model)
+
+ alimeeting = AlimeetingAsrDataModule(args)
+
+ train_cuts = alimeeting.train_cuts()
+ train_dl = alimeeting.train_dataloaders(
+ train_cuts, sampler_state_dict=sampler_state_dict
+ )
+
+ valid_cuts = alimeeting.eval_ihm_cuts()
+ valid_dl = alimeeting.valid_dataloaders(valid_cuts)
+
+ # if not params.print_diagnostics:
+ # scan_pessimistic_batches_for_oom(
+ # model=model,
+ # train_dl=train_dl,
+ # optimizer=optimizer,
+ # graph_compiler=graph_compiler,
+ # params=params,
+ # )
+
+ scaler = GradScaler(enabled=params.use_fp16, init_scale=1.0)
+ if checkpoints and "grad_scaler" in checkpoints:
+ logging.info("Loading grad scaler state dict")
+ scaler.load_state_dict(checkpoints["grad_scaler"])
+
+ for epoch in range(params.start_epoch, params.num_epochs + 1):
+ scheduler.step_epoch(epoch - 1)
+ fix_random_seed(params.seed + epoch - 1)
+ train_dl.sampler.set_epoch(epoch - 1)
+
+ if tb_writer is not None:
+ tb_writer.add_scalar("train/epoch", epoch, params.batch_idx_train)
+
+ params.cur_epoch = epoch
+
+ train_one_epoch(
+ params=params,
+ model=model,
+ model_avg=model_avg,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ graph_compiler=graph_compiler,
+ train_dl=train_dl,
+ valid_dl=valid_dl,
+ scaler=scaler,
+ tb_writer=tb_writer,
+ world_size=world_size,
+ rank=rank,
+ )
+
+ if params.print_diagnostics:
+ diagnostic.print_diagnostics()
+ break
+
+ save_checkpoint(
+ params=params,
+ model=model,
+ model_avg=model_avg,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=train_dl.sampler,
+ scaler=scaler,
+ rank=rank,
+ )
+
+ logging.info("Done!")
+
+ if world_size > 1:
+ torch.distributed.barrier()
+ cleanup_dist()
+
+
+def display_and_save_batch(
+ batch: dict,
+ params: AttributeDict,
+ graph_compiler: CharCtcTrainingGraphCompiler,
+) -> None:
+ """Display the batch statistics and save the batch into disk.
+
+ Args:
+ batch:
+ A batch of data. See `lhotse.dataset.K2SpeechRecognitionDataset()`
+ for the content in it.
+ params:
+ Parameters for training. See :func:`get_params`.
+ sp:
+ The BPE model.
+ """
+ from lhotse.utils import uuid4
+
+ filename = f"{params.exp_dir}/batch-{uuid4()}.pt"
+ logging.info(f"Saving batch to {filename}")
+ torch.save(batch, filename)
+
+ supervisions = batch["supervisions"]
+ features = batch["inputs"]
+
+ logging.info(f"features shape: {features.shape}")
+
+
+def scan_pessimistic_batches_for_oom(
+ model: Union[nn.Module, DDP],
+ train_dl: torch.utils.data.DataLoader,
+ optimizer: torch.optim.Optimizer,
+ graph_compiler: CharCtcTrainingGraphCompiler,
+ params: AttributeDict,
+):
+ from lhotse.dataset import find_pessimistic_batches
+
+ logging.info(
+ "Sanity check -- see if any of the batches in epoch 1 would cause OOM."
+ )
+ batches, crit_values = find_pessimistic_batches(train_dl.sampler)
+ for criterion, cuts in batches.items():
+ batch = train_dl.dataset[cuts]
+ try:
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, _ = compute_loss(
+ params=params,
+ model=model,
+ graph_compiler=graph_compiler,
+ batch=batch,
+ is_training=True,
+ )
+ loss.backward()
+ optimizer.zero_grad()
+ except Exception as e:
+ if "CUDA out of memory" in str(e):
+ logging.error(
+ "Your GPU ran out of memory with the current "
+ "max_duration setting. We recommend decreasing "
+ "max_duration and trying again.\n"
+ f"Failing criterion: {criterion} "
+ f"(={crit_values[criterion]}) ..."
+ )
+ display_and_save_batch(batch, params=params, graph_compiler=graph_compiler)
+ raise
+ logging.info(
+ f"Maximum memory allocated so far is {torch.cuda.max_memory_allocated()//1000000}MB"
+ )
+
+
+def main():
+ parser = get_parser()
+ AlimeetingAsrDataModule.add_arguments(parser)
+ args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ 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)
+
+
+torch.set_num_threads(1)
+torch.set_num_interop_threads(1)
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/zipformer.py b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/zipformer.py
new file mode 120000
index 000000000..f2f66041e
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/zipformer.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7/zipformer.py
\ No newline at end of file
diff --git a/egs/alimeeting/ASR_v2/shared b/egs/alimeeting/ASR_v2/shared
new file mode 120000
index 000000000..3a3b28f96
--- /dev/null
+++ b/egs/alimeeting/ASR_v2/shared
@@ -0,0 +1 @@
+../../../egs/aishell/ASR/shared
\ No newline at end of file
diff --git a/egs/gigaspeech/ASR/.gitignore b/egs/gigaspeech/ASR/.gitignore
index 5592679cc..8dec2d86d 100644
--- a/egs/gigaspeech/ASR/.gitignore
+++ b/egs/gigaspeech/ASR/.gitignore
@@ -1 +1,2 @@
log-*
+.DS_Store
\ No newline at end of file
diff --git a/egs/librispeech/ASR/.gitignore b/egs/librispeech/ASR/.gitignore
index 5592679cc..8dec2d86d 100644
--- a/egs/librispeech/ASR/.gitignore
+++ b/egs/librispeech/ASR/.gitignore
@@ -1 +1,2 @@
log-*
+.DS_Store
\ No newline at end of file
diff --git a/egs/librispeech/ASR/README.md b/egs/librispeech/ASR/README.md
index caa23a49f..94cb445a8 100644
--- a/egs/librispeech/ASR/README.md
+++ b/egs/librispeech/ASR/README.md
@@ -19,18 +19,36 @@ The following table lists the differences among them.
| `pruned_transducer_stateless` | Conformer | Embedding + Conv1d | Using k2 pruned RNN-T loss |
| `pruned_transducer_stateless2` | Conformer(modified) | Embedding + Conv1d | Using k2 pruned RNN-T loss |
| `pruned_transducer_stateless3` | Conformer(modified) | Embedding + Conv1d | Using k2 pruned RNN-T loss + using GigaSpeech as extra training data |
-| `pruned_transducer_stateless4` | Conformer(modified) | Embedding + Conv1d | same as pruned_transducer_stateless2 + save averaged models periodically during training |
+| `pruned_transducer_stateless4` | Conformer(modified) | Embedding + Conv1d | same as pruned_transducer_stateless2 + save averaged models periodically during training + delay penalty |
| `pruned_transducer_stateless5` | Conformer(modified) | Embedding + Conv1d | same as pruned_transducer_stateless4 + more layers + random combiner|
| `pruned_transducer_stateless6` | Conformer(modified) | Embedding + Conv1d | same as pruned_transducer_stateless4 + distillation with hubert|
| `pruned_transducer_stateless7` | Zipformer | Embedding + Conv1d | First experiment with Zipformer from Dan|
| `pruned_transducer_stateless7_ctc` | Zipformer | Embedding + Conv1d | Same as pruned_transducer_stateless7, but with extra CTC head|
+| `pruned_transducer_stateless7_ctc_bs` | Zipformer | Embedding + Conv1d | pruned_transducer_stateless7_ctc + blank skip |
+| `pruned_transducer_stateless7_streaming` | Streaming Zipformer | Embedding + Conv1d | streaming version of pruned_transducer_stateless7 |
| `pruned_transducer_stateless8` | Zipformer | Embedding + Conv1d | Same as pruned_transducer_stateless7, but using extra data from GigaSpeech|
| `pruned_stateless_emformer_rnnt2` | Emformer(from torchaudio) | Embedding + Conv1d | Using Emformer from torchaudio for streaming ASR|
| `conv_emformer_transducer_stateless` | ConvEmformer | Embedding + Conv1d | Using ConvEmformer for streaming ASR + mechanisms in reworked model |
| `conv_emformer_transducer_stateless2` | ConvEmformer | Embedding + Conv1d | Using ConvEmformer with simplified memory for streaming ASR + mechanisms in reworked model |
| `lstm_transducer_stateless` | LSTM | Embedding + Conv1d | Using LSTM with mechanisms in reworked model |
-| `lstm_transducer_stateless2` | LSTM | Embedding + Conv1d | Using LSTM with mechanisms in reworked model + gigaspeech (multi-dataset setup) |
+| `lstm_transducer_stateless2` | LSTM | Embedding + Conv1d | Using LSTM with mechanisms in reworked model + gigaspeech (multi-dataset setup) |
+| `lstm_transducer_stateless3` | LSTM | Embedding + Conv1d | Using LSTM with mechanisms in reworked model + gradient filter + delay penalty |
The decoder in `transducer_stateless` is modified from the paper
[Rnn-Transducer with Stateless Prediction Network](https://ieeexplore.ieee.org/document/9054419/).
We place an additional Conv1d layer right after the input embedding layer.
+
+# CTC
+
+| | Encoder | Comment |
+|------------------------------|--------------------|------------------------------|
+| `conformer-ctc` | Conformer | Use auxiliary attention head |
+| `conformer-ctc2` | Reworked Conformer | Use auxiliary attention head |
+| `conformer-ctc3` | Reworked Conformer | Streaming version + delay penalty |
+
+# MMI
+
+| | Encoder | Comment |
+|------------------------------|-----------|---------------------------------------------------|
+| `conformer-mmi` | Conformer | |
+| `zipformer-mmi` | Zipformer | CTC warmup + use HP as decoding graph for decoding |
diff --git a/egs/librispeech/ASR/RESULTS.md b/egs/librispeech/ASR/RESULTS.md
index 9e5669f6d..b30cf7c1f 100644
--- a/egs/librispeech/ASR/RESULTS.md
+++ b/egs/librispeech/ASR/RESULTS.md
@@ -1,5 +1,140 @@
## Results
+### Streaming Zipformer-Transducer (Pruned Stateless Transducer + Streaming Zipformer)
+
+#### [pruned_transducer_stateless7_streaming](./pruned_transducer_stateless7_streaming)
+
+See for more details.
+
+You can find a pretrained model, training logs, decoding logs, and decoding
+results at:
+
+
+Number of model parameters: 70369391, i.e., 70.37 M
+
+##### training on full librispeech
+
+The WERs are:
+
+| decoding method | chunk size | test-clean | test-other | comment | decoding mode |
+|----------------------|------------|------------|------------|---------------------|----------------------|
+| greedy search | 320ms | 3.15 | 8.09 | --epoch 30 --avg 9 | simulated streaming |
+| greedy search | 320ms | 3.17 | 8.24 | --epoch 30 --avg 9 | chunk-wise |
+| fast beam search | 320ms | 3.2 | 8.04 | --epoch 30 --avg 9 | simulated streaming |
+| fast beam search | 320ms | 3.36 | 8.19 | --epoch 30 --avg 9 | chunk-wise |
+| modified beam search | 320ms | 3.11 | 7.93 | --epoch 30 --avg 9 | simulated streaming |
+| modified beam search | 320ms | 3.12 | 8.11 | --epoch 30 --avg 9 | chunk-size |
+| greedy search | 640ms | 2.97 | 7.5 | --epoch 30 --avg 9 | simulated streaming |
+| greedy search | 640ms | 2.98 | 7.67 | --epoch 30 --avg 9 | chunk-wise |
+| fast beam search | 640ms | 3.02 | 7.47 | --epoch 30 --avg 9 | simulated streaming |
+| fast beam search | 640ms | 2.96 | 7.61 | --epoch 30 --avg 9 | chunk-wise |
+| modified beam search | 640ms | 2.94 | 7.36 | --epoch 30 --avg 9 | simulated streaming |
+| modified beam search | 640ms | 2.95 | 7.53 | --epoch 30 --avg 9 | chunk-size |
+
+Note: `simulated streaming` indicates feeding full utterance during decoding using `decode.py`,
+while `chunk-size` indicates feeding certain number of frames at each time using `streaming_decode.py`.
+
+The training command is:
+
+```bash
+./pruned_transducer_stateless7_streaming/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --use-fp16 1 \
+ --exp-dir pruned_transducer_stateless7_streaming/exp \
+ --full-libri 1 \
+ --max-duration 750 \
+ --master-port 12345
+```
+
+The tensorboard log can be found at
+
+
+The simulated streaming decoding command (e.g., chunk-size=320ms) is:
+```bash
+for $m in greedy_search fast_beam_search modified_beam_search; do
+ ./pruned_transducer_stateless7_streaming/decode.py \
+ --epoch 30 \
+ --avg 9 \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --max-duration 600 \
+ --decode-chunk-len 32 \
+ --decoding-method $m
+done
+```
+
+The streaming chunk-size decoding command (e.g., chunk-size=320ms) is:
+```bash
+for m in greedy_search modified_beam_search fast_beam_search; do
+ ./pruned_transducer_stateless7_streaming/streaming_decode.py \
+ --epoch 30 \
+ --avg 9 \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --decoding-method $m \
+ --decode-chunk-len 32 \
+ --num-decode-streams 2000
+done
+```
+
+
+### zipformer_mmi (zipformer with mmi loss)
+
+See for more details.
+
+[zipformer_mmi](./zipformer_mmi)
+
+The tensorboard log can be found at
+
+
+You can find a pretrained model, training logs, decoding logs, and decoding
+results at:
+
+
+Number of model parameters: 69136519, i.e., 69.14 M
+
+| | test-clean | test-other | comment |
+|--------------------------|------------|-------------|---------------------|
+| 1best | 2.54 | 5.65 | --epoch 30 --avg 10 |
+| nbest | 2.54 | 5.66 | --epoch 30 --avg 10 |
+| nbest-rescoring-LG | 2.49 | 5.42 | --epoch 30 --avg 10 |
+| nbest-rescoring-3-gram | 2.52 | 5.62 | --epoch 30 --avg 10 |
+| nbest-rescoring-4-gram | 2.5 | 5.51 | --epoch 30 --avg 10 |
+
+The training commands are:
+```bash
+export CUDA_VISIBLE_DEVICES="0,1,2,3"
+
+./zipformer_mmi/train.py \
+ --world-size 4 \
+ --master-port 12345 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --lang-dir data/lang_bpe_500 \
+ --max-duration 500 \
+ --full-libri 1 \
+ --use-fp16 1 \
+ --exp-dir zipformer_mmi/exp
+```
+
+The decoding commands for the transducer branch are:
+```bash
+export CUDA_VISIBLE_DEVICES="5"
+
+for m in nbest nbest-rescoring-LG nbest-rescoring-3-gram nbest-rescoring-4-gram; do
+ ./zipformer_mmi/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir ./zipformer_mmi/exp/ \
+ --max-duration 100 \
+ --lang-dir data/lang_bpe_500 \
+ --nbest-scale 1.2 \
+ --hp-scale 1.0 \
+ --decoding-method $m
+done
+```
+
+
### pruned_transducer_stateless7_ctc (zipformer with transducer loss and ctc loss)
See for more details.
@@ -261,9 +396,13 @@ Number of model parameters: 70369391, i.e., 70.37 M
| | test-clean | test-other | comment |
|----------------------|------------|-------------|----------------------------------------|
-| greedy search | 2.17 | 5.23 | --epoch 39 --avg 6 --max-duration 600 |
-| modified beam search | 2.15 | 5.20 | --epoch 39 --avg 6 --max-duration 600 |
-| fast beam search | 2.15 | 5.22 | --epoch 39 --avg 6 --max-duration 600 |
+| greedy search | 2.17 | 5.23 | --epoch 30 --avg 9 --max-duration 600 |
+| modified beam search | 2.15 | 5.20 | --epoch 30 --avg 9 --max-duration 600 |
+| modified beam search + RNNLM shallow fusion | 1.99 | 4.73 | --epoch 30 --avg 9 --max-duration 600 |
+| modified beam search + TransformerLM shallow fusion | 1.94 | 4.73 | --epoch 30 --avg 9 --max-duration 600 |
+| modified beam search + RNNLM + LODR | 1.91 | 4.57 | --epoch 30 --avg 9 --max-duration 600 |
+| modified beam search + TransformerLM + LODR | 1.91 | 4.51 | --epoch 30 --avg 9 --max-duration 600 |
+| fast beam search | 2.15 | 5.22 | --epoch 30 --avg 9 --max-duration 600 |
The training commands are:
```bash
@@ -401,7 +540,9 @@ The WERs are:
| greedy search (max sym per frame 1) | 2.78 | 7.36 | --iter 468000 --avg 16 |
| modified_beam_search | 2.73 | 7.15 | --iter 468000 --avg 16 |
| modified_beam_search + RNNLM shallow fusion | 2.42 | 6.46 | --iter 468000 --avg 16 |
-| modified_beam_search + RNNLM shallow fusion | 2.28 | 5.94 | --iter 468000 --avg 16 |
+| modified_beam_search + TransformerLM shallow fusion | 2.37 | 6.48 | --iter 468000 --avg 16 |
+| modified_beam_search + RNNLM + LODR | 2.24 | 5.89 | --iter 468000 --avg 16 |
+| modified_beam_search + TransformerLM + LODR | 2.19 | 5.90 | --iter 468000 --avg 16 |
| fast_beam_search | 2.76 | 7.31 | --iter 468000 --avg 16 |
| greedy search (max sym per frame 1) | 2.77 | 7.35 | --iter 472000 --avg 18 |
| modified_beam_search | 2.75 | 7.08 | --iter 472000 --avg 18 |
@@ -456,9 +597,12 @@ for m in greedy_search fast_beam_search modified_beam_search; do
done
```
-To decode with RNNLM shallow fusion, use the following decoding command. A well-trained RNNLM
-can be found here:
+You may also decode using shallow fusion with external neural network LM. To do so you need to
+download a well-trained NN LM:
+RNN LM:
+Transformer LM:
+```bash
for iter in 472000; do
for avg in 8 10 12 14 16 18; do
./lstm_transducer_stateless2/decode.py \
@@ -466,23 +610,24 @@ for iter in 472000; do
--avg $avg \
--exp-dir ./lstm_transducer_stateless2/exp \
--max-duration 600 \
- --decoding-method modified_beam_search_rnnlm_shallow_fusion \
- --beam 4 \
- --rnn-lm-scale 0.3 \
- --rnn-lm-exp-dir /path/to/RNNLM \
- --rnn-lm-epoch 99 \
- --rnn-lm-avg 1 \
- --rnn-lm-num-layers 3 \
- --rnn-lm-tie-weights 1
+ --decoding-method modified_beam_search_lm_shallow_fusion \
+ --use-shallow-fusion 1 \
+ --lm-type rnn \
+ --lm-exp-dir /ceph-data4/yangxiaoyu/pretrained_models/LM/icefall-librispeech-rnn-lm/exp \
+ --lm-epoch 99 \
+ --lm-scale $lm_scale \
+ --lm-avg 1 \
done
done
+```
-You may also decode using LODR + RNNLM shallow fusion. This decoding method is proposed in .
+You may also decode using LODR + LM shallow fusion. This decoding method is proposed in .
It subtracts the internal language model score during shallow fusion, which is approximated by a bi-gram model. The bi-gram can be
generated by `generate-lm.sh`, or you may download it from .
The decoding command is as follows:
+```bash
for iter in 472000; do
for avg in 8 10 12 14 16 18; do
./lstm_transducer_stateless2/decode.py \
@@ -490,18 +635,22 @@ for iter in 472000; do
--avg $avg \
--exp-dir ./lstm_transducer_stateless2/exp \
--max-duration 600 \
- --decoding-method modified_beam_search_rnnlm_LODR \
+ --decoding-method modified_beam_search_LODR \
--beam 4 \
- --rnn-lm-scale 0.4 \
- --rnn-lm-exp-dir /path/to/RNNLM \
- --rnn-lm-epoch 99 \
- --rnn-lm-avg 1 \
- --rnn-lm-num-layers 3 \
- --rnn-lm-tie-weights 1 \
- --token-ngram 2 \
+ --max-contexts 4 \
+ --use-shallow-fusion 1 \
+ --lm-type rnn \
+ --lm-exp-dir /ceph-data4/yangxiaoyu/pretrained_models/LM/icefall-librispeech-rnn-lm/exp \
+ --lm-epoch 99 \
+ --lm-scale 0.4 \
+ --lm-avg 1 \
+ --tokens-ngram 2 \
--ngram-lm-scale -0.16
done
done
+```
+Note that you can also set `--lm-type transformer` to use transformer LM during LODR. But it will be slower
+because it has not been optimized. The pre-trained transformer LM is available at
Pretrained models, training logs, decoding logs, and decoding results
are available at
@@ -1660,6 +1809,9 @@ layers (24 v.s 12) but a narrower model (1536 feedforward dim and 384 encoder di
| greedy search (max sym per frame 1) | 2.54 | 5.72 | --epoch 30 --avg 10 --max-duration 600 |
| modified beam search | 2.47 | 5.71 | --epoch 30 --avg 10 --max-duration 600 |
| modified beam search + RNNLM shallow fusion | 2.27 | 5.24 | --epoch 30 --avg 10 --max-duration 600 |
+| modified beam search + RNNLM + LODR | 2.23 | 5.17 | --epoch 30 --avg 10 --max-duration 600 |
+| modified beam search + TransformerLM shallow fusion | 2.27 | 5.26 | --epoch 30 --avg 10 --max-duration 600 |
+| modified beam search + TransformerLM + LODR | 2.22 | 5.11 | --epoch 30 --avg 10 --max-duration 600 |
| fast beam search | 2.5 | 5.72 | --epoch 30 --avg 10 --max-duration 600 |
```bash
@@ -2023,7 +2175,8 @@ subset so that the gigaspeech dataloader never exhausts.
| greedy search (max sym per frame 1) | 2.03 | 4.70 | --iter 1224000 --avg 14 --max-duration 600 |
| modified beam search | 2.00 | 4.63 | --iter 1224000 --avg 14 --max-duration 600 |
| modified beam search + rnnlm shallow fusion | 1.94 | 4.2 | --iter 1224000 --avg 14 --max-duration 600 |
-| modified beam search + LODR | 1.83 | 4.03 | --iter 1224000 --avg 14 --max-duration 600 |
+| modified beam search + rnnlm + LODR | 1.77 | 3.99 | --iter 1224000 --avg 14 --max-duration 600 |
+| modified beam search + TransformerLM + LODR | 1.75 | 3.94 | --iter 1224000 --avg 14 --max-duration 600 |
| fast beam search | 2.10 | 4.68 | --iter 1224000 --avg 14 --max-duration 600 |
The training commands are:
@@ -2069,8 +2222,10 @@ for iter in 1224000; do
done
done
```
-You may also decode using shallow fusion with external RNNLM. To do so you need to
-download a well-trained RNNLM from this link
+You may also decode using shallow fusion with external neural network LM. To do so you need to
+download a well-trained NN LM:
+RNN LM:
+Transformer LM:
```bash
rnn_lm_scale=0.3
diff --git a/egs/librispeech/ASR/conformer_ctc/label_smoothing.py b/egs/librispeech/ASR/conformer_ctc/label_smoothing.py
index cb0d6e04d..52d2eda3b 100644
--- a/egs/librispeech/ASR/conformer_ctc/label_smoothing.py
+++ b/egs/librispeech/ASR/conformer_ctc/label_smoothing.py
@@ -44,7 +44,8 @@ class LabelSmoothingLoss(torch.nn.Module):
mean of the output is taken. (3) "sum": the output will be summed.
"""
super().__init__()
- assert 0.0 <= label_smoothing < 1.0
+ assert 0.0 <= label_smoothing < 1.0, f"{label_smoothing}"
+ assert reduction in ("none", "sum", "mean"), reduction
self.ignore_index = ignore_index
self.label_smoothing = label_smoothing
self.reduction = reduction
diff --git a/egs/librispeech/ASR/conformer_ctc2/subsampling.py b/egs/librispeech/ASR/conformer_ctc2/subsampling.py
index 3fcb4196f..85a4dc8df 100644
--- a/egs/librispeech/ASR/conformer_ctc2/subsampling.py
+++ b/egs/librispeech/ASR/conformer_ctc2/subsampling.py
@@ -24,10 +24,9 @@ from scaling import (
ScaledConv2d,
ScaledLinear,
)
-from torch import nn
-class Conv2dSubsampling(nn.Module):
+class Conv2dSubsampling(torch.nn.Module):
"""Convolutional 2D subsampling (to 1/4 length).
Convert an input of shape (N, T, idim) to an output
@@ -61,7 +60,7 @@ class Conv2dSubsampling(nn.Module):
assert in_channels >= 7
super().__init__()
- self.conv = nn.Sequential(
+ self.conv = torch.nn.Sequential(
ScaledConv2d(
in_channels=1,
out_channels=layer1_channels,
diff --git a/egs/librispeech/ASR/conformer_ctc3/jit_pretrained.py b/egs/librispeech/ASR/conformer_ctc3/jit_pretrained.py
index 5be898e37..76db46cc8 100755
--- a/egs/librispeech/ASR/conformer_ctc3/jit_pretrained.py
+++ b/egs/librispeech/ASR/conformer_ctc3/jit_pretrained.py
@@ -291,7 +291,10 @@ def main():
batch_size = nnet_output.shape[0]
supervision_segments = torch.tensor(
- [[i, 0, nnet_output.shape[1]] for i in range(batch_size)],
+ [
+ [i, 0, feature_lengths[i] // params.subsampling_factor]
+ for i in range(batch_size)
+ ],
dtype=torch.int32,
)
diff --git a/egs/librispeech/ASR/conformer_ctc3/pretrained.py b/egs/librispeech/ASR/conformer_ctc3/pretrained.py
index 3628d6a5f..880945ea0 100755
--- a/egs/librispeech/ASR/conformer_ctc3/pretrained.py
+++ b/egs/librispeech/ASR/conformer_ctc3/pretrained.py
@@ -339,7 +339,10 @@ def main():
batch_size = nnet_output.shape[0]
supervision_segments = torch.tensor(
- [[i, 0, nnet_output.shape[1]] for i in range(batch_size)],
+ [
+ [i, 0, feature_lengths[i] // params.subsampling_factor]
+ for i in range(batch_size)
+ ],
dtype=torch.int32,
)
diff --git a/egs/librispeech/ASR/conformer_mmi/decode.py b/egs/librispeech/ASR/conformer_mmi/decode.py
index e3c7b685f..74f6e73fa 100755
--- a/egs/librispeech/ASR/conformer_mmi/decode.py
+++ b/egs/librispeech/ASR/conformer_mmi/decode.py
@@ -660,14 +660,22 @@ def main():
# we need cut ids to display recognition results.
args.return_cuts = True
librispeech = LibriSpeechAsrDataModule(args)
+
+ test_clean_cuts = librispeech.test_clean_cuts()
+ test_other_cuts = librispeech.test_other_cuts()
+
+ test_clean_dl = librispeech.test_dataloaders(test_clean_cuts)
+ test_other_dl = librispeech.test_dataloaders(test_other_cuts)
+
# 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()):
+ test_dls = [test_clean_dl, test_other_dl]
+
+ for test_set, test_dl in zip(test_sets, test_dls):
results_dict = decode_dataset(
dl=test_dl,
params=params,
diff --git a/egs/librispeech/ASR/conformer_mmi/train-with-attention.py b/egs/librispeech/ASR/conformer_mmi/train-with-attention.py
index f8c94cff9..100bc846a 100755
--- a/egs/librispeech/ASR/conformer_mmi/train-with-attention.py
+++ b/egs/librispeech/ASR/conformer_mmi/train-with-attention.py
@@ -30,6 +30,8 @@ import torch.multiprocessing as mp
import torch.nn as nn
from asr_datamodule import LibriSpeechAsrDataModule
from conformer import Conformer
+from lhotse.cut import Cut
+from lhotse.dataset.sampling.base import CutSampler
from lhotse.utils import fix_random_seed
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.nn.utils import clip_grad_norm_
@@ -100,6 +102,41 @@ def get_parser():
""",
)
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="conformer_mmi/exp-attn",
+ help="""The experiment dir.
+ It specifies the directory where all training related
+ files, e.g., checkpoints, log, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=str,
+ default="data/lang_bpe_500",
+ help="""The lang dir
+ It contains language related input files such as
+ "lexicon.txt"
+ """,
+ )
+
+ parser.add_argument(
+ "--seed",
+ type=int,
+ default=42,
+ help="The seed for random generators intended for reproducibility",
+ )
+
+ parser.add_argument(
+ "--use-pruned-intersect",
+ type=str2bool,
+ default=False,
+ help="""Whether to use `intersect_dense_pruned` to get denominator
+ lattice.""",
+ )
+
return parser
@@ -114,12 +151,6 @@ def get_params() -> AttributeDict:
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"
-
- 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.
@@ -164,8 +195,6 @@ def get_params() -> AttributeDict:
"""
params = AttributeDict(
{
- "exp_dir": Path("conformer_mmi/exp_500_with_attention"),
- "lang_dir": Path("data/lang_bpe_500"),
"best_train_loss": float("inf"),
"best_valid_loss": float("inf"),
"best_train_epoch": -1,
@@ -184,15 +213,12 @@ def get_params() -> AttributeDict:
"beam_size": 6, # will change it to 8 after some batches (see code)
"reduction": "sum",
"use_double_scores": True,
- # "att_rate": 0.0,
- # "num_decoder_layers": 0,
"att_rate": 0.7,
"num_decoder_layers": 6,
# parameters for Noam
"weight_decay": 1e-6,
"lr_factor": 5.0,
"warm_step": 80000,
- "use_pruned_intersect": False,
"den_scale": 1.0,
# use alignments before this number of batches
"use_ali_until": 13000,
@@ -661,7 +687,7 @@ def run(rank, world_size, args):
params = get_params()
params.update(vars(args))
- fix_random_seed(42)
+ fix_random_seed(params.seed)
if world_size > 1:
setup_dist(rank, world_size, params.master_port)
@@ -745,8 +771,29 @@ def run(rank, world_size, args):
valid_ali = None
librispeech = LibriSpeechAsrDataModule(args)
- train_dl = librispeech.train_dataloaders()
- valid_dl = librispeech.valid_dataloaders()
+ train_cuts = librispeech.train_clean_100_cuts()
+ if params.full_libri:
+ train_cuts += librispeech.train_clean_360_cuts()
+ train_cuts += librispeech.train_other_500_cuts()
+
+ def remove_short_and_long_utt(c: Cut):
+ # Keep only utterances with duration between 1 second and 20 seconds
+ #
+ # Caution: There is a reason to select 20.0 here. Please see
+ # ../local/display_manifest_statistics.py
+ #
+ # You should use ../local/display_manifest_statistics.py to get
+ # an utterance duration distribution for your dataset to select
+ # the threshold
+ return 1.0 <= c.duration <= 20.0
+
+ train_cuts = train_cuts.filter(remove_short_and_long_utt)
+
+ train_dl = librispeech.train_dataloaders(train_cuts)
+
+ valid_cuts = librispeech.dev_clean_cuts()
+ valid_cuts += librispeech.dev_other_cuts()
+ valid_dl = librispeech.valid_dataloaders(valid_cuts)
for epoch in range(params.start_epoch, params.num_epochs):
train_dl.sampler.set_epoch(epoch)
@@ -796,6 +843,7 @@ def main():
parser = get_parser()
LibriSpeechAsrDataModule.add_arguments(parser)
args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
world_size = args.world_size
assert world_size >= 1
diff --git a/egs/librispeech/ASR/conformer_mmi/train.py b/egs/librispeech/ASR/conformer_mmi/train.py
index 5cfb2bfc7..f9f80632e 100755
--- a/egs/librispeech/ASR/conformer_mmi/train.py
+++ b/egs/librispeech/ASR/conformer_mmi/train.py
@@ -30,6 +30,8 @@ import torch.multiprocessing as mp
import torch.nn as nn
from asr_datamodule import LibriSpeechAsrDataModule
from conformer import Conformer
+from lhotse.cut import Cut
+from lhotse.dataset.sampling.base import CutSampler
from lhotse.utils import fix_random_seed
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.nn.utils import clip_grad_norm_
@@ -100,6 +102,26 @@ def get_parser():
""",
)
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="conformer_mmi/exp",
+ help="""The experiment dir.
+ It specifies the directory where all training related
+ files, e.g., checkpoints, log, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=str,
+ default="data/lang_bpe_500",
+ help="""The lang dir
+ It contains language related input files such as
+ "lexicon.txt"
+ """,
+ )
+
parser.add_argument(
"--seed",
type=int,
@@ -107,6 +129,14 @@ def get_parser():
help="The seed for random generators intended for reproducibility",
)
+ parser.add_argument(
+ "--use-pruned-intersect",
+ type=str2bool,
+ default=False,
+ help="""Whether to use `intersect_dense_pruned` to get denominator
+ lattice.""",
+ )
+
return parser
@@ -121,12 +151,6 @@ def get_params() -> AttributeDict:
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"
-
- 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.
@@ -171,8 +195,6 @@ def get_params() -> AttributeDict:
"""
params = AttributeDict(
{
- "exp_dir": Path("conformer_mmi/exp_500"),
- "lang_dir": Path("data/lang_bpe_500"),
"best_train_loss": float("inf"),
"best_valid_loss": float("inf"),
"best_train_epoch": -1,
@@ -193,13 +215,10 @@ def get_params() -> AttributeDict:
"use_double_scores": True,
"att_rate": 0.0,
"num_decoder_layers": 0,
- # "att_rate": 0.7,
- # "num_decoder_layers": 6,
# parameters for Noam
"weight_decay": 1e-6,
"lr_factor": 5.0,
"warm_step": 80000,
- "use_pruned_intersect": False,
"den_scale": 1.0,
# use alignments before this number of batches
"use_ali_until": 13000,
@@ -752,8 +771,29 @@ def run(rank, world_size, args):
valid_ali = None
librispeech = LibriSpeechAsrDataModule(args)
- train_dl = librispeech.train_dataloaders()
- valid_dl = librispeech.valid_dataloaders()
+ train_cuts = librispeech.train_clean_100_cuts()
+ if params.full_libri:
+ train_cuts += librispeech.train_clean_360_cuts()
+ train_cuts += librispeech.train_other_500_cuts()
+
+ def remove_short_and_long_utt(c: Cut):
+ # Keep only utterances with duration between 1 second and 20 seconds
+ #
+ # Caution: There is a reason to select 20.0 here. Please see
+ # ../local/display_manifest_statistics.py
+ #
+ # You should use ../local/display_manifest_statistics.py to get
+ # an utterance duration distribution for your dataset to select
+ # the threshold
+ return 1.0 <= c.duration <= 20.0
+
+ train_cuts = train_cuts.filter(remove_short_and_long_utt)
+
+ train_dl = librispeech.train_dataloaders(train_cuts)
+
+ valid_cuts = librispeech.dev_clean_cuts()
+ valid_cuts += librispeech.dev_other_cuts()
+ valid_dl = librispeech.valid_dataloaders(valid_cuts)
for epoch in range(params.start_epoch, params.num_epochs):
fix_random_seed(params.seed + epoch)
@@ -804,6 +844,7 @@ def main():
parser = get_parser()
LibriSpeechAsrDataModule.add_arguments(parser)
args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
world_size = args.world_size
assert world_size >= 1
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/emformer2.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/emformer2.py
index 65a7efa77..f0c92a9b4 100644
--- a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/emformer2.py
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/emformer2.py
@@ -1435,7 +1435,7 @@ class EmformerEncoder(nn.Module):
self,
x: torch.Tensor,
states: List[torch.Tensor],
- ) -> Tuple[torch.Tensor, List[torch.Tensor],]:
+ ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
"""Forward pass for streaming inference.
B: batch size;
@@ -1512,24 +1512,6 @@ class EmformerEncoder(nn.Module):
)
return states
- attn_caches = [
- [
- torch.zeros(self.memory_size, self.d_model, device=device),
- torch.zeros(self.left_context_length, self.d_model, device=device),
- torch.zeros(self.left_context_length, self.d_model, device=device),
- ]
- for _ in range(self.num_encoder_layers)
- ]
- conv_caches = [
- torch.zeros(self.d_model, self.cnn_module_kernel - 1, device=device)
- for _ in range(self.num_encoder_layers)
- ]
- states: Tuple[List[List[torch.Tensor]], List[torch.Tensor]] = (
- attn_caches,
- conv_caches,
- )
- return states
-
class Emformer(EncoderInterface):
def __init__(
@@ -1640,7 +1622,7 @@ class Emformer(EncoderInterface):
self,
x: torch.Tensor,
states: List[torch.Tensor],
- ) -> Tuple[torch.Tensor, List[torch.Tensor],]:
+ ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
"""Forward pass for streaming inference.
B: batch size;
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/export-for-ncnn.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/export-for-ncnn.py
index 716de5734..64c16141c 100755
--- a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/export-for-ncnn.py
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/export-for-ncnn.py
@@ -152,7 +152,6 @@ def export_encoder_model_jit_trace(
x = torch.zeros(1, T, 80, dtype=torch.float32)
states = encoder_model.init_states()
- states = encoder_model.init_states()
traced_model = torch.jit.trace(encoder_model, (x, states))
traced_model.save(encoder_filename)
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/streaming-ncnn-decode.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/streaming-ncnn-decode.py
index b21fe5c7e..e4104a5bb 100755
--- a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/streaming-ncnn-decode.py
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/streaming-ncnn-decode.py
@@ -131,6 +131,8 @@ class Model:
encoder_net = ncnn.Net()
encoder_net.opt.use_packing_layout = False
encoder_net.opt.use_fp16_storage = False
+ encoder_net.opt.num_threads = 4
+
encoder_param = args.encoder_param_filename
encoder_model = args.encoder_bin_filename
@@ -144,6 +146,7 @@ class Model:
decoder_model = args.decoder_bin_filename
decoder_net = ncnn.Net()
+ decoder_net.opt.num_threads = 4
decoder_net.load_param(decoder_param)
decoder_net.load_model(decoder_model)
@@ -154,6 +157,8 @@ class Model:
joiner_param = args.joiner_param_filename
joiner_model = args.joiner_bin_filename
joiner_net = ncnn.Net()
+ joiner_net.opt.num_threads = 4
+
joiner_net.load_param(joiner_param)
joiner_net.load_model(joiner_model)
@@ -176,7 +181,6 @@ class Model:
- next_states, a list of tensors containing the next states
"""
with self.encoder_net.create_extractor() as ex:
- ex.set_num_threads(4)
ex.input("in0", ncnn.Mat(x.numpy()).clone())
# layer0 in2-in5
@@ -220,7 +224,6 @@ class Model:
assert decoder_input.dtype == torch.int32
with self.decoder_net.create_extractor() as ex:
- ex.set_num_threads(4)
ex.input("in0", ncnn.Mat(decoder_input.numpy()).clone())
ret, ncnn_out0 = ex.extract("out0")
assert ret == 0, ret
@@ -229,7 +232,6 @@ class Model:
def run_joiner(self, encoder_out, decoder_out):
with self.joiner_net.create_extractor() as ex:
- ex.set_num_threads(4)
ex.input("in0", ncnn.Mat(encoder_out.numpy()).clone())
ex.input("in1", ncnn.Mat(decoder_out.numpy()).clone())
ret, ncnn_out0 = ex.extract("out0")
diff --git a/egs/librispeech/ASR/distillation_with_hubert.sh b/egs/librispeech/ASR/distillation_with_hubert.sh
index 2a69d3921..6aaa0333b 100755
--- a/egs/librispeech/ASR/distillation_with_hubert.sh
+++ b/egs/librispeech/ASR/distillation_with_hubert.sh
@@ -35,7 +35,7 @@ stop_stage=4
# export CUDA_VISIBLE_DEVICES="0"
#
# Suppose GPU 2,3,4,5 are available.
-export CUDA_VISIBLE_DEVICES="0,1,2,3"
+# export CUDA_VISIBLE_DEVICES="0,1,2,3"
exp_dir=./pruned_transducer_stateless6/exp
mkdir -p $exp_dir
@@ -43,13 +43,13 @@ mkdir -p $exp_dir
# full_libri can be "True" or "False"
# "True" -> use full librispeech dataset for distillation
# "False" -> use train-clean-100 subset for distillation
-full_libri=False
+full_libri=True
# use_extracted_codebook can be "True" or "False"
# "True" -> stage 0 and stage 1 would be skipped,
# and directly download the extracted codebook indexes for distillation
# "False" -> start from scratch
-use_extracted_codebook=False
+use_extracted_codebook=True
# teacher_model_id can be one of
# "hubert_xtralarge_ll60k_finetune_ls960" -> fine-tuned model, it is the one we currently use.
@@ -145,8 +145,12 @@ if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
log "Currently we only uploaded codebook indexes from teacher model hubert_xtralarge_ll60k_finetune_ls960"
exit 1
fi
+ # The codebook indexes to be downloaded are generated using the following setup:
+ embedding_layer=36
+ num_codebooks=8
+
mkdir -p $exp_dir/vq
- codebook_dir=$exp_dir/vq/$teacher_model_id
+ codebook_dir=$exp_dir/vq/${teacher_model_id}
mkdir -p codebook_dir
codebook_download_dir=$exp_dir/download_codebook
if [ -d $codebook_download_dir ]; then
@@ -155,11 +159,18 @@ if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
fi
log "Downloading extracted codebook indexes to $codebook_download_dir"
# Make sure you have git-lfs installed (https://git-lfs.github.com)
+ # The codebook indexes are generated using lhotse 1.11.0, to avoid
+ # potential issues, we recommend you to use lhotse version >= 1.11.0
+ lhotse_version=$(python3 -c "import lhotse; from packaging import version; print(version.parse(lhotse.version.__version__)>=version.parse('1.11.0'))")
+ if [ "$lhotse_version" == "False" ]; then
+ log "Expecting lhotse >= 1.11.0. This may lead to potential ID mismatch."
+ fi
git lfs install
- git clone https://huggingface.co/Zengwei/pruned_transducer_stateless6_hubert_xtralarge_ll60k_finetune_ls960 $codebook_download_dir
+ git clone https://huggingface.co/marcoyang/pruned_transducer_stateless6_hubert_xtralarge_ll60k_finetune_ls960 $codebook_download_dir
- mkdir -p data/vq_fbank
- mv $codebook_download_dir/*.jsonl.gz data/vq_fbank/
+ vq_fbank=data/vq_fbank_layer${embedding_layer}_cb${num_codebooks}/
+ mkdir -p $vq_fbank
+ mv $codebook_download_dir/*.jsonl.gz $vq_fbank
mkdir -p $codebook_dir/splits4
mv $codebook_download_dir/*.h5 $codebook_dir/splits4/
log "Remove $codebook_download_dir"
@@ -169,12 +180,21 @@ if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
./pruned_transducer_stateless6/extract_codebook_index.py \
--full-libri $full_libri \
--exp-dir $exp_dir \
- --embedding-layer 36 \
+ --embedding-layer $embedding_layer \
--num-utts 1000 \
- --num-codebooks 8 \
+ --num-codebooks $num_codebooks \
--max-duration 100 \
--teacher-model-id $teacher_model_id \
--use-extracted-codebook $use_extracted_codebook
+
+ if [ "$full_libri" == "True" ]; then
+ # Merge the 3 subsets and create a full one
+ rm ${vq_fbank}/librispeech_cuts_train-all-shuf.jsonl.gz
+ cat <(gunzip -c ${vq_fbank}/librispeech_cuts_train-clean-100.jsonl.gz) \
+ <(gunzip -c ${vq_fbank}/librispeech_cuts_train-clean-360.jsonl.gz) \
+ <(gunzip -c ${vq_fbank}/librispeech_cuts_train-other-500.jsonl.gz) | \
+ shuf | gzip -c > ${vq_fbank}/librispeech_cuts_train-all-shuf.jsonl.gz
+ fi
fi
if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
diff --git a/egs/librispeech/ASR/generate-lm.sh b/egs/librispeech/ASR/generate-lm.sh
index 6baccd381..dacd276d1 100755
--- a/egs/librispeech/ASR/generate-lm.sh
+++ b/egs/librispeech/ASR/generate-lm.sh
@@ -2,7 +2,7 @@
lang_dir=data/lang_bpe_500
-for ngram in 2 3 5; do
+for ngram in 2 3 4 5; do
if [ ! -f $lang_dir/${ngram}gram.arpa ]; then
./shared/make_kn_lm.py \
-ngram-order ${ngram} \
diff --git a/egs/librispeech/ASR/local/compile_hlg.py b/egs/librispeech/ASR/local/compile_hlg.py
index df6c609bb..08dac6a7b 100755
--- a/egs/librispeech/ASR/local/compile_hlg.py
+++ b/egs/librispeech/ASR/local/compile_hlg.py
@@ -24,7 +24,7 @@ This script takes as input lang_dir and generates HLG from
Caution: We use a lexicon that contains disambiguation symbols
- - G, the LM, built from data/lm/G_3_gram.fst.txt
+ - G, the LM, built from data/lm/G_n_gram.fst.txt
The generated HLG is saved in $lang_dir/HLG.pt
"""
diff --git a/egs/librispeech/ASR/local/compile_hlg_using_openfst.py b/egs/librispeech/ASR/local/compile_hlg_using_openfst.py
index 9e5e3df69..15fc47ef1 100755
--- a/egs/librispeech/ASR/local/compile_hlg_using_openfst.py
+++ b/egs/librispeech/ASR/local/compile_hlg_using_openfst.py
@@ -24,7 +24,7 @@ This script takes as input lang_dir and generates HLG from
Caution: We use a lexicon that contains disambiguation symbols
- - G, the LM, built from data/lm/G_3_gram.fst.txt
+ - G, the LM, built from data/lm/G_n_gram.fst.txt
The generated HLG is saved in $lang_dir/HLG_fst.pt
@@ -46,6 +46,13 @@ from icefall.lexicon import Lexicon
def get_args():
parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--lm",
+ type=str,
+ default="G_3_gram",
+ help="""Stem name for LM used in HLG compiling.
+ """,
+ )
parser.add_argument(
"--lang-dir",
type=str,
@@ -56,11 +63,13 @@ def get_args():
return parser.parse_args()
-def compile_HLG(lang_dir: str) -> kaldifst.StdVectorFst:
+def compile_HLG(lang_dir: str, lm: str = "G_3_gram") -> kaldifst.StdVectorFst:
"""
Args:
lang_dir:
The language directory, e.g., data/lang_phone or data/lang_bpe_5000.
+ lm:
+ The language stem base name.
Return:
An FST representing HLG.
@@ -71,8 +80,8 @@ def compile_HLG(lang_dir: str) -> kaldifst.StdVectorFst:
kaldifst.arcsort(L, sort_type="olabel")
logging.info(f"L: #states {L.num_states}")
- G_filename_txt = "data/lm/G_3_gram.fst.txt"
- G_filename_binary = "data/lm/G_3_gram.fst"
+ G_filename_txt = f"data/lm/{lm}.fst.txt"
+ G_filename_binary = f"data/lm/{lm}.fst"
if Path(G_filename_binary).is_file():
logging.info(f"Loading {G_filename_binary}")
G = kaldifst.StdVectorFst.read(G_filename_binary)
@@ -171,7 +180,7 @@ def main():
logging.info(f"{filename} already exists - skipping")
return
- HLG = compile_HLG(lang_dir)
+ HLG = compile_HLG(lang_dir, args.lm)
logging.info(f"Saving HLG to {filename}")
torch.save(HLG.as_dict(), filename)
diff --git a/egs/librispeech/ASR/local/compute_fbank_musan.py b/egs/librispeech/ASR/local/compute_fbank_musan.py
index 4a4093ae4..62036467e 100755
--- a/egs/librispeech/ASR/local/compute_fbank_musan.py
+++ b/egs/librispeech/ASR/local/compute_fbank_musan.py
@@ -28,7 +28,7 @@ import os
from pathlib import Path
import torch
-from lhotse import CutSet, Fbank, FbankConfig, LilcomChunkyWriter, combine
+from lhotse import CutSet, Fbank, FbankConfig, LilcomChunkyWriter, MonoCut, combine
from lhotse.recipes.utils import read_manifests_if_cached
from icefall.utils import get_executor
@@ -41,6 +41,10 @@ torch.set_num_threads(1)
torch.set_num_interop_threads(1)
+def is_cut_long(c: MonoCut) -> bool:
+ return c.duration > 5
+
+
def compute_fbank_musan():
src_dir = Path("data/manifests")
output_dir = Path("data/fbank")
@@ -86,7 +90,7 @@ def compute_fbank_musan():
recordings=combine(part["recordings"] for part in manifests.values())
)
.cut_into_windows(10.0)
- .filter(lambda c: c.duration > 5)
+ .filter(is_cut_long)
.compute_and_store_features(
extractor=extractor,
storage_path=f"{output_dir}/musan_feats",
diff --git a/egs/librispeech/ASR/local/prepare_lang_bpe.py b/egs/librispeech/ASR/local/prepare_lang_bpe.py
index e121aefa9..2a2d9c219 100755
--- a/egs/librispeech/ASR/local/prepare_lang_bpe.py
+++ b/egs/librispeech/ASR/local/prepare_lang_bpe.py
@@ -127,7 +127,7 @@ def lexicon_to_fst_no_sil(
def generate_lexicon(
- model_file: str, words: List[str]
+ model_file: str, words: List[str], oov: str
) -> Tuple[Lexicon, Dict[str, int]]:
"""Generate a lexicon from a BPE model.
@@ -136,6 +136,8 @@ def generate_lexicon(
Path to a sentencepiece model.
words:
A list of strings representing words.
+ oov:
+ The out of vocabulary word in lexicon.
Returns:
Return a tuple with two elements:
- A dict whose keys are words and values are the corresponding
@@ -156,12 +158,9 @@ def generate_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())]))
+ lexicon.append((oov, ["▁", 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
+ token2id: Dict[str, int] = {sp.id_to_piece(i): i for i in range(sp.vocab_size())}
return lexicon, token2id
@@ -176,6 +175,13 @@ def get_args():
""",
)
+ parser.add_argument(
+ "--oov",
+ type=str,
+ default="",
+ help="The out of vocabulary word in lexicon.",
+ )
+
parser.add_argument(
"--debug",
type=str2bool,
@@ -202,12 +208,13 @@ def main():
words = word_sym_table.symbols
- excluded = ["", "!SIL", "", "", "#0", "", ""]
+ excluded = ["", "!SIL", "", args.oov, "#0", "", ""]
+
for w in excluded:
if w in words:
words.remove(w)
- lexicon, token_sym_table = generate_lexicon(model_file, words)
+ lexicon, token_sym_table = generate_lexicon(model_file, words, args.oov)
lexicon_disambig, max_disambig = add_disambig_symbols(lexicon)
diff --git a/egs/librispeech/ASR/lstm_transducer_stateless2/decode.py b/egs/librispeech/ASR/lstm_transducer_stateless2/decode.py
index fa5bf1825..78be9c01f 100755
--- a/egs/librispeech/ASR/lstm_transducer_stateless2/decode.py
+++ b/egs/librispeech/ASR/lstm_transducer_stateless2/decode.py
@@ -93,36 +93,37 @@ Usage:
--max-contexts 8 \
--max-states 64
-(8) modified beam search (with RNNLM shallow fusion)
+(8) modified beam search (with LM shallow fusion)
./lstm_transducer_stateless2/decode.py \
--epoch 35 \
--avg 15 \
--exp-dir ./lstm_transducer_stateless2/exp \
--max-duration 600 \
- --decoding-method modified_beam_search_rnnlm_shallow_fusion \
+ --decoding-method modified_beam_search_lm_shallow_fusion \
--beam 4 \
- --rnn-lm-scale 0.3 \
- --rnn-lm-exp-dir /path/to/RNNLM \
+ --lm-type rnn \
+ --lm-scale 0.3 \
+ --lm-exp-dir /path/to/LM \
--rnn-lm-epoch 99 \
--rnn-lm-avg 1 \
--rnn-lm-num-layers 3 \
--rnn-lm-tie-weights 1
-(9) modified beam search with RNNLM shallow fusion + LODR
+(9) modified beam search with LM shallow fusion + LODR
./lstm_transducer_stateless2/decode.py \
--epoch 35 \
--avg 15 \
--max-duration 600 \
--exp-dir ./lstm_transducer_stateless2/exp \
- --decoding-method modified_beam_search_rnnlm_LODR \
+ --decoding-method modified_beam_search_LODR \
--beam 4 \
- --max-contexts 4 \
- --rnn-lm-scale 0.4 \
- --rnn-lm-exp-dir /path/to/RNNLM/exp \
+ --lm-type rnn \
+ --lm-scale 0.4 \
+ --lm-exp-dir /path/to/LM \
--rnn-lm-epoch 99 \
--rnn-lm-avg 1 \
--rnn-lm-num-layers 3 \
- --rnn-lm-tie-weights 1 \
+ --rnn-lm-tie-weights 1
--tokens-ngram 2 \
--ngram-lm-scale -0.16 \
"""
@@ -148,14 +149,14 @@ from beam_search import (
greedy_search,
greedy_search_batch,
modified_beam_search,
+ modified_beam_search_lm_shallow_fusion,
+ modified_beam_search_LODR,
modified_beam_search_ngram_rescoring,
- modified_beam_search_rnnlm_LODR,
- modified_beam_search_rnnlm_shallow_fusion,
)
from librispeech import LibriSpeech
from train import add_model_arguments, get_params, get_transducer_model
-from icefall import NgramLm
+from icefall import LmScorer, NgramLm
from icefall.checkpoint import (
average_checkpoints,
average_checkpoints_with_averaged_model,
@@ -163,7 +164,6 @@ from icefall.checkpoint import (
load_checkpoint,
)
from icefall.lexicon import Lexicon
-from icefall.rnn_lm.model import RnnLmModel
from icefall.utils import (
AttributeDict,
setup_logger,
@@ -253,8 +253,8 @@ def get_parser():
- fast_beam_search_nbest_oracle
- fast_beam_search_nbest_LG
- modified_beam_search_ngram_rescoring
- - modified_beam_search_rnnlm_shallow_fusion
- - modified_beam_search_rnnlm_LODR
+ - modified_beam_search_lm_shallow_fusion
+ - modified_beam_search_LODR
If you use fast_beam_search_nbest_LG, you have to specify
`--lang-dir`, which should contain `LG.pt`.
""",
@@ -344,67 +344,28 @@ def get_parser():
)
parser.add_argument(
- "--rnn-lm-scale",
- type=float,
- default=0.0,
- help="""Used only when --method is modified-beam-search_rnnlm_shallow_fusion.
- It specifies the path to RNN LM exp dir.
- """,
- )
-
- parser.add_argument(
- "--rnn-lm-exp-dir",
- type=str,
- default="rnn_lm/exp",
- help="""Used only when --method is modified_beam_search_rnnlm_shallow_fusion.
- It specifies the path to RNN LM exp dir.
- """,
- )
-
- parser.add_argument(
- "--rnn-lm-epoch",
- type=int,
- default=7,
- help="""Used only when --method is modified_beam_search_rnnlm_shallow_fusion.
- It specifies the checkpoint to use.
- """,
- )
-
- parser.add_argument(
- "--rnn-lm-avg",
- type=int,
- default=2,
- help="""Used only when --method is modified_beam_search_rnnlm_shallow_fusion.
- It specifies the number of checkpoints to average.
- """,
- )
-
- parser.add_argument(
- "--rnn-lm-embedding-dim",
- type=int,
- default=2048,
- help="Embedding dim of the model",
- )
-
- parser.add_argument(
- "--rnn-lm-hidden-dim",
- type=int,
- default=2048,
- help="Hidden dim of the model",
- )
-
- parser.add_argument(
- "--rnn-lm-num-layers",
- type=int,
- default=4,
- help="Number of RNN layers the model",
- )
- parser.add_argument(
- "--rnn-lm-tie-weights",
+ "--use-shallow-fusion",
type=str2bool,
default=False,
- help="""True to share the weights between the input embedding layer and the
- last output linear layer
+ help="""Use neural network LM for shallow fusion.
+ If you want to use LODR, you will also need to set this to true
+ """,
+ )
+
+ parser.add_argument(
+ "--lm-type",
+ type=str,
+ default="rnn",
+ help="Type of NN lm",
+ choices=["rnn", "transformer"],
+ )
+
+ parser.add_argument(
+ "--lm-scale",
+ type=float,
+ default=0.3,
+ help="""The scale of the neural network LM
+ Used only when `--use-shallow-fusion` is set to True.
""",
)
@@ -440,8 +401,7 @@ def decode_one_batch(
decoding_graph: Optional[k2.Fsa] = None,
ngram_lm: Optional[NgramLm] = None,
ngram_lm_scale: float = 1.0,
- rnnlm: Optional[RnnLmModel] = None,
- rnnlm_scale: float = 1.0,
+ LM: Optional[LmScorer] = None,
) -> Dict[str, List[List[str]]]:
"""Decode one batch and return the result in a dict. The dict has the
following format:
@@ -470,6 +430,9 @@ def decode_one_batch(
The decoding graph. Can be either a `k2.trivial_graph` or LG, Used
only when --decoding_method is fast_beam_search, fast_beam_search_nbest,
fast_beam_search_nbest_oracle, and fast_beam_search_nbest_LG.
+ LM:
+ A neural net LM for shallow fusion. Only used when `--use-shallow-fusion`
+ set to true.
Returns:
Return the decoding result. See above description for the format of
the returned dict.
@@ -581,20 +544,19 @@ def decode_one_batch(
)
for hyp in sp.decode(hyp_tokens):
hyps.append(hyp.split())
- elif params.decoding_method == "modified_beam_search_rnnlm_shallow_fusion":
- hyp_tokens = modified_beam_search_rnnlm_shallow_fusion(
+ elif params.decoding_method == "modified_beam_search_lm_shallow_fusion":
+ hyp_tokens = modified_beam_search_lm_shallow_fusion(
model=model,
encoder_out=encoder_out,
encoder_out_lens=encoder_out_lens,
beam=params.beam_size,
sp=sp,
- rnnlm=rnnlm,
- rnnlm_scale=rnnlm_scale,
+ LM=LM,
)
for hyp in sp.decode(hyp_tokens):
hyps.append(hyp.split())
- elif params.decoding_method == "modified_beam_search_rnnlm_LODR":
- hyp_tokens = modified_beam_search_rnnlm_LODR(
+ elif params.decoding_method == "modified_beam_search_LODR":
+ hyp_tokens = modified_beam_search_LODR(
model=model,
encoder_out=encoder_out,
encoder_out_lens=encoder_out_lens,
@@ -602,8 +564,7 @@ def decode_one_batch(
sp=sp,
LODR_lm=ngram_lm,
LODR_lm_scale=ngram_lm_scale,
- rnnlm=rnnlm,
- rnnlm_scale=rnnlm_scale,
+ LM=LM,
)
for hyp in sp.decode(hyp_tokens):
hyps.append(hyp.split())
@@ -658,8 +619,7 @@ def decode_dataset(
decoding_graph: Optional[k2.Fsa] = None,
ngram_lm: Optional[NgramLm] = None,
ngram_lm_scale: float = 1.0,
- rnnlm: Optional[RnnLmModel] = None,
- rnnlm_scale: float = 1.0,
+ LM: Optional[LmScorer] = None,
) -> Dict[str, List[Tuple[str, List[str], List[str]]]]:
"""Decode dataset.
@@ -678,6 +638,8 @@ def decode_dataset(
The decoding graph. Can be either a `k2.trivial_graph` or LG, Used
only when --decoding_method is fast_beam_search, fast_beam_search_nbest,
fast_beam_search_nbest_oracle, and fast_beam_search_nbest_LG.
+ LM:
+ A neural network LM, used during shallow fusion
Returns:
Return a dict, whose key may be "greedy_search" if greedy search
is used, or it may be "beam_7" if beam size of 7 is used.
@@ -711,8 +673,7 @@ def decode_dataset(
batch=batch,
ngram_lm=ngram_lm,
ngram_lm_scale=ngram_lm_scale,
- rnnlm=rnnlm,
- rnnlm_scale=rnnlm_scale,
+ LM=LM,
)
for name, hyps in hyps_dict.items():
@@ -730,6 +691,7 @@ def decode_dataset(
batch_str = f"{batch_idx}/{num_batches}"
logging.info(f"batch {batch_str}, cuts processed until now is {num_cuts}")
+ logging.info(f"batch {batch_str}, cuts processed until now is {num_cuts}")
return results
@@ -781,6 +743,7 @@ def save_results(
def main():
parser = get_parser()
AsrDataModule.add_arguments(parser)
+ LmScorer.add_arguments(parser)
args = parser.parse_args()
args.exp_dir = Path(args.exp_dir)
@@ -795,9 +758,9 @@ def main():
"fast_beam_search_nbest_LG",
"fast_beam_search_nbest_oracle",
"modified_beam_search",
- "modified_beam_search_rnnlm_LODR",
+ "modified_beam_search_LODR",
+ "modified_beam_search_lm_shallow_fusion",
"modified_beam_search_ngram_rescoring",
- "modified_beam_search_rnnlm_shallow_fusion",
)
params.res_dir = params.exp_dir / params.decoding_method
@@ -820,12 +783,18 @@ def main():
else:
params.suffix += f"-context-{params.context_size}"
params.suffix += f"-max-sym-per-frame-{params.max_sym_per_frame}"
- params.suffix += f"-ngram-lm-scale-{params.ngram_lm_scale}"
- if "rnnlm" in params.decoding_method:
- params.suffix += f"-rnnlm-lm-scale-{params.rnn_lm_scale}"
+ if "ngram" in params.decoding_method:
+ params.suffix += f"-ngram-lm-scale-{params.ngram_lm_scale}"
+ if params.use_shallow_fusion:
+ if params.lm_type == "rnn":
+ params.suffix += f"-rnnlm-lm-scale-{params.lm_scale}"
+ elif params.lm_type == "transformer":
+ params.suffix += f"-transformer-lm-scale-{params.lm_scale}"
- if "LODR" in params.decoding_method:
- params.suffix += "-LODR"
+ if "LODR" in params.decoding_method:
+ params.suffix += (
+ f"-LODR-{params.tokens_ngram}gram-scale-{params.ngram_lm_scale}"
+ )
if params.use_averaged_model:
params.suffix += "-use-averaged-model"
@@ -954,28 +923,19 @@ def main():
ngram_lm = None
ngram_lm_scale = None
- # only load rnnlm if used
- if "rnnlm" in params.decoding_method:
- rnn_lm_scale = params.rnn_lm_scale
-
- rnn_lm_model = RnnLmModel(
- vocab_size=params.vocab_size,
- embedding_dim=params.rnn_lm_embedding_dim,
- hidden_dim=params.rnn_lm_hidden_dim,
- num_layers=params.rnn_lm_num_layers,
- tie_weights=params.rnn_lm_tie_weights,
+ # only load the neural network LM if doing shallow fusion
+ if params.use_shallow_fusion:
+ LM = LmScorer(
+ lm_type=params.lm_type,
+ params=params,
+ device=device,
+ lm_scale=params.lm_scale,
)
- assert params.rnn_lm_avg == 1
+ LM.to(device)
+ LM.eval()
- load_checkpoint(
- f"{params.rnn_lm_exp_dir}/epoch-{params.rnn_lm_epoch}.pt",
- rnn_lm_model,
- )
- rnn_lm_model.to(device)
- rnn_lm_model.eval()
else:
- rnn_lm_model = None
- rnn_lm_scale = 0.0
+ LM = None
if "fast_beam_search" in params.decoding_method:
if params.decoding_method == "fast_beam_search_nbest_LG":
@@ -1003,7 +963,9 @@ def main():
librispeech = LibriSpeech(manifest_dir=args.manifest_dir)
test_clean_cuts = librispeech.test_clean_cuts()
+ # test_clean_cuts = test_clean_cuts.subset(first=500)
test_other_cuts = librispeech.test_other_cuts()
+ # test_other_cuts = test_other_cuts.subset(first=500)
test_clean_dl = asr_datamodule.test_dataloaders(test_clean_cuts)
test_other_dl = asr_datamodule.test_dataloaders(test_other_cuts)
@@ -1021,8 +983,7 @@ def main():
decoding_graph=decoding_graph,
ngram_lm=ngram_lm,
ngram_lm_scale=ngram_lm_scale,
- rnnlm=rnn_lm_model,
- rnnlm_scale=rnn_lm_scale,
+ LM=LM,
)
save_results(
diff --git a/egs/librispeech/ASR/lstm_transducer_stateless2/ncnn-decode.py b/egs/librispeech/ASR/lstm_transducer_stateless2/ncnn-decode.py
index 3b471fa85..3bd1b0a09 100755
--- a/egs/librispeech/ASR/lstm_transducer_stateless2/ncnn-decode.py
+++ b/egs/librispeech/ASR/lstm_transducer_stateless2/ncnn-decode.py
@@ -104,6 +104,8 @@ class Model:
encoder_net = ncnn.Net()
encoder_net.opt.use_packing_layout = False
encoder_net.opt.use_fp16_storage = False
+ encoder_net.opt.num_threads = 4
+
encoder_param = args.encoder_param_filename
encoder_model = args.encoder_bin_filename
@@ -118,6 +120,7 @@ class Model:
decoder_net = ncnn.Net()
decoder_net.opt.use_packing_layout = False
+ decoder_net.opt.num_threads = 4
decoder_net.load_param(decoder_param)
decoder_net.load_model(decoder_model)
@@ -129,6 +132,8 @@ class Model:
joiner_model = args.joiner_bin_filename
joiner_net = ncnn.Net()
joiner_net.opt.use_packing_layout = False
+ joiner_net.opt.num_threads = 4
+
joiner_net.load_param(joiner_param)
joiner_net.load_model(joiner_model)
@@ -136,7 +141,6 @@ class Model:
def run_encoder(self, x, states):
with self.encoder_net.create_extractor() as ex:
- ex.set_num_threads(10)
ex.input("in0", ncnn.Mat(x.numpy()).clone())
x_lens = torch.tensor([x.size(0)], dtype=torch.float32)
ex.input("in1", ncnn.Mat(x_lens.numpy()).clone())
@@ -165,7 +169,6 @@ class Model:
assert decoder_input.dtype == torch.int32
with self.decoder_net.create_extractor() as ex:
- ex.set_num_threads(10)
ex.input("in0", ncnn.Mat(decoder_input.numpy()).clone())
ret, ncnn_out0 = ex.extract("out0")
assert ret == 0, ret
@@ -174,7 +177,6 @@ class Model:
def run_joiner(self, encoder_out, decoder_out):
with self.joiner_net.create_extractor() as ex:
- ex.set_num_threads(10)
ex.input("in0", ncnn.Mat(encoder_out.numpy()).clone())
ex.input("in1", ncnn.Mat(decoder_out.numpy()).clone())
ret, ncnn_out0 = ex.extract("out0")
diff --git a/egs/librispeech/ASR/lstm_transducer_stateless2/streaming-ncnn-decode.py b/egs/librispeech/ASR/lstm_transducer_stateless2/streaming-ncnn-decode.py
index baff15ea6..02ed16a8c 100755
--- a/egs/librispeech/ASR/lstm_transducer_stateless2/streaming-ncnn-decode.py
+++ b/egs/librispeech/ASR/lstm_transducer_stateless2/streaming-ncnn-decode.py
@@ -92,6 +92,8 @@ class Model:
encoder_net = ncnn.Net()
encoder_net.opt.use_packing_layout = False
encoder_net.opt.use_fp16_storage = False
+ encoder_net.opt.num_threads = 4
+
encoder_param = args.encoder_param_filename
encoder_model = args.encoder_bin_filename
@@ -106,6 +108,7 @@ class Model:
decoder_net = ncnn.Net()
decoder_net.opt.use_packing_layout = False
+ decoder_net.opt.num_threads = 4
decoder_net.load_param(decoder_param)
decoder_net.load_model(decoder_model)
@@ -117,6 +120,8 @@ class Model:
joiner_model = args.joiner_bin_filename
joiner_net = ncnn.Net()
joiner_net.opt.use_packing_layout = False
+ joiner_net.opt.num_threads = 4
+
joiner_net.load_param(joiner_param)
joiner_net.load_model(joiner_model)
@@ -124,7 +129,6 @@ class Model:
def run_encoder(self, x, states):
with self.encoder_net.create_extractor() as ex:
- # ex.set_num_threads(10)
ex.input("in0", ncnn.Mat(x.numpy()).clone())
x_lens = torch.tensor([x.size(0)], dtype=torch.float32)
ex.input("in1", ncnn.Mat(x_lens.numpy()).clone())
@@ -153,7 +157,6 @@ class Model:
assert decoder_input.dtype == torch.int32
with self.decoder_net.create_extractor() as ex:
- # ex.set_num_threads(10)
ex.input("in0", ncnn.Mat(decoder_input.numpy()).clone())
ret, ncnn_out0 = ex.extract("out0")
assert ret == 0, ret
@@ -162,7 +165,6 @@ class Model:
def run_joiner(self, encoder_out, decoder_out):
with self.joiner_net.create_extractor() as ex:
- # ex.set_num_threads(10)
ex.input("in0", ncnn.Mat(encoder_out.numpy()).clone())
ex.input("in1", ncnn.Mat(decoder_out.numpy()).clone())
ret, ncnn_out0 = ex.extract("out0")
diff --git a/egs/librispeech/ASR/prepare.sh b/egs/librispeech/ASR/prepare.sh
index 11c8e1066..b1d207049 100755
--- a/egs/librispeech/ASR/prepare.sh
+++ b/egs/librispeech/ASR/prepare.sh
@@ -123,10 +123,12 @@ if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
touch data/fbank/.librispeech.done
fi
- cat <(gunzip -c data/fbank/librispeech_cuts_train-clean-100.jsonl.gz) \
- <(gunzip -c data/fbank/librispeech_cuts_train-clean-360.jsonl.gz) \
- <(gunzip -c data/fbank/librispeech_cuts_train-other-500.jsonl.gz) | \
- shuf | gzip -c > data/fbank/librispeech_cuts_train-all-shuf.jsonl.gz
+ if [ ! -f data/fbank/librispeech_cuts_train-all-shuf.jsonl.gz ]; then
+ cat <(gunzip -c data/fbank/librispeech_cuts_train-clean-100.jsonl.gz) \
+ <(gunzip -c data/fbank/librispeech_cuts_train-clean-360.jsonl.gz) \
+ <(gunzip -c data/fbank/librispeech_cuts_train-other-500.jsonl.gz) | \
+ shuf | gzip -c > data/fbank/librispeech_cuts_train-all-shuf.jsonl.gz
+ fi
if [ ! -e data/fbank/.librispeech-validated.done ]; then
log "Validating data/fbank for LibriSpeech"
@@ -244,7 +246,7 @@ if [ $stage -le 6 ] && [ $stop_stage -ge 6 ]; then
fi
if [ $stage -le 7 ] && [ $stop_stage -ge 7 ]; then
- log "Stage 7: Prepare bigram P"
+ log "Stage 7: Prepare bigram token-level P for MMI training"
for vocab_size in ${vocab_sizes[@]}; do
lang_dir=data/lang_bpe_${vocab_size}
@@ -302,13 +304,20 @@ fi
if [ $stage -le 9 ] && [ $stop_stage -ge 9 ]; then
log "Stage 9: Compile HLG"
./local/compile_hlg.py --lang-dir data/lang_phone
- ./local/compile_hlg_using_openfst.py --lang-dir data/lang_phone
+
+ # Note If ./local/compile_hlg.py throws OOM,
+ # please switch to the following command
+ #
+ # ./local/compile_hlg_using_openfst.py --lang-dir data/lang_phone
for vocab_size in ${vocab_sizes[@]}; do
lang_dir=data/lang_bpe_${vocab_size}
./local/compile_hlg.py --lang-dir $lang_dir
- ./local/compile_hlg_using_openfst.py --lang-dir $lang_dir
+ # Note If ./local/compile_hlg.py throws OOM,
+ # please switch to the following command
+ #
+ # ./local/compile_hlg_using_openfst.py --lang-dir $lang_dir
done
fi
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless2/beam_search.py b/egs/librispeech/ASR/pruned_transducer_stateless2/beam_search.py
index b324cc9b7..7388af389 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless2/beam_search.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless2/beam_search.py
@@ -26,7 +26,9 @@ from model import Transducer
from icefall import NgramLm, NgramLmStateCost
from icefall.decode import Nbest, one_best_decoding
+from icefall.lm_wrapper import LmScorer
from icefall.rnn_lm.model import RnnLmModel
+from icefall.transformer_lm.model import TransformerLM
from icefall.utils import (
DecodingResults,
add_eos,
@@ -1846,254 +1848,14 @@ def modified_beam_search_ngram_rescoring(
return ans
-def modified_beam_search_rnnlm_shallow_fusion(
- model: Transducer,
- encoder_out: torch.Tensor,
- encoder_out_lens: torch.Tensor,
- sp: spm.SentencePieceProcessor,
- rnnlm: RnnLmModel,
- rnnlm_scale: float,
- beam: int = 4,
- return_timestamps: bool = False,
-) -> List[List[int]]:
- """Modified_beam_search + RNNLM shallow fusion
-
- Args:
- model (Transducer):
- The transducer model
- encoder_out (torch.Tensor):
- Encoder output in (N,T,C)
- encoder_out_lens (torch.Tensor):
- A 1-D tensor of shape (N,), containing the number of
- valid frames in encoder_out before padding.
- sp:
- Sentence piece generator.
- rnnlm (RnnLmModel):
- RNNLM
- rnnlm_scale (float):
- scale of RNNLM in shallow fusion
- beam (int, optional):
- Beam size. Defaults to 4.
-
- Returns:
- Return a list-of-list of token IDs. ans[i] is the decoding results
- for the i-th utterance.
- """
- assert encoder_out.ndim == 3, encoder_out.shape
- assert encoder_out.size(0) >= 1, encoder_out.size(0)
- assert rnnlm is not None
- lm_scale = rnnlm_scale
- vocab_size = rnnlm.vocab_size
-
- packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
- input=encoder_out,
- lengths=encoder_out_lens.cpu(),
- batch_first=True,
- enforce_sorted=False,
- )
-
- blank_id = model.decoder.blank_id
- sos_id = sp.piece_to_id("")
- unk_id = getattr(model, "unk_id", blank_id)
- context_size = model.decoder.context_size
- device = next(model.parameters()).device
-
- batch_size_list = packed_encoder_out.batch_sizes.tolist()
- N = encoder_out.size(0)
- assert torch.all(encoder_out_lens > 0), encoder_out_lens
- assert N == batch_size_list[0], (N, batch_size_list)
-
- # get initial lm score and lm state by scoring the "sos" token
- sos_token = torch.tensor([[sos_id]]).to(torch.int64).to(device)
- init_score, init_states = rnnlm.score_token(sos_token)
-
- B = [HypothesisList() for _ in range(N)]
- for i in range(N):
- B[i].add(
- Hypothesis(
- ys=[blank_id] * context_size,
- log_prob=torch.zeros(1, dtype=torch.float32, device=device),
- state=init_states,
- lm_score=init_score.reshape(-1),
- timestamp=[],
- )
- )
-
- rnnlm.clean_cache()
- encoder_out = model.joiner.encoder_proj(packed_encoder_out.data)
-
- offset = 0
- finalized_B = []
- for (t, batch_size) in enumerate(batch_size_list):
- start = offset
- end = offset + batch_size
- current_encoder_out = encoder_out.data[start:end] # get batch
- current_encoder_out = current_encoder_out.unsqueeze(1).unsqueeze(1)
- # current_encoder_out's shape is (batch_size, 1, 1, encoder_out_dim)
- offset = end
-
- finalized_B = B[batch_size:] + finalized_B
- B = B[:batch_size]
-
- hyps_shape = get_hyps_shape(B).to(device)
-
- A = [list(b) for b in B]
- B = [HypothesisList() for _ in range(batch_size)]
-
- ys_log_probs = torch.cat(
- [hyp.log_prob.reshape(1, 1) for hyps in A for hyp in hyps]
- )
-
- decoder_input = torch.tensor(
- [hyp.ys[-context_size:] for hyps in A for hyp in hyps],
- device=device,
- dtype=torch.int64,
- ) # (num_hyps, context_size)
-
- decoder_out = model.decoder(decoder_input, need_pad=False).unsqueeze(1)
- decoder_out = model.joiner.decoder_proj(decoder_out)
-
- current_encoder_out = torch.index_select(
- current_encoder_out,
- dim=0,
- index=hyps_shape.row_ids(1).to(torch.int64),
- ) # (num_hyps, 1, 1, encoder_out_dim)
-
- logits = model.joiner(
- current_encoder_out,
- decoder_out,
- project_input=False,
- ) # (num_hyps, 1, 1, vocab_size)
-
- logits = logits.squeeze(1).squeeze(1) # (num_hyps, vocab_size)
-
- log_probs = logits.log_softmax(dim=-1) # (num_hyps, vocab_size)
-
- log_probs.add_(ys_log_probs)
-
- vocab_size = log_probs.size(-1)
-
- log_probs = log_probs.reshape(-1)
-
- row_splits = hyps_shape.row_splits(1) * vocab_size
- log_probs_shape = k2.ragged.create_ragged_shape2(
- row_splits=row_splits, cached_tot_size=log_probs.numel()
- )
- ragged_log_probs = k2.RaggedTensor(shape=log_probs_shape, value=log_probs)
- """
- for all hyps with a non-blank new token, score this token.
- It is a little confusing here because this for-loop
- looks very similar to the one below. Here, we go through all
- top-k tokens and only add the non-blanks ones to the token_list.
- The RNNLM will score those tokens given the LM states. Note that
- the variable `scores` is the LM score after seeing the new
- non-blank token.
- """
- token_list = []
- hs = []
- cs = []
- for i in range(batch_size):
- topk_log_probs, topk_indexes = ragged_log_probs[i].topk(beam)
-
- with warnings.catch_warnings():
- warnings.simplefilter("ignore")
- topk_hyp_indexes = (topk_indexes // vocab_size).tolist()
- topk_token_indexes = (topk_indexes % vocab_size).tolist()
- for k in range(len(topk_hyp_indexes)):
- hyp_idx = topk_hyp_indexes[k]
- hyp = A[i][hyp_idx]
-
- new_token = topk_token_indexes[k]
- if new_token not in (blank_id, unk_id):
- assert new_token != 0, new_token
- token_list.append([new_token])
- # store the LSTM states
- hs.append(hyp.state[0])
- cs.append(hyp.state[1])
-
- # forward RNNLM to get new states and scores
- if len(token_list) != 0:
- tokens_to_score = (
- torch.tensor(token_list).to(torch.int64).to(device).reshape(-1, 1)
- )
-
- hs = torch.cat(hs, dim=1).to(device)
- cs = torch.cat(cs, dim=1).to(device)
- scores, lm_states = rnnlm.score_token(tokens_to_score, (hs, cs))
-
- count = 0 # index, used to locate score and lm states
- for i in range(batch_size):
- topk_log_probs, topk_indexes = ragged_log_probs[i].topk(beam)
-
- with warnings.catch_warnings():
- warnings.simplefilter("ignore")
- topk_hyp_indexes = (topk_indexes // vocab_size).tolist()
- topk_token_indexes = (topk_indexes % vocab_size).tolist()
-
- for k in range(len(topk_hyp_indexes)):
- hyp_idx = topk_hyp_indexes[k]
- hyp = A[i][hyp_idx]
-
- ys = hyp.ys[:]
-
- lm_score = hyp.lm_score
- state = hyp.state
-
- hyp_log_prob = topk_log_probs[k] # get score of current hyp
- new_token = topk_token_indexes[k]
- new_timestamp = hyp.timestamp[:]
- if new_token not in (blank_id, unk_id):
-
- ys.append(new_token)
- new_timestamp.append(t)
- hyp_log_prob += lm_score[new_token] * lm_scale # add the lm score
-
- lm_score = scores[count]
- state = (
- lm_states[0][:, count, :].unsqueeze(1),
- lm_states[1][:, count, :].unsqueeze(1),
- )
- count += 1
-
- new_hyp = Hypothesis(
- ys=ys,
- log_prob=hyp_log_prob,
- state=state,
- lm_score=lm_score,
- timestamp=new_timestamp,
- )
- B[i].add(new_hyp)
-
- B = B + finalized_B
- best_hyps = [b.get_most_probable(length_norm=True) for b in B]
-
- sorted_ans = [h.ys[context_size:] for h in best_hyps]
- sorted_timestamps = [h.timestamp for h in best_hyps]
- ans = []
- ans_timestamps = []
- unsorted_indices = packed_encoder_out.unsorted_indices.tolist()
- for i in range(N):
- ans.append(sorted_ans[unsorted_indices[i]])
- ans_timestamps.append(sorted_timestamps[unsorted_indices[i]])
-
- if not return_timestamps:
- return ans
- else:
- return DecodingResults(
- tokens=ans,
- timestamps=ans_timestamps,
- )
-
-
-def modified_beam_search_rnnlm_LODR(
+def modified_beam_search_LODR(
model: Transducer,
encoder_out: torch.Tensor,
encoder_out_lens: torch.Tensor,
sp: spm.SentencePieceProcessor,
LODR_lm: NgramLm,
LODR_lm_scale: float,
- rnnlm: RnnLmModel,
- rnnlm_scale: float,
+ LM: LmScorer,
beam: int = 4,
) -> List[List[int]]:
"""This function implements LODR (https://arxiv.org/abs/2203.16776) with
@@ -2113,13 +1875,11 @@ def modified_beam_search_rnnlm_LODR(
sp:
Sentence piece generator.
LODR_lm:
- A low order n-gram LM
+ A low order n-gram LM, whose score will be subtracted during shallow fusion
LODR_lm_scale:
The scale of the LODR_lm
- rnnlm (RnnLmModel):
- RNNLM, the external language model
- rnnlm_scale (float):
- scale of RNNLM in shallow fusion
+ LM:
+ A neural net LM, e.g an RNNLM or transformer LM
beam (int, optional):
Beam size. Defaults to 4.
@@ -2130,9 +1890,8 @@ def modified_beam_search_rnnlm_LODR(
"""
assert encoder_out.ndim == 3, encoder_out.shape
assert encoder_out.size(0) >= 1, encoder_out.size(0)
- assert rnnlm is not None
- lm_scale = rnnlm_scale
- vocab_size = rnnlm.vocab_size
+ assert LM is not None
+ lm_scale = LM.lm_scale
packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
input=encoder_out,
@@ -2154,7 +1913,8 @@ def modified_beam_search_rnnlm_LODR(
# get initial lm score and lm state by scoring the "sos" token
sos_token = torch.tensor([[sos_id]]).to(torch.int64).to(device)
- init_score, init_states = rnnlm.score_token(sos_token)
+ lens = torch.tensor([1]).to(device)
+ init_score, init_states = LM.score_token(sos_token, lens)
B = [HypothesisList() for _ in range(N)]
for i in range(N):
@@ -2162,7 +1922,7 @@ def modified_beam_search_rnnlm_LODR(
Hypothesis(
ys=[blank_id] * context_size,
log_prob=torch.zeros(1, dtype=torch.float32, device=device),
- state=init_states, # state of the RNNLM
+ state=init_states, # state of the NN LM
lm_score=init_score.reshape(-1),
state_cost=NgramLmStateCost(
LODR_lm
@@ -2170,7 +1930,6 @@ def modified_beam_search_rnnlm_LODR(
)
)
- rnnlm.clean_cache()
encoder_out = model.joiner.encoder_proj(packed_encoder_out.data)
offset = 0
@@ -2236,7 +1995,7 @@ def modified_beam_search_rnnlm_LODR(
It is a little confusing here because this for-loop
looks very similar to the one below. Here, we go through all
top-k tokens and only add the non-blanks ones to the token_list.
- The RNNLM will score those tokens given the LM states. Note that
+ LM will score those tokens given the LM states. Note that
the variable `scores` is the LM score after seeing the new
non-blank token.
"""
@@ -2256,21 +2015,41 @@ def modified_beam_search_rnnlm_LODR(
new_token = topk_token_indexes[k]
if new_token not in (blank_id, unk_id):
- assert new_token != 0, new_token
- token_list.append([new_token])
- # store the LSTM states
- hs.append(hyp.state[0])
- cs.append(hyp.state[1])
+ if LM.lm_type == "rnn":
+ token_list.append([new_token])
+ # store the LSTM states
+ hs.append(hyp.state[0])
+ cs.append(hyp.state[1])
+ else:
+ # for transformer LM
+ token_list.append(
+ [sos_id] + hyp.ys[context_size:] + [new_token]
+ )
- # forward RNNLM to get new states and scores
+ # forward NN LM to get new states and scores
if len(token_list) != 0:
- tokens_to_score = (
- torch.tensor(token_list).to(torch.int64).to(device).reshape(-1, 1)
- )
+ x_lens = torch.tensor([len(tokens) for tokens in token_list]).to(device)
+ if LM.lm_type == "rnn":
+ tokens_to_score = (
+ torch.tensor(token_list).to(torch.int64).to(device).reshape(-1, 1)
+ )
+ hs = torch.cat(hs, dim=1).to(device)
+ cs = torch.cat(cs, dim=1).to(device)
+ state = (hs, cs)
+ else:
+ # for transformer LM
+ tokens_list = [torch.tensor(tokens) for tokens in token_list]
+ tokens_to_score = (
+ torch.nn.utils.rnn.pad_sequence(
+ tokens_list, batch_first=True, padding_value=0.0
+ )
+ .to(device)
+ .to(torch.int64)
+ )
- hs = torch.cat(hs, dim=1).to(device)
- cs = torch.cat(cs, dim=1).to(device)
- scores, lm_states = rnnlm.score_token(tokens_to_score, (hs, cs))
+ state = None
+
+ scores, lm_states = LM.score_token(tokens_to_score, x_lens, state)
count = 0 # index, used to locate score and lm states
for i in range(batch_size):
@@ -2305,18 +2084,19 @@ def modified_beam_search_rnnlm_LODR(
state_cost.lm_score,
hyp.state_cost.lm_score,
)
- # score = score + RNNLM_score - LODR_score
- # LODR_LM_scale is a negative number here
+ # score = score + TDLM_score - LODR_score
+ # LODR_LM_scale should be a negative number here
hyp_log_prob += (
lm_score[new_token] * lm_scale
+ LODR_lm_scale * current_ngram_score
) # add the lm score
lm_score = scores[count]
- state = (
- lm_states[0][:, count, :].unsqueeze(1),
- lm_states[1][:, count, :].unsqueeze(1),
- )
+ if LM.lm_type == "rnn":
+ state = (
+ lm_states[0][:, count, :].unsqueeze(1),
+ lm_states[1][:, count, :].unsqueeze(1),
+ )
count += 1
else:
state_cost = hyp.state_cost
@@ -2340,3 +2120,263 @@ def modified_beam_search_rnnlm_LODR(
ans.append(sorted_ans[unsorted_indices[i]])
return ans
+
+
+def modified_beam_search_lm_shallow_fusion(
+ model: Transducer,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ sp: spm.SentencePieceProcessor,
+ LM: LmScorer,
+ beam: int = 4,
+ return_timestamps: bool = False,
+) -> List[List[int]]:
+ """Modified_beam_search + NN LM shallow fusion
+
+ Args:
+ model (Transducer):
+ The transducer model
+ encoder_out (torch.Tensor):
+ Encoder output in (N,T,C)
+ encoder_out_lens (torch.Tensor):
+ A 1-D tensor of shape (N,), containing the number of
+ valid frames in encoder_out before padding.
+ sp:
+ Sentence piece generator.
+ LM (LmScorer):
+ A neural net LM, e.g RNN or Transformer
+ beam (int, optional):
+ Beam size. Defaults to 4.
+
+ Returns:
+ Return a list-of-list of token IDs. ans[i] is the decoding results
+ for the i-th utterance.
+ """
+ assert encoder_out.ndim == 3, encoder_out.shape
+ assert encoder_out.size(0) >= 1, encoder_out.size(0)
+ assert LM is not None
+ lm_scale = LM.lm_scale
+
+ packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
+ input=encoder_out,
+ lengths=encoder_out_lens.cpu(),
+ batch_first=True,
+ enforce_sorted=False,
+ )
+
+ blank_id = model.decoder.blank_id
+ sos_id = sp.piece_to_id("")
+ unk_id = getattr(model, "unk_id", blank_id)
+ context_size = model.decoder.context_size
+ device = next(model.parameters()).device
+
+ batch_size_list = packed_encoder_out.batch_sizes.tolist()
+ N = encoder_out.size(0)
+ assert torch.all(encoder_out_lens > 0), encoder_out_lens
+ assert N == batch_size_list[0], (N, batch_size_list)
+
+ # get initial lm score and lm state by scoring the "sos" token
+ sos_token = torch.tensor([[sos_id]]).to(torch.int64).to(device)
+ lens = torch.tensor([1]).to(device)
+ init_score, init_states = LM.score_token(sos_token, lens)
+
+ B = [HypothesisList() for _ in range(N)]
+ for i in range(N):
+ B[i].add(
+ Hypothesis(
+ ys=[blank_id] * context_size,
+ log_prob=torch.zeros(1, dtype=torch.float32, device=device),
+ state=init_states,
+ lm_score=init_score.reshape(-1),
+ timestamp=[],
+ )
+ )
+
+ encoder_out = model.joiner.encoder_proj(packed_encoder_out.data)
+
+ offset = 0
+ finalized_B = []
+ for (t, batch_size) in enumerate(batch_size_list):
+ start = offset
+ end = offset + batch_size
+ current_encoder_out = encoder_out.data[start:end] # get batch
+ current_encoder_out = current_encoder_out.unsqueeze(1).unsqueeze(1)
+ # current_encoder_out's shape is (batch_size, 1, 1, encoder_out_dim)
+ offset = end
+
+ finalized_B = B[batch_size:] + finalized_B
+ B = B[:batch_size]
+
+ hyps_shape = get_hyps_shape(B).to(device)
+
+ A = [list(b) for b in B]
+ B = [HypothesisList() for _ in range(batch_size)]
+
+ ys_log_probs = torch.cat(
+ [hyp.log_prob.reshape(1, 1) for hyps in A for hyp in hyps]
+ )
+
+ lm_scores = torch.cat(
+ [hyp.lm_score.reshape(1, -1) for hyps in A for hyp in hyps]
+ )
+
+ decoder_input = torch.tensor(
+ [hyp.ys[-context_size:] for hyps in A for hyp in hyps],
+ device=device,
+ dtype=torch.int64,
+ ) # (num_hyps, context_size)
+
+ decoder_out = model.decoder(decoder_input, need_pad=False).unsqueeze(1)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+
+ current_encoder_out = torch.index_select(
+ current_encoder_out,
+ dim=0,
+ index=hyps_shape.row_ids(1).to(torch.int64),
+ ) # (num_hyps, 1, 1, encoder_out_dim)
+
+ logits = model.joiner(
+ current_encoder_out,
+ decoder_out,
+ project_input=False,
+ ) # (num_hyps, 1, 1, vocab_size)
+
+ logits = logits.squeeze(1).squeeze(1) # (num_hyps, vocab_size)
+
+ log_probs = logits.log_softmax(dim=-1) # (num_hyps, vocab_size)
+
+ log_probs.add_(ys_log_probs)
+
+ vocab_size = log_probs.size(-1)
+
+ log_probs = log_probs.reshape(-1)
+
+ row_splits = hyps_shape.row_splits(1) * vocab_size
+ log_probs_shape = k2.ragged.create_ragged_shape2(
+ row_splits=row_splits, cached_tot_size=log_probs.numel()
+ )
+ ragged_log_probs = k2.RaggedTensor(shape=log_probs_shape, value=log_probs)
+ """
+ for all hyps with a non-blank new token, score this token.
+ It is a little confusing here because this for-loop
+ looks very similar to the one below. Here, we go through all
+ top-k tokens and only add the non-blanks ones to the token_list.
+ `LM` will score those tokens given the LM states. Note that
+ the variable `scores` is the LM score after seeing the new
+ non-blank token.
+ """
+ token_list = [] # a list of list
+ hs = []
+ cs = []
+ for i in range(batch_size):
+ topk_log_probs, topk_indexes = ragged_log_probs[i].topk(beam)
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ topk_hyp_indexes = (topk_indexes // vocab_size).tolist()
+ topk_token_indexes = (topk_indexes % vocab_size).tolist()
+ for k in range(len(topk_hyp_indexes)):
+ hyp_idx = topk_hyp_indexes[k]
+ hyp = A[i][hyp_idx]
+
+ new_token = topk_token_indexes[k]
+ if new_token not in (blank_id, unk_id):
+ if LM.lm_type == "rnn":
+ token_list.append([new_token])
+ # store the LSTM states
+ hs.append(hyp.state[0])
+ cs.append(hyp.state[1])
+ else:
+ # for transformer LM
+ token_list.append(
+ [sos_id] + hyp.ys[context_size:] + [new_token]
+ )
+
+ if len(token_list) != 0:
+ x_lens = torch.tensor([len(tokens) for tokens in token_list]).to(device)
+ if LM.lm_type == "rnn":
+ tokens_to_score = (
+ torch.tensor(token_list).to(torch.int64).to(device).reshape(-1, 1)
+ )
+ hs = torch.cat(hs, dim=1).to(device)
+ cs = torch.cat(cs, dim=1).to(device)
+ state = (hs, cs)
+ else:
+ # for transformer LM
+ tokens_list = [torch.tensor(tokens) for tokens in token_list]
+ tokens_to_score = (
+ torch.nn.utils.rnn.pad_sequence(
+ tokens_list, batch_first=True, padding_value=0.0
+ )
+ .to(device)
+ .to(torch.int64)
+ )
+
+ state = None
+
+ scores, lm_states = LM.score_token(tokens_to_score, x_lens, state)
+
+ count = 0 # index, used to locate score and lm states
+ for i in range(batch_size):
+ topk_log_probs, topk_indexes = ragged_log_probs[i].topk(beam)
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ topk_hyp_indexes = (topk_indexes // vocab_size).tolist()
+ topk_token_indexes = (topk_indexes % vocab_size).tolist()
+
+ for k in range(len(topk_hyp_indexes)):
+ hyp_idx = topk_hyp_indexes[k]
+ hyp = A[i][hyp_idx]
+
+ ys = hyp.ys[:]
+
+ lm_score = hyp.lm_score
+ state = hyp.state
+
+ hyp_log_prob = topk_log_probs[k] # get score of current hyp
+ new_token = topk_token_indexes[k]
+ new_timestamp = hyp.timestamp[:]
+ if new_token not in (blank_id, unk_id):
+
+ ys.append(new_token)
+ new_timestamp.append(t)
+
+ hyp_log_prob += lm_score[new_token] * lm_scale # add the lm score
+
+ lm_score = scores[count]
+ if LM.lm_type == "rnn":
+ state = (
+ lm_states[0][:, count, :].unsqueeze(1),
+ lm_states[1][:, count, :].unsqueeze(1),
+ )
+ count += 1
+
+ new_hyp = Hypothesis(
+ ys=ys,
+ log_prob=hyp_log_prob,
+ state=state,
+ lm_score=lm_score,
+ timestamp=new_timestamp,
+ )
+ B[i].add(new_hyp)
+
+ B = B + finalized_B
+ best_hyps = [b.get_most_probable(length_norm=True) for b in B]
+
+ sorted_ans = [h.ys[context_size:] for h in best_hyps]
+ sorted_timestamps = [h.timestamp for h in best_hyps]
+ ans = []
+ ans_timestamps = []
+ unsorted_indices = packed_encoder_out.unsorted_indices.tolist()
+ for i in range(N):
+ ans.append(sorted_ans[unsorted_indices[i]])
+ ans_timestamps.append(sorted_timestamps[unsorted_indices[i]])
+
+ if not return_timestamps:
+ return ans
+ else:
+ return DecodingResults(
+ tokens=ans,
+ timestamps=ans_timestamps,
+ )
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless2/scaling.py b/egs/librispeech/ASR/pruned_transducer_stateless2/scaling.py
index c802ecf89..963ebdc2d 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless2/scaling.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless2/scaling.py
@@ -652,16 +652,16 @@ class ActivationBalancer(torch.nn.Module):
def forward(self, x: Tensor) -> Tensor:
if random.random() >= self.balance_prob:
return x
- else:
- return ActivationBalancerFunction.apply(
- x,
- self.channel_dim,
- self.min_positive,
- self.max_positive,
- self.max_factor / self.balance_prob,
- self.min_abs,
- self.max_abs,
- )
+
+ return ActivationBalancerFunction.apply(
+ x,
+ self.channel_dim,
+ self.min_positive,
+ self.max_positive,
+ self.max_factor / self.balance_prob,
+ self.min_abs,
+ self.max_abs,
+ )
class DoubleSwishFunction(torch.autograd.Function):
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless3/decode.py b/egs/librispeech/ASR/pruned_transducer_stateless3/decode.py
index e00aab34a..109a94a69 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless3/decode.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless3/decode.py
@@ -92,36 +92,37 @@ Usage:
--max-contexts 8 \
--max-states 64
-(8) modified beam search (with RNNLM shallow fusion)
+(8) modified beam search (with LM shallow fusion)
./pruned_transducer_stateless3/decode.py \
--epoch 28 \
--avg 15 \
--exp-dir ./pruned_transducer_stateless3/exp \
--max-duration 600 \
- --decoding-method modified_beam_search_rnnlm_shallow_fusion \
- --beam 4 \
- --rnn-lm-scale 0.3 \
- --rnn-lm-exp-dir /path/to/RNNLM \
+ --decoding-method modified_beam_search_lm_shallow_fusion \
+ --beam-size 4 \
+ --lm-type rnn \
+ --lm-scale 0.3 \
+ --lm-exp-dir /path/to/LM \
--rnn-lm-epoch 99 \
--rnn-lm-avg 1 \
--rnn-lm-num-layers 3 \
--rnn-lm-tie-weights 1
-(9) modified beam search with RNNLM shallow fusion + LODR
+(9) modified beam search with LM shallow fusion + LODR
./pruned_transducer_stateless3/decode.py \
--epoch 28 \
--avg 15 \
--max-duration 600 \
--exp-dir ./pruned_transducer_stateless3/exp \
- --decoding-method modified_beam_search_rnnlm_LODR \
- --beam 4 \
- --max-contexts 4 \
- --rnn-lm-scale 0.4 \
- --rnn-lm-exp-dir /path/to/RNNLM/exp \
+ --decoding-method modified_beam_search_LODR \
+ --beam-size 4 \
+ --lm-type rnn \
+ --lm-scale 0.4 \
+ --lm-exp-dir /path/to/LM \
--rnn-lm-epoch 99 \
--rnn-lm-avg 1 \
--rnn-lm-num-layers 3 \
- --rnn-lm-tie-weights 1 \
+ --rnn-lm-tie-weights 1
--tokens-ngram 2 \
--ngram-lm-scale -0.16 \
"""
@@ -149,14 +150,14 @@ from beam_search import (
greedy_search,
greedy_search_batch,
modified_beam_search,
+ modified_beam_search_lm_shallow_fusion,
+ modified_beam_search_LODR,
modified_beam_search_ngram_rescoring,
- modified_beam_search_rnnlm_LODR,
- modified_beam_search_rnnlm_shallow_fusion,
)
from librispeech import LibriSpeech
from train import add_model_arguments, get_params, get_transducer_model
-from icefall import NgramLm
+from icefall import LmScorer, NgramLm
from icefall.checkpoint import average_checkpoints, find_checkpoints, load_checkpoint
from icefall.lexicon import Lexicon
from icefall.rnn_lm.model import RnnLmModel
@@ -240,8 +241,8 @@ def get_parser():
- fast_beam_search_nbest_oracle
- fast_beam_search_nbest_LG
- modified_beam_search_ngram_rescoring
- - modified_beam_search_rnnlm_shallow_fusion
- - modified_beam_search_rnnlm_LODR
+ - modified_beam_search_lm_shallow_fusion
+ - modified_beam_search_LODR
If you use fast_beam_search_nbest_LG, you have to specify
`--lang-dir`, which should contain `LG.pt`.
""",
@@ -392,58 +393,28 @@ def get_parser():
)
parser.add_argument(
- "--rnn-lm-exp-dir",
- type=str,
- default="rnn_lm/exp",
- help="""Used only when --method is rnn-lm.
- It specifies the path to RNN LM exp dir.
- """,
- )
-
- parser.add_argument(
- "--rnn-lm-epoch",
- type=int,
- default=7,
- help="""Used only when --method is rnn-lm.
- It specifies the checkpoint to use.
- """,
- )
-
- parser.add_argument(
- "--rnn-lm-avg",
- type=int,
- default=2,
- help="""Used only when --method is rnn-lm.
- It specifies the number of checkpoints to average.
- """,
- )
-
- parser.add_argument(
- "--rnn-lm-embedding-dim",
- type=int,
- default=2048,
- help="Embedding dim of the model",
- )
-
- parser.add_argument(
- "--rnn-lm-hidden-dim",
- type=int,
- default=2048,
- help="Hidden dim of the model",
- )
-
- parser.add_argument(
- "--rnn-lm-num-layers",
- type=int,
- default=4,
- help="Number of RNN layers the model",
- )
- parser.add_argument(
- "--rnn-lm-tie-weights",
+ "--use-shallow-fusion",
type=str2bool,
- default=True,
- help="""True to share the weights between the input embedding layer and the
- last output linear layer
+ default=False,
+ help="""Use neural network LM for shallow fusion.
+ If you want to use LODR, you will also need to set this to true
+ """,
+ )
+
+ parser.add_argument(
+ "--lm-type",
+ type=str,
+ default="rnn",
+ help="Type of NN lm",
+ choices=["rnn", "transformer"],
+ )
+
+ parser.add_argument(
+ "--lm-scale",
+ type=float,
+ default=0.3,
+ help="""The scale of the neural network LM
+ Used only when `--use-shallow-fusion` is set to True.
""",
)
@@ -481,7 +452,7 @@ def decode_one_batch(
ngram_lm: Optional[NgramLm] = None,
ngram_lm_scale: float = 1.0,
rnn_lm_model: Optional[RnnLmModel] = None,
- rnnlm_scale: float = 1.0,
+ LM: Optional[LmScorer] = None,
) -> Dict[str, List[List[str]]]:
"""Decode one batch and return the result in a dict. The dict has the
following format:
@@ -515,10 +486,9 @@ def decode_one_batch(
fast_beam_search_nbest, fast_beam_search_nbest_oracle,
or fast_beam_search_with_nbest_rescoring.
It an FsaVec containing an acceptor.
- rnn_lm_model:
- A rnnlm which can be used for rescoring or shallow fusion
- rnnlm_scale:
- The scale of the rnnlm.
+ LM:
+ A neural net LM for shallow fusion. Only used when `--use-shallow-fusion`
+ set to true.
ngram_lm:
A ngram lm. Used in LODR decoding.
ngram_lm_scale:
@@ -697,20 +667,19 @@ def decode_one_batch(
)
for hyp in sp.decode(hyp_tokens):
hyps.append(hyp.split())
- elif params.decoding_method == "modified_beam_search_rnnlm_shallow_fusion":
- hyp_tokens = modified_beam_search_rnnlm_shallow_fusion(
+ elif params.decoding_method == "modified_beam_search_lm_shallow_fusion":
+ hyp_tokens = modified_beam_search_lm_shallow_fusion(
model=model,
encoder_out=encoder_out,
encoder_out_lens=encoder_out_lens,
beam=params.beam_size,
sp=sp,
- rnnlm=rnn_lm_model,
- rnnlm_scale=rnnlm_scale,
+ LM=LM,
)
for hyp in sp.decode(hyp_tokens):
hyps.append(hyp.split())
- elif params.decoding_method == "modified_beam_search_rnnlm_LODR":
- hyp_tokens = modified_beam_search_rnnlm_LODR(
+ elif params.decoding_method == "modified_beam_search_LODR":
+ hyp_tokens = modified_beam_search_LODR(
model=model,
encoder_out=encoder_out,
encoder_out_lens=encoder_out_lens,
@@ -718,8 +687,7 @@ def decode_one_batch(
sp=sp,
LODR_lm=ngram_lm,
LODR_lm_scale=ngram_lm_scale,
- rnnlm=rnn_lm_model,
- rnnlm_scale=rnnlm_scale,
+ LM=LM,
)
for hyp in sp.decode(hyp_tokens):
hyps.append(hyp.split())
@@ -812,7 +780,7 @@ def decode_dataset(
ngram_lm: Optional[NgramLm] = None,
ngram_lm_scale: float = 1.0,
rnn_lm_model: Optional[RnnLmModel] = None,
- rnnlm_scale: float = 1.0,
+ LM: Optional[LmScorer] = None,
) -> Dict[str, List[Tuple[List[str], List[str]]]]:
"""Decode dataset.
@@ -836,6 +804,8 @@ def decode_dataset(
fast_beam_search_nbest, fast_beam_search_nbest_oracle,
or fast_beam_search_with_nbest_rescoring.
It's an FsaVec containing an acceptor.
+ LM:
+ A neural network LM, used during shallow fusion
Returns:
Return a dict, whose key may be "greedy_search" if greedy search
is used, or it may be "beam_7" if beam size of 7 is used.
@@ -871,7 +841,7 @@ def decode_dataset(
ngram_lm=ngram_lm,
ngram_lm_scale=ngram_lm_scale,
rnn_lm_model=rnn_lm_model,
- rnnlm_scale=rnnlm_scale,
+ LM=LM,
)
for name, hyps in hyps_dict.items():
@@ -1005,6 +975,7 @@ def load_ngram_LM(
def main():
parser = get_parser()
AsrDataModule.add_arguments(parser)
+ LmScorer.add_arguments(parser)
args = parser.parse_args()
args.exp_dir = Path(args.exp_dir)
@@ -1022,9 +993,9 @@ def main():
"modified_beam_search",
"fast_beam_search_with_nbest_rescoring",
"fast_beam_search_with_nbest_rnn_rescoring",
- "modified_beam_search_rnnlm_LODR",
+ "modified_beam_search_LODR",
+ "modified_beam_search_lm_shallow_fusion",
"modified_beam_search_ngram_rescoring",
- "modified_beam_search_rnnlm_shallow_fusion",
)
params.res_dir = params.exp_dir / params.decoding_method
@@ -1055,12 +1026,18 @@ def main():
params.suffix += f"-max-sym-per-frame-{params.max_sym_per_frame}"
params.suffix += f"-temperature-{params.temperature}"
- if "rnnlm" in params.decoding_method:
- params.suffix += f"-rnnlm-lm-scale-{params.rnn_lm_scale}"
- if "LODR" in params.decoding_method:
- params.suffix += "-LODR"
if "ngram" in params.decoding_method:
params.suffix += f"-ngram-lm-scale-{params.ngram_lm_scale}"
+ if params.use_shallow_fusion:
+ if params.lm_type == "rnn":
+ params.suffix += f"-rnnlm-lm-scale-{params.lm_scale}"
+ elif params.lm_type == "transformer":
+ params.suffix += f"-transformer-lm-scale-{params.lm_scale}"
+
+ if "LODR" in params.decoding_method:
+ params.suffix += (
+ f"-LODR-{params.tokens_ngram}gram-scale-{params.ngram_lm_scale}"
+ )
setup_logger(f"{params.res_dir}/log-decode-{params.suffix}")
logging.info("Decoding started")
@@ -1195,28 +1172,19 @@ def main():
ngram_lm = None
ngram_lm_scale = None
- # only load rnnlm if used
- if "rnnlm" in params.decoding_method:
- rnn_lm_scale = params.rnn_lm_scale
-
- rnn_lm_model = RnnLmModel(
- vocab_size=params.vocab_size,
- embedding_dim=params.rnn_lm_embedding_dim,
- hidden_dim=params.rnn_lm_hidden_dim,
- num_layers=params.rnn_lm_num_layers,
- tie_weights=params.rnn_lm_tie_weights,
+ # only load the neural network LM if doing shallow fusion
+ if params.use_shallow_fusion:
+ LM = LmScorer(
+ lm_type=params.lm_type,
+ params=params,
+ device=device,
+ lm_scale=params.lm_scale,
)
- assert params.rnn_lm_avg == 1
+ LM.to(device)
+ LM.eval()
- load_checkpoint(
- f"{params.rnn_lm_exp_dir}/epoch-{params.rnn_lm_epoch}.pt",
- rnn_lm_model,
- )
- rnn_lm_model.to(device)
- rnn_lm_model.eval()
else:
- rnn_lm_model = None
- rnn_lm_scale = 0.0
+ LM = None
num_param = sum([p.numel() for p in model.parameters()])
logging.info(f"Number of model parameters: {num_param}")
@@ -1247,7 +1215,7 @@ def main():
ngram_lm=ngram_lm,
ngram_lm_scale=ngram_lm_scale,
rnn_lm_model=rnn_lm_model,
- rnnlm_scale=rnn_lm_scale,
+ LM=LM,
)
save_results(
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless3/librispeech.py b/egs/librispeech/ASR/pruned_transducer_stateless3/librispeech.py
index 6dba8e9fe..9f2cb6225 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless3/librispeech.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless3/librispeech.py
@@ -72,3 +72,12 @@ class LibriSpeech:
f = self.manifest_dir / "librispeech_cuts_dev-other.jsonl.gz"
logging.info(f"About to get dev-other cuts from {f}")
return load_manifest_lazy(f)
+
+ def train_all_shuf_cuts(self) -> CutSet:
+ logging.info(
+ "About to get the shuffled train-clean-100, \
+ train-clean-360 and train-other-500 cuts"
+ )
+ return load_manifest_lazy(
+ self.manifest_dir / "librispeech_cuts_train-all-shuf.jsonl.gz"
+ )
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless3/scaling_converter.py b/egs/librispeech/ASR/pruned_transducer_stateless3/scaling_converter.py
index b712eeda0..a6540c584 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless3/scaling_converter.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless3/scaling_converter.py
@@ -282,7 +282,7 @@ def convert_scaled_to_non_scaled(
if not inplace:
model = copy.deepcopy(model)
- excluded_patterns = r"self_attn\.(in|out)_proj"
+ excluded_patterns = r"(self|src)_attn\.(in|out)_proj"
p = re.compile(excluded_patterns)
d = {}
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless5/decode.py b/egs/librispeech/ASR/pruned_transducer_stateless5/decode.py
index 8b993f638..90b0fcf4b 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless5/decode.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless5/decode.py
@@ -87,22 +87,39 @@ Usage:
--max-contexts 8 \
--max-states 64
-(8) modified beam search with RNNLM shallow fusion (with LG)
+(8) modified beam search with RNNLM shallow fusion
./pruned_transducer_stateless5/decode.py \
--epoch 35 \
--avg 15 \
--exp-dir ./pruned_transducer_stateless5/exp \
--max-duration 600 \
- --decoding-method fast_beam_search_nbest_LG \
- --beam 4 \
- --max-contexts 4 \
- --rnn-lm-scale 0.4 \
- --rnn-lm-exp-dir /path/to/RNNLM/exp \
+ --decoding-method modified_beam_search_lm_shallow_fusion \
+ --beam-size 4 \
+ --lm-type rnn \
+ --lm-scale 0.3 \
+ --lm-exp-dir /path/to/LM \
--rnn-lm-epoch 99 \
--rnn-lm-avg 1 \
--rnn-lm-num-layers 3 \
--rnn-lm-tie-weights 1
+(9) modified beam search with LM shallow fusion + LODR
+./pruned_transducer_stateless5/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --max-duration 600 \
+ --exp-dir ./pruned_transducer_stateless5/exp \
+ --decoding-method modified_beam_search_LODR \
+ --beam-size 4 \
+ --lm-type rnn \
+ --lm-scale 0.4 \
+ --lm-exp-dir /path/to/LM \
+ --rnn-lm-epoch 99 \
+ --rnn-lm-avg 1 \
+ --rnn-lm-num-layers 3 \
+ --rnn-lm-tie-weights 1
+ --tokens-ngram 2 \
+ --ngram-lm-scale -0.16 \
"""
@@ -128,10 +145,13 @@ from beam_search import (
greedy_search,
greedy_search_batch,
modified_beam_search,
- modified_beam_search_rnnlm_shallow_fusion,
+ modified_beam_search_lm_shallow_fusion,
+ modified_beam_search_LODR,
+ modified_beam_search_ngram_rescoring,
)
from train import add_model_arguments, get_params, get_transducer_model
+from icefall import LmScorer, NgramLm
from icefall.checkpoint import (
average_checkpoints,
average_checkpoints_with_averaged_model,
@@ -139,7 +159,6 @@ from icefall.checkpoint import (
load_checkpoint,
)
from icefall.lexicon import Lexicon
-from icefall.rnn_lm.model import RnnLmModel
from icefall.utils import (
AttributeDict,
setup_logger,
@@ -229,7 +248,8 @@ def get_parser():
- fast_beam_search_nbest
- fast_beam_search_nbest_oracle
- fast_beam_search_nbest_LG
- - modified_beam_search_rnnlm_shallow_fusion # for rnn lm shallow fusion
+ - modified_beam_search_lm_shallow_fusion # for rnn lm shallow fusion
+ - modified_beam_search_LODR
If you use fast_beam_search_nbest_LG, you have to specify
`--lang-dir`, which should contain `LG.pt`.
""",
@@ -342,69 +362,49 @@ def get_parser():
)
parser.add_argument(
- "--rnn-lm-scale",
- type=float,
- default=0.0,
- help="""Used only when --method is modified_beam_search_rnnlm_shallow_fusion.
- It specifies the path to RNN LM exp dir.
- """,
- )
-
- parser.add_argument(
- "--rnn-lm-exp-dir",
- type=str,
- default="rnn_lm/exp",
- help="""Used only when --method is modified_beam_search_rnnlm_shallow_fusion.
- It specifies the path to RNN LM exp dir.
- """,
- )
-
- parser.add_argument(
- "--rnn-lm-epoch",
- type=int,
- default=7,
- help="""Used only when --method is modified_beam_search_rnnlm_shallow_fusion.
- It specifies the checkpoint to use.
- """,
- )
-
- parser.add_argument(
- "--rnn-lm-avg",
- type=int,
- default=2,
- help="""Used only when --method is modified_beam_search_rnnlm_shallow_fusion.
- It specifies the number of checkpoints to average.
- """,
- )
-
- parser.add_argument(
- "--rnn-lm-embedding-dim",
- type=int,
- default=2048,
- help="Embedding dim of the model",
- )
-
- parser.add_argument(
- "--rnn-lm-hidden-dim",
- type=int,
- default=2048,
- help="Hidden dim of the model",
- )
-
- parser.add_argument(
- "--rnn-lm-num-layers",
- type=int,
- default=4,
- help="Number of RNN layers the model",
- )
- parser.add_argument(
- "--rnn-lm-tie-weights",
+ "--use-shallow-fusion",
type=str2bool,
default=False,
- help="""True to share the weights between the input embedding layer and the
- last output linear layer
+ help="""Use neural network LM for shallow fusion.
+ If you want to use LODR, you will also need to set this to true
""",
)
+
+ parser.add_argument(
+ "--lm-type",
+ type=str,
+ default="rnn",
+ help="Type of NN lm",
+ choices=["rnn", "transformer"],
+ )
+
+ parser.add_argument(
+ "--lm-scale",
+ type=float,
+ default=0.3,
+ help="""The scale of the neural network LM
+ Used only when `--use-shallow-fusion` is set to True.
+ """,
+ )
+
+ parser.add_argument(
+ "--tokens-ngram",
+ type=int,
+ default=3,
+ help="""Token Ngram used for rescoring.
+ Used only when the decoding method is
+ modified_beam_search_ngram_rescoring, or LODR
+ """,
+ )
+
+ parser.add_argument(
+ "--backoff-id",
+ type=int,
+ default=500,
+ help="""ID of the backoff symbol.
+ Used only when the decoding method is
+ modified_beam_search_ngram_rescoring""",
+ )
add_model_arguments(parser)
return parser
@@ -417,8 +417,9 @@ def decode_one_batch(
batch: dict,
word_table: Optional[k2.SymbolTable] = None,
decoding_graph: Optional[k2.Fsa] = None,
- rnnlm: Optional[RnnLmModel] = None,
- rnnlm_scale: float = 1.0,
+ ngram_lm: Optional[NgramLm] = None,
+ ngram_lm_scale: float = 1.0,
+ LM: Optional[LmScorer] = None,
) -> Dict[str, List[List[str]]]:
"""Decode one batch and return the result in a dict. The dict has the
following format:
@@ -447,6 +448,13 @@ def decode_one_batch(
The decoding graph. Can be either a `k2.trivial_graph` or LG, Used
only when --decoding_method is fast_beam_search, fast_beam_search_LG, fast_beam_search_nbest,
fast_beam_search_nbest_oracle, and fast_beam_search_nbest_LG.
+ LM:
+ A neural net LM for shallow fusion. Only used when `--use-shallow-fusion`
+ set to true.
+ ngram_lm:
+ A ngram lm. Used in LODR decoding.
+ ngram_lm_scale:
+ The scale of the ngram language model.
Returns:
Return the decoding result. See above description for the format of
the returned dict.
@@ -559,15 +567,38 @@ def decode_one_batch(
)
for hyp in sp.decode(hyp_tokens):
hyps.append(hyp.split())
- elif params.decoding_method == "modified_beam_search_rnnlm_shallow_fusion":
- hyp_tokens = modified_beam_search_rnnlm_shallow_fusion(
+ elif params.decoding_method == "modified_beam_search_ngram_rescoring":
+ hyp_tokens = modified_beam_search_ngram_rescoring(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ ngram_lm=ngram_lm,
+ ngram_lm_scale=ngram_lm_scale,
+ beam=params.beam_size,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.decoding_method == "modified_beam_search_lm_shallow_fusion":
+ hyp_tokens = modified_beam_search_lm_shallow_fusion(
model=model,
encoder_out=encoder_out,
encoder_out_lens=encoder_out_lens,
beam=params.beam_size,
sp=sp,
- rnnlm=rnnlm,
- rnnlm_scale=rnnlm_scale,
+ LM=LM,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.decoding_method == "modified_beam_search_LODR":
+ hyp_tokens = modified_beam_search_LODR(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam_size,
+ sp=sp,
+ LODR_lm=ngram_lm,
+ LODR_lm_scale=ngram_lm_scale,
+ LM=LM,
)
for hyp in sp.decode(hyp_tokens):
hyps.append(hyp.split())
@@ -620,8 +651,9 @@ def decode_dataset(
sp: spm.SentencePieceProcessor,
word_table: Optional[k2.SymbolTable] = None,
decoding_graph: Optional[k2.Fsa] = None,
- rnnlm: Optional[RnnLmModel] = None,
- rnnlm_scale: float = 1.0,
+ ngram_lm: Optional[NgramLm] = None,
+ ngram_lm_scale: float = 1.0,
+ LM: Optional[LmScorer] = None,
) -> Dict[str, List[Tuple[str, List[str], List[str]]]]:
"""Decode dataset.
@@ -640,6 +672,8 @@ def decode_dataset(
The decoding graph. Can be either a `k2.trivial_graph` or LG, Used
only when --decoding_method is fast_beam_search, fast_beam_search_LG, fast_beam_search_nbest,
fast_beam_search_nbest_oracle, and fast_beam_search_nbest_LG.
+ LM:
+ A neural network LM, used during shallow fusion
Returns:
Return a dict, whose key may be "greedy_search" if greedy search
is used, or it may be "beam_7" if beam size of 7 is used.
@@ -663,7 +697,6 @@ def decode_dataset(
for batch_idx, batch in enumerate(dl):
texts = batch["supervisions"]["text"]
cut_ids = [cut.id for cut in batch["supervisions"]["cut"]]
- logging.info(f"Decoding {batch_idx}-th batch")
hyps_dict = decode_one_batch(
params=params,
@@ -672,8 +705,9 @@ def decode_dataset(
decoding_graph=decoding_graph,
word_table=word_table,
batch=batch,
- rnnlm=rnnlm,
- rnnlm_scale=rnnlm_scale,
+ ngram_lm=ngram_lm,
+ ngram_lm_scale=ngram_lm_scale,
+ LM=LM,
)
for name, hyps in hyps_dict.items():
@@ -742,6 +776,7 @@ def save_results(
def main():
parser = get_parser()
LibriSpeechAsrDataModule.add_arguments(parser)
+ LmScorer.add_arguments(parser)
args = parser.parse_args()
args.exp_dir = Path(args.exp_dir)
@@ -757,7 +792,8 @@ def main():
"fast_beam_search_nbest_LG",
"fast_beam_search_nbest_oracle",
"modified_beam_search",
- "modified_beam_search_rnnlm_shallow_fusion",
+ "modified_beam_search_lm_shallow_fusion",
+ "modified_beam_search_LODR",
)
params.res_dir = params.exp_dir / params.decoding_method
@@ -783,7 +819,18 @@ def main():
params.suffix += f"-context-{params.context_size}"
params.suffix += f"-max-sym-per-frame-{params.max_sym_per_frame}"
- params.suffix += f"-rnnlm-lm-scale-{params.rnn_lm_scale}"
+ if "ngram" in params.decoding_method:
+ params.suffix += f"-ngram-lm-scale-{params.ngram_lm_scale}"
+ if params.use_shallow_fusion:
+ if params.lm_type == "rnn":
+ params.suffix += f"-rnnlm-lm-scale-{params.lm_scale}"
+ elif params.lm_type == "transformer":
+ params.suffix += f"-transformer-lm-scale-{params.lm_scale}"
+
+ if "LODR" in params.decoding_method:
+ params.suffix += (
+ f"-LODR-{params.tokens_ngram}gram-scale-{params.ngram_lm_scale}"
+ )
if params.use_averaged_model:
params.suffix += "-use-averaged-model"
@@ -895,24 +942,34 @@ def main():
model.to(device)
model.eval()
- rnn_lm_model = None
- rnn_lm_scale = params.rnn_lm_scale
- if params.decoding_method == "modified_beam_search_rnnlm_shallow_fusion":
- rnn_lm_model = RnnLmModel(
- vocab_size=params.vocab_size,
- embedding_dim=params.rnn_lm_embedding_dim,
- hidden_dim=params.rnn_lm_hidden_dim,
- num_layers=params.rnn_lm_num_layers,
- tie_weights=params.rnn_lm_tie_weights,
+ # only load N-gram LM when needed
+ if "ngram" in params.decoding_method or "LODR" in params.decoding_method:
+ lm_filename = f"{params.tokens_ngram}gram.fst.txt"
+ logging.info(f"lm filename: {lm_filename}")
+ ngram_lm = NgramLm(
+ str(params.lang_dir / lm_filename),
+ backoff_id=params.backoff_id,
+ is_binary=False,
)
- assert params.rnn_lm_avg == 1
+ logging.info(f"num states: {ngram_lm.lm.num_states}")
+ ngram_lm_scale = params.ngram_lm_scale
+ else:
+ ngram_lm = None
+ ngram_lm_scale = None
- load_checkpoint(
- f"{params.rnn_lm_exp_dir}/epoch-{params.rnn_lm_epoch}.pt",
- rnn_lm_model,
+ # only load the neural network LM if doing shallow fusion
+ if params.use_shallow_fusion:
+ LM = LmScorer(
+ lm_type=params.lm_type,
+ params=params,
+ device=device,
+ lm_scale=params.lm_scale,
)
- rnn_lm_model.to(device)
- rnn_lm_model.eval()
+ LM.to(device)
+ LM.eval()
+
+ else:
+ LM = None
if "fast_beam_search" in params.decoding_method:
if "LG" in params.decoding_method:
@@ -955,8 +1012,9 @@ def main():
sp=sp,
word_table=word_table,
decoding_graph=decoding_graph,
- rnnlm=rnn_lm_model,
- rnnlm_scale=rnn_lm_scale,
+ ngram_lm=ngram_lm,
+ ngram_lm_scale=ngram_lm_scale,
+ LM=LM,
)
save_results(
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless6/vq_utils.py b/egs/librispeech/ASR/pruned_transducer_stateless6/vq_utils.py
index 97a83b974..14ff86f23 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless6/vq_utils.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless6/vq_utils.py
@@ -68,7 +68,10 @@ class CodebookIndexExtractor:
def init_dirs(self):
# vq_dir is the root dir for quantization, containing:
# training data, trained quantizer, and extracted codebook indexes
- self.vq_dir = self.params.exp_dir / f"vq/{self.params.teacher_model_id}/"
+ self.vq_dir = (
+ self.params.exp_dir
+ / f"vq/{self.params.teacher_model_id}_layer{self.params.embedding_layer}_cb{self.params.num_codebooks}/"
+ )
self.vq_dir.mkdir(parents=True, exist_ok=True)
# manifest_dir contains:
@@ -79,7 +82,10 @@ class CodebookIndexExtractor:
# It's doesn't matter whether ori_manifest_dir is str or Path.
# Set it to Path to be consistent.
self.ori_manifest_dir = Path("./data/fbank/")
- self.dst_manifest_dir = Path("./data/vq_fbank/")
+ self.dst_manifest_dir = Path(
+ f"./data/vq_fbank_layer"
+ + f"{self.params.embedding_layer}_cb{self.params.num_codebooks}/"
+ )
self.dst_manifest_dir.mkdir(parents=True, exist_ok=True)
@@ -244,10 +250,36 @@ class CodebookIndexExtractor:
)
cuts_vq = load_manifest(vq_manifest_path)
cuts_ori = load_manifest(ori_manifest_path)
- cuts_vq = cuts_vq.sort_like(cuts_ori)
- for cut_idx, (cut_vq, cut_ori) in enumerate(zip(cuts_vq, cuts_ori)):
- assert cut_vq.id == cut_ori.id
- cut_ori.codebook_indexes = cut_vq.codebook_indexes
+ assert len(cuts_vq) == len(cuts_ori), "Cuts should have the same length!"
+
+ if set(cuts_vq.ids) == set(cuts_ori.ids):
+ # IDs match exactly
+ cuts_vq = cuts_vq.sort_like(cuts_ori)
+ for cut_idx, (cut_vq, cut_ori) in enumerate(zip(cuts_vq, cuts_ori)):
+ assert cut_vq.id == cut_ori.id, (cut_vq.id, cut_ori.id)
+ cut_ori.codebook_indexes = cut_vq.codebook_indexes
+ else:
+ # in case of ID mismatch, remap them
+ # get the mapping between audio and cut ID
+ logging
+ ori_id_map = {}
+ for id in cuts_ori.ids:
+ # some text normalization
+ if "sp" in id:
+ clean_id = "-".join(id.split("-")[:3]) + "_" + id.split("_")[-1]
+ else:
+ clean_id = "-".join(id.split("-")[:3])
+ ori_id_map[clean_id] = id
+
+ for id in cuts_vq.ids:
+ if "sp" in id:
+ clean_id = "-".join(id.split("-")[:3]) + "_" + id.split("_")[-1]
+ else:
+ clean_id = "-".join(id.split("-")[:3])
+ assert clean_id in ori_id_map, clean_id
+ cuts_ori[ori_id_map[clean_id]].codebook_indexes = cuts_vq[
+ id
+ ].codebook_indexes
CutSet.from_cuts(cuts_ori).to_jsonl(dst_vq_manifest_path)
logging.info(f"Processed {subset}.")
@@ -258,7 +290,10 @@ class CodebookIndexExtractor:
Merge generated vq included manfiests and storage to self.dst_manifest_dir.
"""
for subset in self.params.subsets:
- vq_manifests = f"{self.manifest_dir}/with_codebook_indexes-librispeech-cuts_train-{subset}*.jsonl.gz"
+ vq_manifests = (
+ f"{self.manifest_dir}/"
+ + f"with_codebook_indexes-librispeech-cuts_train-{subset}*.jsonl.gz"
+ )
dst_vq_manifest = (
self.dst_manifest_dir / f"librispeech_cuts_train-{subset}-vq.jsonl.gz"
)
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/decode.py b/egs/librispeech/ASR/pruned_transducer_stateless7/decode.py
index bc15948fc..b9bce465f 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless7/decode.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/decode.py
@@ -1,7 +1,8 @@
#!/usr/bin/env python3
#
# Copyright 2021-2022 Xiaomi Corporation (Author: Fangjun Kuang,
-# Zengwei Yao)
+# Zengwei Yao,
+# Xiaoyu Yang)
#
# See ../../../../LICENSE for clarification regarding multiple authors
#
@@ -91,6 +92,41 @@ Usage:
--beam 20.0 \
--max-contexts 8 \
--max-states 64
+
+(8) modified beam search with RNNLM shallow fusion
+./pruned_transducer_stateless5/decode.py \
+ --epoch 35 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless5/exp \
+ --max-duration 600 \
+ --decoding-method modified_beam_search_lm_shallow_fusion \
+ --beam-size 4 \
+ --lm-type rnn \
+ --lm-scale 0.3 \
+ --lm-exp-dir /path/to/LM \
+ --rnn-lm-epoch 99 \
+ --rnn-lm-avg 1 \
+ --rnn-lm-num-layers 3 \
+ --rnn-lm-tie-weights 1
+
+(9) modified beam search with LM shallow fusion + LODR
+./pruned_transducer_stateless5/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --max-duration 600 \
+ --exp-dir ./pruned_transducer_stateless5/exp \
+ --decoding-method modified_beam_search_LODR \
+ --beam-size 4 \
+ --lm-type rnn \
+ --lm-scale 0.4 \
+ --lm-exp-dir /path/to/LM \
+ --rnn-lm-epoch 99 \
+ --rnn-lm-avg 1 \
+ --rnn-lm-num-layers 3 \
+ --rnn-lm-tie-weights 1
+ --tokens-ngram 2 \
+ --ngram-lm-scale -0.16 \
+
"""
@@ -115,9 +151,13 @@ from beam_search import (
greedy_search,
greedy_search_batch,
modified_beam_search,
+ modified_beam_search_lm_shallow_fusion,
+ modified_beam_search_LODR,
+ modified_beam_search_ngram_rescoring,
)
from train import add_model_arguments, get_params, get_transducer_model
+from icefall import LmScorer, NgramLm
from icefall.checkpoint import (
average_checkpoints,
average_checkpoints_with_averaged_model,
@@ -213,6 +253,8 @@ def get_parser():
- fast_beam_search_nbest
- fast_beam_search_nbest_oracle
- fast_beam_search_nbest_LG
+ - modified_beam_search_lm_shallow_fusion # for rnn lm shallow fusion
+ - modified_beam_search_LODR
If you use fast_beam_search_nbest_LG, you have to specify
`--lang-dir`, which should contain `LG.pt`.
""",
@@ -274,6 +316,7 @@ def get_parser():
default=2,
help="The context size in the decoder. 1 means bigram; 2 means tri-gram",
)
+
parser.add_argument(
"--max-sym-per-frame",
type=int,
@@ -323,6 +366,50 @@ def get_parser():
help="left context can be seen during decoding (in frames after subsampling)",
)
+ parser.add_argument(
+ "--use-shallow-fusion",
+ type=str2bool,
+ default=False,
+ help="""Use neural network LM for shallow fusion.
+ If you want to use LODR, you will also need to set this to true
+ """,
+ )
+
+ parser.add_argument(
+ "--lm-type",
+ type=str,
+ default="rnn",
+ help="Type of NN lm",
+ choices=["rnn", "transformer"],
+ )
+
+ parser.add_argument(
+ "--lm-scale",
+ type=float,
+ default=0.3,
+ help="""The scale of the neural network LM
+ Used only when `--use-shallow-fusion` is set to True.
+ """,
+ )
+
+ parser.add_argument(
+ "--tokens-ngram",
+ type=int,
+ default=3,
+ help="""Token Ngram used for rescoring.
+ Used only when the decoding method is
+ modified_beam_search_ngram_rescoring, or LODR
+ """,
+ )
+
+ parser.add_argument(
+ "--backoff-id",
+ type=int,
+ default=500,
+ help="""ID of the backoff symbol.
+ Used only when the decoding method is
+ modified_beam_search_ngram_rescoring""",
+ )
add_model_arguments(parser)
return parser
@@ -335,6 +422,9 @@ def decode_one_batch(
batch: dict,
word_table: Optional[k2.SymbolTable] = None,
decoding_graph: Optional[k2.Fsa] = None,
+ ngram_lm: Optional[NgramLm] = None,
+ ngram_lm_scale: float = 1.0,
+ LM: Optional[LmScorer] = None,
) -> Dict[str, List[List[str]]]:
"""Decode one batch and return the result in a dict. The dict has the
following format:
@@ -363,6 +453,13 @@ def decode_one_batch(
The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
only when --decoding_method is fast_beam_search, fast_beam_search_nbest,
fast_beam_search_nbest_oracle, and fast_beam_search_nbest_LG.
+ LM:
+ A neural net LM for shallow fusion. Only used when `--use-shallow-fusion`
+ set to true.
+ ngram_lm:
+ A ngram lm. Used in LODR decoding.
+ ngram_lm_scale:
+ The scale of the ngram language model.
Returns:
Return the decoding result. See above description for the format of
the returned dict.
@@ -468,6 +565,30 @@ def decode_one_batch(
)
for hyp in sp.decode(hyp_tokens):
hyps.append(hyp.split())
+ elif params.decoding_method == "modified_beam_search_lm_shallow_fusion":
+ hyp_tokens = modified_beam_search_lm_shallow_fusion(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam_size,
+ sp=sp,
+ LM=LM,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.decoding_method == "modified_beam_search_LODR":
+ hyp_tokens = modified_beam_search_LODR(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam_size,
+ sp=sp,
+ LODR_lm=ngram_lm,
+ LODR_lm_scale=ngram_lm_scale,
+ LM=LM,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
else:
batch_size = encoder_out.size(0)
@@ -517,6 +638,9 @@ def decode_dataset(
sp: spm.SentencePieceProcessor,
word_table: Optional[k2.SymbolTable] = None,
decoding_graph: Optional[k2.Fsa] = None,
+ ngram_lm: Optional[NgramLm] = None,
+ ngram_lm_scale: float = 1.0,
+ LM: Optional[LmScorer] = None,
) -> Dict[str, List[Tuple[str, List[str], List[str]]]]:
"""Decode dataset.
@@ -535,6 +659,8 @@ def decode_dataset(
The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
only when --decoding_method is fast_beam_search, fast_beam_search_nbest,
fast_beam_search_nbest_oracle, and fast_beam_search_nbest_LG.
+ LM:
+ A neural network LM, used during shallow fusion
Returns:
Return a dict, whose key may be "greedy_search" if greedy search
is used, or it may be "beam_7" if beam size of 7 is used.
@@ -566,6 +692,9 @@ def decode_dataset(
decoding_graph=decoding_graph,
word_table=word_table,
batch=batch,
+ ngram_lm=ngram_lm,
+ ngram_lm_scale=ngram_lm_scale,
+ LM=LM,
)
for name, hyps in hyps_dict.items():
@@ -634,6 +763,7 @@ def save_results(
def main():
parser = get_parser()
LibriSpeechAsrDataModule.add_arguments(parser)
+ LmScorer.add_arguments(parser)
args = parser.parse_args()
args.exp_dir = Path(args.exp_dir)
@@ -648,6 +778,8 @@ def main():
"fast_beam_search_nbest_LG",
"fast_beam_search_nbest_oracle",
"modified_beam_search",
+ "modified_beam_search_lm_shallow_fusion",
+ "modified_beam_search_LODR",
)
params.res_dir = params.exp_dir / params.decoding_method
@@ -675,6 +807,19 @@ def main():
params.suffix += f"-context-{params.context_size}"
params.suffix += f"-max-sym-per-frame-{params.max_sym_per_frame}"
+ if "ngram" in params.decoding_method:
+ params.suffix += f"-ngram-lm-scale-{params.ngram_lm_scale}"
+ if params.use_shallow_fusion:
+ if params.lm_type == "rnn":
+ params.suffix += f"-rnnlm-lm-scale-{params.lm_scale}"
+ elif params.lm_type == "transformer":
+ params.suffix += f"-transformer-lm-scale-{params.lm_scale}"
+
+ if "LODR" in params.decoding_method:
+ params.suffix += (
+ f"-LODR-{params.tokens_ngram}gram-scale-{params.ngram_lm_scale}"
+ )
+
if params.use_averaged_model:
params.suffix += "-use-averaged-model"
@@ -785,6 +930,34 @@ def main():
model.to(device)
model.eval()
+ # only load N-gram LM when needed
+ if "ngram" in params.decoding_method or "LODR" in params.decoding_method:
+ lm_filename = f"{params.tokens_ngram}gram.fst.txt"
+ logging.info(f"lm filename: {lm_filename}")
+ ngram_lm = NgramLm(
+ str(params.lang_dir / lm_filename),
+ backoff_id=params.backoff_id,
+ is_binary=False,
+ )
+ logging.info(f"num states: {ngram_lm.lm.num_states}")
+ ngram_lm_scale = params.ngram_lm_scale
+ else:
+ ngram_lm = None
+ ngram_lm_scale = None
+
+ # only load the neural network LM if doing shallow fusion
+ if params.use_shallow_fusion:
+ LM = LmScorer(
+ lm_type=params.lm_type,
+ params=params,
+ device=device,
+ lm_scale=params.lm_scale,
+ )
+ LM.to(device)
+ LM.eval()
+
+ else:
+ LM = None
if "fast_beam_search" in params.decoding_method:
if params.decoding_method == "fast_beam_search_nbest_LG":
lexicon = Lexicon(params.lang_dir)
@@ -826,6 +999,9 @@ def main():
sp=sp,
word_table=word_table,
decoding_graph=decoding_graph,
+ ngram_lm=ngram_lm,
+ ngram_lm_scale=ngram_lm_scale,
+ LM=LM,
)
save_results(
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/export.py b/egs/librispeech/ASR/pruned_transducer_stateless7/export.py
index 9a6f3ed37..db8b5eb2b 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless7/export.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/export.py
@@ -41,7 +41,31 @@ Check
https://github.com/k2-fsa/sherpa
for how to use the exported models outside of icefall.
-(2) Export `model.state_dict()`
+(2) Export to ONNX format
+
+./pruned_transducer_stateless7/export.py \
+ --exp-dir ./pruned_transducer_stateless7/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 20 \
+ --avg 10 \
+ --onnx 1
+
+It will generate the following files in the given `exp_dir`.
+Check `onnx_check.py` for how to use them.
+
+ - encoder.onnx
+ - decoder.onnx
+ - joiner.onnx
+ - joiner_encoder_proj.onnx
+ - joiner_decoder_proj.onnx
+
+Please see ./onnx_pretrained.py for usage of the generated files
+
+Check
+https://github.com/k2-fsa/sherpa-onnx
+for how to use the exported models outside of icefall.
+
+(3) Export `model.state_dict()`
./pruned_transducer_stateless7/export.py \
--exp-dir ./pruned_transducer_stateless7/exp \
@@ -172,6 +196,23 @@ def get_parser():
""",
)
+ parser.add_argument(
+ "--onnx",
+ type=str2bool,
+ default=False,
+ help="""If True, --jit is ignored and it exports the model
+ to onnx format. It will generate the following files:
+
+ - encoder.onnx
+ - decoder.onnx
+ - joiner.onnx
+ - joiner_encoder_proj.onnx
+ - joiner_decoder_proj.onnx
+
+ Refer to ./onnx_check.py and ./onnx_pretrained.py for how to use them.
+ """,
+ )
+
parser.add_argument(
"--context-size",
type=int,
@@ -184,6 +225,204 @@ def get_parser():
return parser
+def export_encoder_model_onnx(
+ encoder_model: nn.Module,
+ encoder_filename: str,
+ opset_version: int = 11,
+) -> None:
+ """Export the given encoder model to ONNX format.
+ The exported model has two inputs:
+
+ - x, a tensor of shape (N, T, C); dtype is torch.float32
+ - x_lens, a tensor of shape (N,); dtype is torch.int64
+
+ and it has two outputs:
+
+ - encoder_out, a tensor of shape (N, T, C)
+ - encoder_out_lens, a tensor of shape (N,)
+
+ Note: The warmup argument is fixed to 1.
+
+ Args:
+ encoder_model:
+ The input encoder model
+ encoder_filename:
+ The filename to save the exported ONNX model.
+ opset_version:
+ The opset version to use.
+ """
+ x = torch.zeros(1, 101, 80, dtype=torch.float32)
+ x_lens = torch.tensor([101], dtype=torch.int64)
+
+ # encoder_model = torch.jit.script(encoder_model)
+ # It throws the following error for the above statement
+ #
+ # RuntimeError: Exporting the operator __is_ to ONNX opset version
+ # 11 is not supported. Please feel free to request support or
+ # submit a pull request on PyTorch GitHub.
+ #
+ # I cannot find which statement causes the above error.
+ # torch.onnx.export() will use torch.jit.trace() internally, which
+ # works well for the current reworked model
+ torch.onnx.export(
+ encoder_model,
+ (x, x_lens),
+ encoder_filename,
+ verbose=False,
+ opset_version=opset_version,
+ input_names=["x", "x_lens"],
+ output_names=["encoder_out", "encoder_out_lens"],
+ dynamic_axes={
+ "x": {0: "N", 1: "T"},
+ "x_lens": {0: "N"},
+ "encoder_out": {0: "N", 1: "T"},
+ "encoder_out_lens": {0: "N"},
+ },
+ )
+ logging.info(f"Saved to {encoder_filename}")
+
+
+def export_decoder_model_onnx(
+ decoder_model: nn.Module,
+ decoder_filename: str,
+ opset_version: int = 11,
+) -> None:
+ """Export the decoder model to ONNX format.
+
+ The exported model has one input:
+
+ - y: a torch.int64 tensor of shape (N, decoder_model.context_size)
+
+ and has one output:
+
+ - decoder_out: a torch.float32 tensor of shape (N, 1, C)
+
+ Note: The argument need_pad is fixed to False.
+
+ Args:
+ decoder_model:
+ The decoder model to be exported.
+ decoder_filename:
+ Filename to save the exported ONNX model.
+ opset_version:
+ The opset version to use.
+ """
+ y = torch.zeros(10, decoder_model.context_size, dtype=torch.int64)
+ need_pad = False # Always False, so we can use torch.jit.trace() here
+ # Note(fangjun): torch.jit.trace() is more efficient than torch.jit.script()
+ # in this case
+ torch.onnx.export(
+ decoder_model,
+ (y, need_pad),
+ decoder_filename,
+ verbose=False,
+ opset_version=opset_version,
+ input_names=["y", "need_pad"],
+ output_names=["decoder_out"],
+ dynamic_axes={
+ "y": {0: "N"},
+ "decoder_out": {0: "N"},
+ },
+ )
+ logging.info(f"Saved to {decoder_filename}")
+
+
+def export_joiner_model_onnx(
+ joiner_model: nn.Module,
+ joiner_filename: str,
+ opset_version: int = 11,
+) -> None:
+ """Export the joiner model to ONNX format.
+ The exported joiner model has two inputs:
+
+ - projected_encoder_out: a tensor of shape (N, joiner_dim)
+ - projected_decoder_out: a tensor of shape (N, joiner_dim)
+
+ and produces one output:
+
+ - logit: a tensor of shape (N, vocab_size)
+
+ The exported encoder_proj model has one input:
+
+ - encoder_out: a tensor of shape (N, encoder_out_dim)
+
+ and produces one output:
+
+ - projected_encoder_out: a tensor of shape (N, joiner_dim)
+
+ The exported decoder_proj model has one input:
+
+ - decoder_out: a tensor of shape (N, decoder_out_dim)
+
+ and produces one output:
+
+ - projected_decoder_out: a tensor of shape (N, joiner_dim)
+ """
+ encoder_proj_filename = str(joiner_filename).replace(".onnx", "_encoder_proj.onnx")
+ decoder_proj_filename = str(joiner_filename).replace(".onnx", "_decoder_proj.onnx")
+
+ encoder_out_dim = joiner_model.encoder_proj.weight.shape[1]
+ decoder_out_dim = joiner_model.decoder_proj.weight.shape[1]
+ joiner_dim = joiner_model.decoder_proj.weight.shape[0]
+
+ projected_encoder_out = torch.rand(1, 1, 1, joiner_dim, dtype=torch.float32)
+ projected_decoder_out = torch.rand(1, 1, 1, joiner_dim, dtype=torch.float32)
+
+ project_input = False
+ # Note: It uses torch.jit.trace() internally
+ torch.onnx.export(
+ joiner_model,
+ (projected_encoder_out, projected_decoder_out, project_input),
+ joiner_filename,
+ verbose=False,
+ opset_version=opset_version,
+ input_names=[
+ "encoder_out",
+ "decoder_out",
+ "project_input",
+ ],
+ output_names=["logit"],
+ dynamic_axes={
+ "encoder_out": {0: "N"},
+ "decoder_out": {0: "N"},
+ "logit": {0: "N"},
+ },
+ )
+ logging.info(f"Saved to {joiner_filename}")
+
+ encoder_out = torch.rand(1, encoder_out_dim, dtype=torch.float32)
+ torch.onnx.export(
+ joiner_model.encoder_proj,
+ encoder_out,
+ encoder_proj_filename,
+ verbose=False,
+ opset_version=opset_version,
+ input_names=["encoder_out"],
+ output_names=["projected_encoder_out"],
+ dynamic_axes={
+ "encoder_out": {0: "N"},
+ "projected_encoder_out": {0: "N"},
+ },
+ )
+ logging.info(f"Saved to {encoder_proj_filename}")
+
+ decoder_out = torch.rand(1, decoder_out_dim, dtype=torch.float32)
+ torch.onnx.export(
+ joiner_model.decoder_proj,
+ decoder_out,
+ decoder_proj_filename,
+ verbose=False,
+ opset_version=opset_version,
+ input_names=["decoder_out"],
+ output_names=["projected_decoder_out"],
+ dynamic_axes={
+ "decoder_out": {0: "N"},
+ "projected_decoder_out": {0: "N"},
+ },
+ )
+ logging.info(f"Saved to {decoder_proj_filename}")
+
+
@torch.no_grad()
def main():
args = get_parser().parse_args()
@@ -292,9 +531,32 @@ def main():
model.to("cpu")
model.eval()
- if params.jit is True:
+ if params.onnx is True:
+ convert_scaled_to_non_scaled(model, inplace=True)
+ opset_version = 13
+ logging.info("Exporting to onnx format")
+ encoder_filename = params.exp_dir / "encoder.onnx"
+ export_encoder_model_onnx(
+ model.encoder,
+ encoder_filename,
+ opset_version=opset_version,
+ )
+
+ decoder_filename = params.exp_dir / "decoder.onnx"
+ export_decoder_model_onnx(
+ model.decoder,
+ decoder_filename,
+ opset_version=opset_version,
+ )
+
+ joiner_filename = params.exp_dir / "joiner.onnx"
+ export_joiner_model_onnx(
+ model.joiner,
+ joiner_filename,
+ opset_version=opset_version,
+ )
+ elif params.jit is True:
convert_scaled_to_non_scaled(model, inplace=True)
- logging.info("Using torch.jit.script()")
# We won't use the forward() method of the model in C++, so just ignore
# it here.
# Otherwise, one of its arguments is a ragged tensor and is not
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/onnx_check.py b/egs/librispeech/ASR/pruned_transducer_stateless7/onnx_check.py
new file mode 100755
index 000000000..63acc0922
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/onnx_check.py
@@ -0,0 +1,286 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022 Xiaomi Corporation (Author: Fangjun Kuang)
+#
+# 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.
+
+"""
+This script checks that exported onnx models produce the same output
+with the given torchscript model for the same input.
+"""
+
+import argparse
+import logging
+
+import onnxruntime as ort
+import torch
+
+from icefall import is_module_available
+
+if not is_module_available("onnxruntime"):
+ raise ValueError("Please 'pip install onnxruntime' first.")
+
+
+ort.set_default_logger_severity(3)
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--jit-filename",
+ required=True,
+ type=str,
+ help="Path to the torchscript model",
+ )
+
+ parser.add_argument(
+ "--onnx-encoder-filename",
+ required=True,
+ type=str,
+ help="Path to the onnx encoder model",
+ )
+
+ parser.add_argument(
+ "--onnx-decoder-filename",
+ required=True,
+ type=str,
+ help="Path to the onnx decoder model",
+ )
+
+ parser.add_argument(
+ "--onnx-joiner-filename",
+ required=True,
+ type=str,
+ help="Path to the onnx joiner model",
+ )
+
+ parser.add_argument(
+ "--onnx-joiner-encoder-proj-filename",
+ required=True,
+ type=str,
+ help="Path to the onnx joiner encoder projection model",
+ )
+
+ parser.add_argument(
+ "--onnx-joiner-decoder-proj-filename",
+ required=True,
+ type=str,
+ help="Path to the onnx joiner decoder projection model",
+ )
+
+ return parser
+
+
+def test_encoder(
+ model: torch.jit.ScriptModule,
+ encoder_session: ort.InferenceSession,
+):
+ inputs = encoder_session.get_inputs()
+ outputs = encoder_session.get_outputs()
+ input_names = [n.name for n in inputs]
+ output_names = [n.name for n in outputs]
+
+ assert inputs[0].shape == ["N", "T", 80]
+ assert inputs[1].shape == ["N"]
+
+ for N in [1, 5]:
+ for T in [12, 50]:
+ print("N, T", N, T)
+ x = torch.rand(N, T, 80, dtype=torch.float32)
+ x_lens = torch.randint(low=10, high=T + 1, size=(N,))
+ x_lens[0] = T
+
+ encoder_inputs = {
+ input_names[0]: x.numpy(),
+ input_names[1]: x_lens.numpy(),
+ }
+
+ torch_encoder_out, torch_encoder_out_lens = model.encoder(x, x_lens)
+
+ encoder_out, encoder_out_lens = encoder_session.run(
+ output_names,
+ encoder_inputs,
+ )
+
+ torch_encoder_out, torch_encoder_out_lens = model.encoder(x, x_lens)
+
+ encoder_out = torch.from_numpy(encoder_out)
+ assert torch.allclose(encoder_out, torch_encoder_out, atol=1e-05), (
+ (encoder_out - torch_encoder_out).abs().max(),
+ encoder_out.shape,
+ torch_encoder_out.shape,
+ )
+
+
+def test_decoder(
+ model: torch.jit.ScriptModule,
+ decoder_session: ort.InferenceSession,
+):
+ inputs = decoder_session.get_inputs()
+ outputs = decoder_session.get_outputs()
+ input_names = [n.name for n in inputs]
+ output_names = [n.name for n in outputs]
+
+ assert inputs[0].shape == ["N", 2]
+ for N in [1, 5, 10]:
+ y = torch.randint(low=1, high=500, size=(10, 2))
+
+ decoder_inputs = {input_names[0]: y.numpy()}
+ decoder_out = decoder_session.run(
+ output_names,
+ decoder_inputs,
+ )[0]
+ decoder_out = torch.from_numpy(decoder_out)
+
+ torch_decoder_out = model.decoder(y, need_pad=False)
+ assert torch.allclose(decoder_out, torch_decoder_out, atol=1e-5), (
+ (decoder_out - torch_decoder_out).abs().max()
+ )
+
+
+def test_joiner(
+ model: torch.jit.ScriptModule,
+ joiner_session: ort.InferenceSession,
+ joiner_encoder_proj_session: ort.InferenceSession,
+ joiner_decoder_proj_session: ort.InferenceSession,
+):
+ joiner_inputs = joiner_session.get_inputs()
+ joiner_outputs = joiner_session.get_outputs()
+ joiner_input_names = [n.name for n in joiner_inputs]
+ joiner_output_names = [n.name for n in joiner_outputs]
+
+ assert joiner_inputs[0].shape == ["N", 1, 1, 512]
+ assert joiner_inputs[1].shape == ["N", 1, 1, 512]
+
+ joiner_encoder_proj_inputs = joiner_encoder_proj_session.get_inputs()
+ encoder_proj_input_name = joiner_encoder_proj_inputs[0].name
+
+ assert joiner_encoder_proj_inputs[0].shape == ["N", 384]
+
+ joiner_encoder_proj_outputs = joiner_encoder_proj_session.get_outputs()
+ encoder_proj_output_name = joiner_encoder_proj_outputs[0].name
+
+ joiner_decoder_proj_inputs = joiner_decoder_proj_session.get_inputs()
+ decoder_proj_input_name = joiner_decoder_proj_inputs[0].name
+
+ assert joiner_decoder_proj_inputs[0].shape == ["N", 512]
+
+ joiner_decoder_proj_outputs = joiner_decoder_proj_session.get_outputs()
+ decoder_proj_output_name = joiner_decoder_proj_outputs[0].name
+
+ for N in [1, 5, 10]:
+ encoder_out = torch.rand(N, 384)
+ decoder_out = torch.rand(N, 512)
+
+ projected_encoder_out = torch.rand(N, 1, 1, 512)
+ projected_decoder_out = torch.rand(N, 1, 1, 512)
+
+ joiner_inputs = {
+ joiner_input_names[0]: projected_encoder_out.numpy(),
+ joiner_input_names[1]: projected_decoder_out.numpy(),
+ }
+ joiner_out = joiner_session.run(joiner_output_names, joiner_inputs)[0]
+ joiner_out = torch.from_numpy(joiner_out)
+
+ torch_joiner_out = model.joiner(
+ projected_encoder_out,
+ projected_decoder_out,
+ project_input=False,
+ )
+ assert torch.allclose(joiner_out, torch_joiner_out, atol=1e-5), (
+ (joiner_out - torch_joiner_out).abs().max()
+ )
+
+ # Now test encoder_proj
+ joiner_encoder_proj_inputs = {encoder_proj_input_name: encoder_out.numpy()}
+ joiner_encoder_proj_out = joiner_encoder_proj_session.run(
+ [encoder_proj_output_name], joiner_encoder_proj_inputs
+ )[0]
+ joiner_encoder_proj_out = torch.from_numpy(joiner_encoder_proj_out)
+
+ torch_joiner_encoder_proj_out = model.joiner.encoder_proj(encoder_out)
+ assert torch.allclose(
+ joiner_encoder_proj_out, torch_joiner_encoder_proj_out, atol=1e-5
+ ), ((joiner_encoder_proj_out - torch_joiner_encoder_proj_out).abs().max())
+
+ # Now test decoder_proj
+ joiner_decoder_proj_inputs = {decoder_proj_input_name: decoder_out.numpy()}
+ joiner_decoder_proj_out = joiner_decoder_proj_session.run(
+ [decoder_proj_output_name], joiner_decoder_proj_inputs
+ )[0]
+ joiner_decoder_proj_out = torch.from_numpy(joiner_decoder_proj_out)
+
+ torch_joiner_decoder_proj_out = model.joiner.decoder_proj(decoder_out)
+ assert torch.allclose(
+ joiner_decoder_proj_out, torch_joiner_decoder_proj_out, atol=1e-5
+ ), ((joiner_decoder_proj_out - torch_joiner_decoder_proj_out).abs().max())
+
+
+@torch.no_grad()
+def main():
+ args = get_parser().parse_args()
+ logging.info(vars(args))
+
+ model = torch.jit.load(args.jit_filename)
+
+ options = ort.SessionOptions()
+ options.inter_op_num_threads = 1
+ options.intra_op_num_threads = 1
+
+ logging.info("Test encoder")
+ encoder_session = ort.InferenceSession(
+ args.onnx_encoder_filename,
+ sess_options=options,
+ )
+ test_encoder(model, encoder_session)
+
+ logging.info("Test decoder")
+ decoder_session = ort.InferenceSession(
+ args.onnx_decoder_filename,
+ sess_options=options,
+ )
+ test_decoder(model, decoder_session)
+
+ logging.info("Test joiner")
+ joiner_session = ort.InferenceSession(
+ args.onnx_joiner_filename,
+ sess_options=options,
+ )
+ joiner_encoder_proj_session = ort.InferenceSession(
+ args.onnx_joiner_encoder_proj_filename,
+ sess_options=options,
+ )
+ joiner_decoder_proj_session = ort.InferenceSession(
+ args.onnx_joiner_decoder_proj_filename,
+ sess_options=options,
+ )
+ test_joiner(
+ model,
+ joiner_session,
+ joiner_encoder_proj_session,
+ joiner_decoder_proj_session,
+ )
+ logging.info("Finished checking ONNX models")
+
+
+if __name__ == "__main__":
+ torch.manual_seed(20220727)
+ 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/pruned_transducer_stateless7/onnx_pretrained.py b/egs/librispeech/ASR/pruned_transducer_stateless7/onnx_pretrained.py
new file mode 100755
index 000000000..3a06ee293
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/onnx_pretrained.py
@@ -0,0 +1,388 @@
+#!/usr/bin/env python3
+# Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang)
+#
+# 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.
+"""
+This script loads ONNX models and uses them to decode waves.
+You can use the following command to get the exported models:
+
+./pruned_transducer_stateless7/export.py \
+ --exp-dir ./pruned_transducer_stateless7/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 20 \
+ --avg 10 \
+ --onnx 1
+
+Usage of this script:
+
+./pruned_transducer_stateless7/onnx_pretrained.py \
+ --encoder-model-filename ./pruned_transducer_stateless7/exp/encoder.onnx \
+ --decoder-model-filename ./pruned_transducer_stateless7/exp/decoder.onnx \
+ --joiner-model-filename ./pruned_transducer_stateless7/exp/joiner.onnx \
+ --joiner-encoder-proj-model-filename ./pruned_transducer_stateless7/exp/joiner_encoder_proj.onnx \
+ --joiner-decoder-proj-model-filename ./pruned_transducer_stateless7/exp/joiner_decoder_proj.onnx \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+"""
+
+import argparse
+import logging
+import math
+from typing import List
+
+import kaldifeat
+import numpy as np
+import onnxruntime as ort
+import sentencepiece as spm
+import torch
+import torchaudio
+from torch.nn.utils.rnn import pad_sequence
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--encoder-model-filename",
+ type=str,
+ required=True,
+ help="Path to the encoder onnx model. ",
+ )
+
+ parser.add_argument(
+ "--decoder-model-filename",
+ type=str,
+ required=True,
+ help="Path to the decoder onnx model. ",
+ )
+
+ parser.add_argument(
+ "--joiner-model-filename",
+ type=str,
+ required=True,
+ help="Path to the joiner onnx model. ",
+ )
+
+ parser.add_argument(
+ "--joiner-encoder-proj-model-filename",
+ type=str,
+ required=True,
+ help="Path to the joiner encoder_proj onnx model. ",
+ )
+
+ parser.add_argument(
+ "--joiner-decoder-proj-model-filename",
+ type=str,
+ required=True,
+ help="Path to the joiner decoder_proj onnx model. ",
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ help="""Path to bpe.model.""",
+ )
+
+ parser.add_argument(
+ "sound_files",
+ type=str,
+ nargs="+",
+ help="The input sound file(s) to transcribe. "
+ "Supported formats are those supported by torchaudio.load(). "
+ "For example, wav and flac are supported. "
+ "The sample rate has to be 16kHz.",
+ )
+
+ parser.add_argument(
+ "--sample-rate",
+ type=int,
+ default=16000,
+ help="The sample rate of the input sound file",
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="Context size of the decoder model",
+ )
+
+ return parser
+
+
+def read_sound_files(
+ filenames: List[str], expected_sample_rate: float
+) -> List[torch.Tensor]:
+ """Read a list of sound files into a list 1-D float32 torch tensors.
+ Args:
+ filenames:
+ A list of sound filenames.
+ expected_sample_rate:
+ The expected sample rate of the sound files.
+ Returns:
+ Return a list of 1-D float32 torch tensors.
+ """
+ ans = []
+ for f in filenames:
+ wave, sample_rate = torchaudio.load(f)
+ assert (
+ sample_rate == expected_sample_rate
+ ), f"expected sample rate: {expected_sample_rate}. Given: {sample_rate}"
+ # We use only the first channel
+ ans.append(wave[0])
+ return ans
+
+
+def greedy_search(
+ decoder: ort.InferenceSession,
+ joiner: ort.InferenceSession,
+ joiner_encoder_proj: ort.InferenceSession,
+ joiner_decoder_proj: ort.InferenceSession,
+ encoder_out: np.ndarray,
+ encoder_out_lens: np.ndarray,
+ context_size: int,
+) -> List[List[int]]:
+ """Greedy search in batch mode. It hardcodes --max-sym-per-frame=1.
+ Args:
+ decoder:
+ The decoder model.
+ joiner:
+ The joiner model.
+ joiner_encoder_proj:
+ The joiner encoder projection model.
+ joiner_decoder_proj:
+ The joiner decoder projection model.
+ encoder_out:
+ A 3-D tensor of shape (N, T, C)
+ encoder_out_lens:
+ A 1-D tensor of shape (N,).
+ context_size:
+ The context size of the decoder model.
+ Returns:
+ Return the decoded results for each utterance.
+ """
+ encoder_out = torch.from_numpy(encoder_out)
+ encoder_out_lens = torch.from_numpy(encoder_out_lens)
+ assert encoder_out.ndim == 3
+ assert encoder_out.size(0) >= 1, encoder_out.size(0)
+
+ packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
+ input=encoder_out,
+ lengths=encoder_out_lens.cpu(),
+ batch_first=True,
+ enforce_sorted=False,
+ )
+
+ projected_encoder_out = joiner_encoder_proj.run(
+ [joiner_encoder_proj.get_outputs()[0].name],
+ {joiner_encoder_proj.get_inputs()[0].name: packed_encoder_out.data.numpy()},
+ )[0]
+
+ blank_id = 0 # hard-code to 0
+
+ batch_size_list = packed_encoder_out.batch_sizes.tolist()
+ N = encoder_out.size(0)
+
+ assert torch.all(encoder_out_lens > 0), encoder_out_lens
+ assert N == batch_size_list[0], (N, batch_size_list)
+
+ hyps = [[blank_id] * context_size for _ in range(N)]
+
+ decoder_input_nodes = decoder.get_inputs()
+ decoder_output_nodes = decoder.get_outputs()
+
+ joiner_input_nodes = joiner.get_inputs()
+ joiner_output_nodes = joiner.get_outputs()
+
+ decoder_input = torch.tensor(
+ hyps,
+ dtype=torch.int64,
+ ) # (N, context_size)
+
+ decoder_out = decoder.run(
+ [decoder_output_nodes[0].name],
+ {
+ decoder_input_nodes[0].name: decoder_input.numpy(),
+ },
+ )[0].squeeze(1)
+ projected_decoder_out = joiner_decoder_proj.run(
+ [joiner_decoder_proj.get_outputs()[0].name],
+ {joiner_decoder_proj.get_inputs()[0].name: decoder_out},
+ )[0]
+
+ projected_decoder_out = torch.from_numpy(projected_decoder_out)
+
+ offset = 0
+ for batch_size in batch_size_list:
+ start = offset
+ end = offset + batch_size
+ current_encoder_out = projected_encoder_out[start:end]
+ # current_encoder_out's shape: (batch_size, encoder_out_dim)
+ offset = end
+
+ projected_decoder_out = projected_decoder_out[:batch_size]
+
+ logits = joiner.run(
+ [joiner_output_nodes[0].name],
+ {
+ joiner_input_nodes[0].name: np.expand_dims(
+ np.expand_dims(current_encoder_out, axis=1), axis=1
+ ),
+ joiner_input_nodes[1]
+ .name: projected_decoder_out.unsqueeze(1)
+ .unsqueeze(1)
+ .numpy(),
+ },
+ )[0]
+ logits = torch.from_numpy(logits).squeeze(1).squeeze(1)
+ # logits'shape (batch_size, vocab_size)
+
+ assert logits.ndim == 2, logits.shape
+ y = logits.argmax(dim=1).tolist()
+ emitted = False
+ for i, v in enumerate(y):
+ if v != blank_id:
+ hyps[i].append(v)
+ emitted = True
+ if emitted:
+ # update decoder output
+ decoder_input = [h[-context_size:] for h in hyps[:batch_size]]
+ decoder_input = torch.tensor(
+ decoder_input,
+ dtype=torch.int64,
+ )
+ decoder_out = decoder.run(
+ [decoder_output_nodes[0].name],
+ {
+ decoder_input_nodes[0].name: decoder_input.numpy(),
+ },
+ )[0].squeeze(1)
+ projected_decoder_out = joiner_decoder_proj.run(
+ [joiner_decoder_proj.get_outputs()[0].name],
+ {joiner_decoder_proj.get_inputs()[0].name: decoder_out},
+ )[0]
+ projected_decoder_out = torch.from_numpy(projected_decoder_out)
+
+ sorted_ans = [h[context_size:] for h in hyps]
+ ans = []
+ unsorted_indices = packed_encoder_out.unsorted_indices.tolist()
+ for i in range(N):
+ ans.append(sorted_ans[unsorted_indices[i]])
+
+ return ans
+
+
+@torch.no_grad()
+def main():
+ parser = get_parser()
+ args = parser.parse_args()
+ logging.info(vars(args))
+
+ session_opts = ort.SessionOptions()
+ session_opts.inter_op_num_threads = 1
+ session_opts.intra_op_num_threads = 1
+
+ encoder = ort.InferenceSession(
+ args.encoder_model_filename,
+ sess_options=session_opts,
+ )
+
+ decoder = ort.InferenceSession(
+ args.decoder_model_filename,
+ sess_options=session_opts,
+ )
+
+ joiner = ort.InferenceSession(
+ args.joiner_model_filename,
+ sess_options=session_opts,
+ )
+
+ joiner_encoder_proj = ort.InferenceSession(
+ args.joiner_encoder_proj_model_filename,
+ sess_options=session_opts,
+ )
+
+ joiner_decoder_proj = ort.InferenceSession(
+ args.joiner_decoder_proj_model_filename,
+ sess_options=session_opts,
+ )
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(args.bpe_model)
+
+ logging.info("Constructing Fbank computer")
+ opts = kaldifeat.FbankOptions()
+ opts.device = "cpu"
+ opts.frame_opts.dither = 0
+ opts.frame_opts.snip_edges = False
+ opts.frame_opts.samp_freq = args.sample_rate
+ opts.mel_opts.num_bins = 80
+
+ fbank = kaldifeat.Fbank(opts)
+
+ logging.info(f"Reading sound files: {args.sound_files}")
+ waves = read_sound_files(
+ filenames=args.sound_files,
+ expected_sample_rate=args.sample_rate,
+ )
+
+ logging.info("Decoding started")
+ features = fbank(waves)
+ feature_lengths = [f.size(0) for f in features]
+
+ features = pad_sequence(
+ features,
+ batch_first=True,
+ padding_value=math.log(1e-10),
+ )
+
+ feature_lengths = torch.tensor(feature_lengths, dtype=torch.int64)
+
+ encoder_input_nodes = encoder.get_inputs()
+ encoder_out_nodes = encoder.get_outputs()
+ encoder_out, encoder_out_lens = encoder.run(
+ [encoder_out_nodes[0].name, encoder_out_nodes[1].name],
+ {
+ encoder_input_nodes[0].name: features.numpy(),
+ encoder_input_nodes[1].name: feature_lengths.numpy(),
+ },
+ )
+
+ hyps = greedy_search(
+ decoder=decoder,
+ joiner=joiner,
+ joiner_encoder_proj=joiner_encoder_proj,
+ joiner_decoder_proj=joiner_decoder_proj,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ context_size=args.context_size,
+ )
+ s = "\n"
+ for filename, hyp in zip(args.sound_files, hyps):
+ words = sp.decode(hyp)
+ s += f"{filename}:\n{words}\n\n"
+ logging.info(s)
+
+ logging.info("Decoding Done")
+
+
+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/pruned_transducer_stateless7/optim.py b/egs/librispeech/ASR/pruned_transducer_stateless7/optim.py
index ff8fbb32c..374b78cb3 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless7/optim.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/optim.py
@@ -1,4 +1,4 @@
-# Copyright 2022 Xiaomi Corp. (authors: Daniel Povey)
+# Copyright 2022 Xiaomi Corp. (authors: Daniel Povey)
#
# See ../LICENSE for clarification regarding multiple authors
#
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/scaling.py b/egs/librispeech/ASR/pruned_transducer_stateless7/scaling.py
index 042c9c3e4..156b91f09 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless7/scaling.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/scaling.py
@@ -261,7 +261,7 @@ class RandomGrad(torch.nn.Module):
self.min_abs = min_abs
def forward(self, x: Tensor):
- if torch.jit.is_scripting() or not self.training:
+ if torch.jit.is_scripting() or not self.training or torch.jit.is_tracing():
return x
else:
return RandomGradFunction.apply(x, self.min_abs)
@@ -298,7 +298,7 @@ class SoftmaxFunction(torch.autograd.Function):
def softmax(x: Tensor, dim: int):
- if torch.jit.is_scripting():
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
return x.softmax(dim)
return SoftmaxFunction.apply(x, dim)
@@ -530,7 +530,7 @@ class ActivationBalancer(torch.nn.Module):
self.register_buffer("count", torch.tensor(0, dtype=torch.int64))
def forward(self, x: Tensor) -> Tensor:
- if torch.jit.is_scripting() or not x.requires_grad:
+ if torch.jit.is_scripting() or not x.requires_grad or torch.jit.is_tracing():
return _no_op(x)
count = self.cpu_count
@@ -783,14 +783,14 @@ class WithLoss(torch.autograd.Function):
def with_loss(x, y):
- if torch.jit.is_scripting():
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
return x
# returns x but adds y.sum() to the loss function.
return WithLoss.apply(x, y)
def _no_op(x: Tensor) -> Tensor:
- if torch.jit.is_scripting():
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
return x
else:
# a no-op function that will have a node in the autograd graph,
@@ -862,6 +862,7 @@ class MaxEig(torch.nn.Module):
torch.jit.is_scripting()
or self.max_var_per_eig <= 0
or random.random() > self.cur_prob
+ or torch.jit.is_tracing()
):
return _no_op(x)
@@ -1013,7 +1014,7 @@ class DoubleSwish(torch.nn.Module):
"""Return double-swish activation function which is an approximation to Swish(Swish(x)),
that we approximate closely with x * sigmoid(x-1).
"""
- if torch.jit.is_scripting():
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
return x * torch.sigmoid(x - 1.0)
return DoubleSwishFunction.apply(x)
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/test_model.py b/egs/librispeech/ASR/pruned_transducer_stateless7/test_model.py
index db7fb7b3e..cdf914df3 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless7/test_model.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/test_model.py
@@ -20,19 +20,21 @@
To run this file, do:
cd icefall/egs/librispeech/ASR
- python ./pruned_transducer_stateless4/test_model.py
+ python ./pruned_transducer_stateless7/test_model.py
"""
+import torch
+
+from scaling_converter import convert_scaled_to_non_scaled
from train import get_params, get_transducer_model
-def test_model_1():
+def test_model():
params = get_params()
params.vocab_size = 500
params.blank_id = 0
params.context_size = 2
params.num_encoder_layers = "2,4,3,2,4"
- # params.feedforward_dims = "1024,1024,1536,1536,1024"
params.feedforward_dims = "1024,1024,2048,2048,1024"
params.nhead = "8,8,8,8,8"
params.encoder_dims = "384,384,384,384,384"
@@ -47,9 +49,19 @@ def test_model_1():
num_param = sum([p.numel() for p in model.parameters()])
print(f"Number of model parameters: {num_param}")
+ # Test jit script
+ convert_scaled_to_non_scaled(model, inplace=True)
+ # We won't use the forward() method of the model in C++, so just ignore
+ # it here.
+ # Otherwise, one of its arguments is a ragged tensor and is not
+ # torch scriptabe.
+ model.__class__.forward = torch.jit.ignore(model.__class__.forward)
+ print("Using torch.jit.script")
+ model = torch.jit.script(model)
+
def main():
- test_model_1()
+ test_model()
if __name__ == "__main__":
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/test_onnx.py b/egs/librispeech/ASR/pruned_transducer_stateless7/test_onnx.py
new file mode 100644
index 000000000..2440d267c
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/test_onnx.py
@@ -0,0 +1,374 @@
+#!/usr/bin/env python3
+# Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang)
+#
+# 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.
+
+
+"""
+This file is to test that models can be exported to onnx.
+"""
+import os
+
+from icefall import is_module_available
+
+if not is_module_available("onnxruntime"):
+ raise ValueError("Please 'pip install onnxruntime' first.")
+
+import onnxruntime as ort
+import torch
+from scaling_converter import convert_scaled_to_non_scaled
+from zipformer import (
+ Conv2dSubsampling,
+ RelPositionalEncoding,
+ Zipformer,
+ ZipformerEncoder,
+ ZipformerEncoderLayer,
+)
+
+ort.set_default_logger_severity(3)
+
+
+def test_conv2d_subsampling():
+ filename = "conv2d_subsampling.onnx"
+ opset_version = 13
+ N = 30
+ T = 50
+ num_features = 80
+ d_model = 512
+ x = torch.rand(N, T, num_features)
+
+ encoder_embed = Conv2dSubsampling(num_features, d_model)
+ encoder_embed.eval()
+ encoder_embed = convert_scaled_to_non_scaled(encoder_embed, inplace=True)
+
+ torch.onnx.export(
+ encoder_embed,
+ x,
+ filename,
+ verbose=False,
+ opset_version=opset_version,
+ input_names=["x"],
+ output_names=["y"],
+ dynamic_axes={
+ "x": {0: "N", 1: "T"},
+ "y": {0: "N", 1: "T"},
+ },
+ )
+
+ options = ort.SessionOptions()
+ options.inter_op_num_threads = 1
+ options.intra_op_num_threads = 1
+
+ session = ort.InferenceSession(
+ filename,
+ sess_options=options,
+ )
+
+ input_nodes = session.get_inputs()
+ assert input_nodes[0].name == "x"
+ assert input_nodes[0].shape == ["N", "T", num_features]
+
+ inputs = {input_nodes[0].name: x.numpy()}
+
+ onnx_y = session.run(["y"], inputs)[0]
+
+ onnx_y = torch.from_numpy(onnx_y)
+ torch_y = encoder_embed(x)
+ assert torch.allclose(onnx_y, torch_y, atol=1e-05), (onnx_y - torch_y).abs().max()
+
+ os.remove(filename)
+
+
+def test_rel_pos():
+ filename = "rel_pos.onnx"
+
+ opset_version = 13
+ N = 30
+ T = 50
+ num_features = 80
+ d_model = 512
+ x = torch.rand(N, T, num_features)
+
+ encoder_pos = RelPositionalEncoding(d_model, dropout_rate=0.1)
+ encoder_pos.eval()
+ encoder_pos = convert_scaled_to_non_scaled(encoder_pos, inplace=True)
+
+ x = x.permute(1, 0, 2)
+
+ torch.onnx.export(
+ encoder_pos,
+ x,
+ filename,
+ verbose=False,
+ opset_version=opset_version,
+ input_names=["x"],
+ output_names=["pos_emb"],
+ dynamic_axes={
+ "x": {0: "N", 1: "T"},
+ "pos_emb": {0: "N", 1: "T"},
+ },
+ )
+
+ options = ort.SessionOptions()
+ options.inter_op_num_threads = 1
+ options.intra_op_num_threads = 1
+
+ session = ort.InferenceSession(
+ filename,
+ sess_options=options,
+ )
+
+ input_nodes = session.get_inputs()
+ assert input_nodes[0].name == "x"
+ assert input_nodes[0].shape == ["N", "T", num_features]
+
+ inputs = {input_nodes[0].name: x.numpy()}
+ onnx_pos_emb = session.run(["pos_emb"], inputs)
+ onnx_pos_emb = torch.from_numpy(onnx_pos_emb[0])
+
+ torch_pos_emb = encoder_pos(x)
+ assert torch.allclose(onnx_pos_emb, torch_pos_emb, atol=1e-05), (
+ (onnx_pos_emb - torch_pos_emb).abs().max()
+ )
+ print(onnx_pos_emb.abs().sum(), torch_pos_emb.abs().sum())
+
+ os.remove(filename)
+
+
+def test_zipformer_encoder_layer():
+ filename = "zipformer_encoder_layer.onnx"
+ opset_version = 13
+ N = 30
+ T = 50
+
+ d_model = 384
+ attention_dim = 192
+ nhead = 8
+ feedforward_dim = 1024
+ dropout = 0.1
+ cnn_module_kernel = 31
+ pos_dim = 4
+
+ x = torch.rand(N, T, d_model)
+
+ encoder_pos = RelPositionalEncoding(d_model, dropout)
+ encoder_pos.eval()
+ encoder_pos = convert_scaled_to_non_scaled(encoder_pos, inplace=True)
+
+ x = x.permute(1, 0, 2)
+ pos_emb = encoder_pos(x)
+
+ encoder_layer = ZipformerEncoderLayer(
+ d_model,
+ attention_dim,
+ nhead,
+ feedforward_dim,
+ dropout,
+ cnn_module_kernel,
+ pos_dim,
+ )
+ encoder_layer.eval()
+ encoder_layer = convert_scaled_to_non_scaled(encoder_layer, inplace=True)
+
+ torch.onnx.export(
+ encoder_layer,
+ (x, pos_emb),
+ filename,
+ verbose=False,
+ opset_version=opset_version,
+ input_names=["x", "pos_emb"],
+ output_names=["y"],
+ dynamic_axes={
+ "x": {0: "T", 1: "N"},
+ "pos_emb": {0: "N", 1: "T"},
+ "y": {0: "T", 1: "N"},
+ },
+ )
+
+ options = ort.SessionOptions()
+ options.inter_op_num_threads = 1
+ options.intra_op_num_threads = 1
+
+ session = ort.InferenceSession(
+ filename,
+ sess_options=options,
+ )
+
+ input_nodes = session.get_inputs()
+ inputs = {
+ input_nodes[0].name: x.numpy(),
+ input_nodes[1].name: pos_emb.numpy(),
+ }
+ onnx_y = session.run(["y"], inputs)[0]
+ onnx_y = torch.from_numpy(onnx_y)
+
+ torch_y = encoder_layer(x, pos_emb)
+ assert torch.allclose(onnx_y, torch_y, atol=1e-05), (onnx_y - torch_y).abs().max()
+
+ print(onnx_y.abs().sum(), torch_y.abs().sum(), onnx_y.shape, torch_y.shape)
+
+ os.remove(filename)
+
+
+def test_zipformer_encoder():
+ filename = "zipformer_encoder.onnx"
+
+ opset_version = 13
+ N = 3
+ T = 15
+
+ d_model = 512
+ attention_dim = 192
+ nhead = 8
+ feedforward_dim = 1024
+ dropout = 0.1
+ cnn_module_kernel = 31
+ pos_dim = 4
+ num_encoder_layers = 12
+
+ warmup_batches = 4000.0
+ warmup_begin = warmup_batches / (num_encoder_layers + 1)
+ warmup_end = warmup_batches / (num_encoder_layers + 1)
+
+ x = torch.rand(N, T, d_model)
+
+ encoder_layer = ZipformerEncoderLayer(
+ d_model,
+ attention_dim,
+ nhead,
+ feedforward_dim,
+ dropout,
+ cnn_module_kernel,
+ pos_dim,
+ )
+ encoder = ZipformerEncoder(
+ encoder_layer, num_encoder_layers, dropout, warmup_begin, warmup_end
+ )
+ encoder.eval()
+ encoder = convert_scaled_to_non_scaled(encoder, inplace=True)
+
+ # jit_model = torch.jit.trace(encoder, (pos_emb))
+
+ torch_y = encoder(x)
+
+ torch.onnx.export(
+ encoder,
+ (x),
+ filename,
+ verbose=False,
+ opset_version=opset_version,
+ input_names=["x"],
+ output_names=["y"],
+ dynamic_axes={
+ "x": {0: "T", 1: "N"},
+ "y": {0: "T", 1: "N"},
+ },
+ )
+
+ options = ort.SessionOptions()
+ options.inter_op_num_threads = 1
+ options.intra_op_num_threads = 1
+
+ session = ort.InferenceSession(
+ filename,
+ sess_options=options,
+ )
+
+ input_nodes = session.get_inputs()
+ inputs = {
+ input_nodes[0].name: x.numpy(),
+ }
+ onnx_y = session.run(["y"], inputs)[0]
+ onnx_y = torch.from_numpy(onnx_y)
+
+ torch_y = encoder(x)
+ assert torch.allclose(onnx_y, torch_y, atol=1e-05), (onnx_y - torch_y).abs().max()
+
+ print(onnx_y.abs().sum(), torch_y.abs().sum(), onnx_y.shape, torch_y.shape)
+
+ os.remove(filename)
+
+
+def test_zipformer():
+ filename = "zipformer.onnx"
+ opset_version = 11
+ N = 3
+ T = 15
+ num_features = 80
+ x = torch.rand(N, T, num_features)
+ x_lens = torch.full((N,), fill_value=T, dtype=torch.int64)
+
+ zipformer = Zipformer(num_features=num_features)
+ zipformer.eval()
+ zipformer = convert_scaled_to_non_scaled(zipformer, inplace=True)
+
+ # jit_model = torch.jit.trace(zipformer, (x, x_lens))
+ torch.onnx.export(
+ zipformer,
+ (x, x_lens),
+ filename,
+ verbose=False,
+ opset_version=opset_version,
+ input_names=["x", "x_lens"],
+ output_names=["y", "y_lens"],
+ dynamic_axes={
+ "x": {0: "N", 1: "T"},
+ "x_lens": {0: "N"},
+ "y": {0: "N", 1: "T"},
+ "y_lens": {0: "N"},
+ },
+ )
+ options = ort.SessionOptions()
+ options.inter_op_num_threads = 1
+ options.intra_op_num_threads = 1
+
+ session = ort.InferenceSession(
+ filename,
+ sess_options=options,
+ )
+
+ input_nodes = session.get_inputs()
+ inputs = {
+ input_nodes[0].name: x.numpy(),
+ input_nodes[1].name: x_lens.numpy(),
+ }
+ onnx_y, onnx_y_lens = session.run(["y", "y_lens"], inputs)
+ onnx_y = torch.from_numpy(onnx_y)
+ onnx_y_lens = torch.from_numpy(onnx_y_lens)
+
+ torch_y, torch_y_lens = zipformer(x, x_lens)
+ assert torch.allclose(onnx_y, torch_y, atol=1e-05), (onnx_y - torch_y).abs().max()
+
+ assert torch.allclose(onnx_y_lens, torch_y_lens, atol=1e-05), (
+ (onnx_y_lens - torch_y_lens).abs().max()
+ )
+ print(onnx_y.abs().sum(), torch_y.abs().sum(), onnx_y.shape, torch_y.shape)
+ print(onnx_y_lens, torch_y_lens)
+
+ os.remove(filename)
+
+
+@torch.no_grad()
+def main():
+ test_conv2d_subsampling()
+ test_rel_pos()
+ test_zipformer_encoder_layer()
+ test_zipformer_encoder()
+ test_zipformer()
+
+
+if __name__ == "__main__":
+ torch.manual_seed(20221011)
+ main()
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
index b007a7308..b1717ec64 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
@@ -1,5 +1,5 @@
#!/usr/bin/env python3
-# Copyright (c) 2021 University of Chinese Academy of Sciences (author: Han Zhu)
+# Copyright 2022 Xiaomi Corp. (authors: Daniel Povey)
#
# See ../../../../LICENSE for clarification regarding multiple authors
#
@@ -81,7 +81,6 @@ class Zipformer(EncoderInterface):
super(Zipformer, self).__init__()
self.num_features = num_features
- self.encoder_unmasked_dims = encoder_unmasked_dims
assert 0 < encoder_dims[0] <= encoder_dims[1]
self.encoder_dims = encoder_dims
self.encoder_unmasked_dims = encoder_unmasked_dims
@@ -211,7 +210,7 @@ class Zipformer(EncoderInterface):
(num_frames, batch_size, encoder_dims0)
"""
num_encoders = len(self.encoder_dims)
- if torch.jit.is_scripting() or not self.training:
+ if torch.jit.is_scripting() or not self.training or torch.jit.is_tracing():
return [1.0] * num_encoders
(num_frames0, batch_size, _encoder_dims0) = x.shape
@@ -294,7 +293,7 @@ class Zipformer(EncoderInterface):
k = self.skip_layers[i]
if isinstance(k, int):
layer_skip_dropout_prob = self._get_layer_skip_dropout_prob()
- if torch.jit.is_scripting():
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
x = skip_module(outputs[k], x)
elif (not self.training) or random.random() > layer_skip_dropout_prob:
x = skip_module(outputs[k], x)
@@ -387,7 +386,7 @@ class ZipformerEncoderLayer(nn.Module):
)
def get_bypass_scale(self):
- if torch.jit.is_scripting() or not self.training:
+ if torch.jit.is_scripting() or not self.training or torch.jit.is_tracing():
return self.bypass_scale
if random.random() < 0.1:
# ensure we get grads if self.bypass_scale becomes out of range
@@ -408,7 +407,7 @@ class ZipformerEncoderLayer(nn.Module):
# return dropout rate for the dynamic modules (self_attn, pooling, convolution); this
# starts at 0.2 and rapidly decreases to 0. Its purpose is to keep the training stable
# at the beginning, by making the network focus on the feedforward modules.
- if torch.jit.is_scripting() or not self.training:
+ if torch.jit.is_scripting() or not self.training or torch.jit.is_tracing():
return 0.0
warmup_period = 2000.0
initial_dropout_rate = 0.2
@@ -453,12 +452,12 @@ class ZipformerEncoderLayer(nn.Module):
dynamic_dropout = self.get_dynamic_dropout_rate()
# pooling module
- if torch.jit.is_scripting():
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
src = src + self.pooling(src, key_padding_mask=src_key_padding_mask)
- elif random.random() > dynamic_dropout:
+ elif random.random() >= dynamic_dropout:
src = src + self.pooling(src, key_padding_mask=src_key_padding_mask)
- if torch.jit.is_scripting():
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
src_att, attn_weights = self.self_attn(
src,
pos_emb=pos_emb,
@@ -479,7 +478,7 @@ class ZipformerEncoderLayer(nn.Module):
src, src_key_padding_mask=src_key_padding_mask
)
else:
- use_self_attn = random.random() > dynamic_dropout
+ use_self_attn = random.random() >= dynamic_dropout
if use_self_attn:
src_att, attn_weights = self.self_attn(
src,
@@ -489,7 +488,7 @@ class ZipformerEncoderLayer(nn.Module):
)
src = src + src_att
- if random.random() > dynamic_dropout:
+ if random.random() >= dynamic_dropout:
src = src + self.conv_module1(
src, src_key_padding_mask=src_key_padding_mask
)
@@ -498,7 +497,7 @@ class ZipformerEncoderLayer(nn.Module):
if use_self_attn:
src = src + self.self_attn.forward2(src, attn_weights)
- if random.random() > dynamic_dropout:
+ if random.random() >= dynamic_dropout:
src = src + self.conv_module2(
src, src_key_padding_mask=src_key_padding_mask
)
@@ -659,7 +658,7 @@ class ZipformerEncoder(nn.Module):
pos_emb = self.encoder_pos(src)
output = src
- if torch.jit.is_scripting():
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
layers_to_drop = []
else:
rnd_seed = src.numel() + random.randint(0, 1000)
@@ -668,7 +667,7 @@ class ZipformerEncoder(nn.Module):
output = output * feature_mask
for i, mod in enumerate(self.layers):
- if not torch.jit.is_scripting():
+ if not torch.jit.is_scripting() and not torch.jit.is_tracing():
if i in layers_to_drop:
continue
output = mod(
@@ -742,7 +741,7 @@ class DownsampledZipformerEncoder(nn.Module):
src,
feature_mask=feature_mask,
mask=mask,
- src_key_padding_mask=mask,
+ src_key_padding_mask=src_key_padding_mask,
)
src = self.upsample(src)
# remove any extra frames that are not a multiple of downsample_factor
@@ -865,7 +864,7 @@ class SimpleCombiner(torch.nn.Module):
assert src1.shape[:-1] == src2.shape[:-1], (src1.shape, src2.shape)
weight1 = self.weight1
- if not torch.jit.is_scripting():
+ if not torch.jit.is_scripting() and not torch.jit.is_tracing():
if (
self.training
and random.random() < 0.25
@@ -908,7 +907,7 @@ class RelPositionalEncoding(torch.nn.Module):
self.d_model = d_model
self.dropout = torch.nn.Dropout(dropout_rate)
self.pe = None
- self.extend_pe(torch.tensor(0.0).expand(1, max_len))
+ self.extend_pe(torch.tensor(0.0).expand(max_len))
def extend_pe(self, x: Tensor) -> None:
"""Reset the positional encodings."""
@@ -1259,21 +1258,31 @@ class RelPositionMultiheadAttention(nn.Module):
# the following .as_strided() expression converts the last axis of pos_weights from relative
# to absolute position. I don't know whether I might have got the time-offsets backwards or
# not, but let this code define which way round it is supposed to be.
- pos_weights = pos_weights.as_strided(
- (bsz, num_heads, seq_len, seq_len),
- (
- pos_weights.stride(0),
- pos_weights.stride(1),
- pos_weights.stride(2) - pos_weights.stride(3),
- pos_weights.stride(3),
- ),
- storage_offset=pos_weights.stride(3) * (seq_len - 1),
- )
+ if torch.jit.is_tracing():
+ (batch_size, num_heads, time1, n) = pos_weights.shape
+ rows = torch.arange(start=time1 - 1, end=-1, step=-1)
+ cols = torch.arange(seq_len)
+ rows = rows.repeat(batch_size * num_heads).unsqueeze(-1)
+ indexes = rows + cols
+ pos_weights = pos_weights.reshape(-1, n)
+ pos_weights = torch.gather(pos_weights, dim=1, index=indexes)
+ pos_weights = pos_weights.reshape(batch_size, num_heads, time1, seq_len)
+ else:
+ pos_weights = pos_weights.as_strided(
+ (bsz, num_heads, seq_len, seq_len),
+ (
+ pos_weights.stride(0),
+ pos_weights.stride(1),
+ pos_weights.stride(2) - pos_weights.stride(3),
+ pos_weights.stride(3),
+ ),
+ storage_offset=pos_weights.stride(3) * (seq_len - 1),
+ )
# caution: they are really scores at this point.
attn_output_weights = torch.matmul(q, k) + pos_weights
- if not torch.jit.is_scripting():
+ if not torch.jit.is_scripting() and not torch.jit.is_tracing():
if training and random.random() < 0.1:
# This is a harder way of limiting the attention scores to not be too large.
# It incurs a penalty if any of them has an absolute value greater than 50.0.
@@ -1290,17 +1299,13 @@ class RelPositionMultiheadAttention(nn.Module):
bsz * num_heads, seq_len, seq_len
)
- assert list(attn_output_weights.size()) == [
- bsz * num_heads,
- seq_len,
- seq_len,
- ]
-
if attn_mask is not None:
if attn_mask.dtype == torch.bool:
- attn_output_weights.masked_fill_(attn_mask, float("-inf"))
+ attn_output_weights = attn_output_weights.masked_fill(
+ attn_mask, float("-inf")
+ )
else:
- attn_output_weights += attn_mask
+ attn_output_weights = attn_output_weights + attn_mask
if key_padding_mask is not None:
attn_output_weights = attn_output_weights.view(
@@ -1320,6 +1325,34 @@ class RelPositionMultiheadAttention(nn.Module):
# only storing the half-precision output for backprop purposes.
attn_output_weights = softmax(attn_output_weights, dim=-1)
+ # If we are using chunk-wise attention mask and setting a limited
+ # num_left_chunks, the attention may only see the padding values which
+ # will also be masked out by `key_padding_mask`. At this circumstances,
+ # the whole column of `attn_output_weights` will be `-inf`
+ # (i.e. be `nan` after softmax). So we fill `0.0` at the masking
+ # positions to avoid invalid loss value below.
+ if (
+ attn_mask is not None
+ and attn_mask.dtype == torch.bool
+ and key_padding_mask is not None
+ ):
+ if attn_mask.size(0) != 1:
+ attn_mask = attn_mask.view(bsz, num_heads, seq_len, seq_len)
+ combined_mask = attn_mask | key_padding_mask.unsqueeze(1).unsqueeze(2)
+ else:
+ # attn_mask.shape == (1, tgt_len, src_len)
+ combined_mask = attn_mask.unsqueeze(0) | key_padding_mask.unsqueeze(
+ 1
+ ).unsqueeze(2)
+
+ attn_output_weights = attn_output_weights.view(
+ bsz, num_heads, seq_len, seq_len
+ )
+ attn_output_weights = attn_output_weights.masked_fill(combined_mask, 0.0)
+ attn_output_weights = attn_output_weights.view(
+ bsz * num_heads, seq_len, seq_len
+ )
+
attn_output_weights = nn.functional.dropout(
attn_output_weights, p=dropout_p, training=training
)
@@ -1360,7 +1393,7 @@ class RelPositionMultiheadAttention(nn.Module):
# now v: (bsz * num_heads, seq_len, head_dim // 2)
attn_output = torch.bmm(attn_weights, v)
- if not torch.jit.is_scripting():
+ if not torch.jit.is_scripting() and not torch.jit.is_tracing():
if random.random() < 0.001 or __name__ == "__main__":
self._print_attn_stats(attn_weights, attn_output)
@@ -1435,7 +1468,10 @@ class PoolingModule(nn.Module):
a Tensor of shape (1, N, C)
"""
if key_padding_mask is not None:
- pooling_mask = key_padding_mask.logical_not().to(x.dtype) # (N, T)
+ if torch.jit.is_tracing():
+ pooling_mask = (~key_padding_mask).to(x.dtype)
+ else:
+ pooling_mask = key_padding_mask.logical_not().to(x.dtype) # (N, T)
pooling_mask = pooling_mask / pooling_mask.sum(dim=1, keepdim=True)
pooling_mask = pooling_mask.transpose(0, 1).contiguous().unsqueeze(-1)
# now pooling_mask: (T, N, 1)
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/ctc_decode.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/ctc_decode.py
index 9c23e7d66..4b373e4c7 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/ctc_decode.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/ctc_decode.py
@@ -44,7 +44,7 @@ Usage:
--exp-dir ./pruned_transducer_stateless7_ctc/exp \
--max-duration 600 \
--hlg-scale 0.8 \
- --decoding-method 1best
+ --decoding-method nbest
(4) nbest-rescoring
./pruned_transducer_stateless7_ctc/ctc_decode.py \
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/export.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/export.py
index 59a393739..c1607699f 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/export.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/export.py
@@ -72,14 +72,14 @@ Check ./pretrained.py for its usage.
Note: If you don't want to train a model from scratch, we have
provided one for you. You can get it at
-https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless7-2022-11-11
+https://huggingface.co/Zengwei/icefall-asr-librispeech-pruned-transducer-stateless7-ctc-2022-12-01
with the following commands:
sudo apt-get install git-lfs
git lfs install
- git clone https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless7-2022-11-11
- # You will find the pre-trained model in icefall-asr-librispeech-pruned-transducer-stateless7-2022-11-11/exp
+ git clone https://huggingface.co/Zengwei/icefall-asr-librispeech-pruned-transducer-stateless7-ctc-2022-12-01
+ # You will find the pre-trained model in icefall-asr-librispeech-pruned-transducer-stateless7-ctc-2022-12-01/exp
"""
import argparse
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/jit_pretrained_ctc.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/jit_pretrained_ctc.py
index d3343d34a..d50d231d5 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/jit_pretrained_ctc.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/jit_pretrained_ctc.py
@@ -31,7 +31,7 @@ Usage of this script:
(1) ctc-decoding
./pruned_transducer_stateless7_ctc/jit_pretrained_ctc.py \
- --nn-model-filename ./pruned_transducer_stateless7_ctc/exp/cpu_jit.pt \
+ --model-filename ./pruned_transducer_stateless7_ctc/exp/cpu_jit.pt \
--bpe-model data/lang_bpe_500/bpe.model \
--method ctc-decoding \
--sample-rate 16000 \
@@ -40,7 +40,7 @@ Usage of this script:
(2) 1best
./pruned_transducer_stateless7_ctc/jit_pretrained_ctc.py \
- --nn-model-filename ./pruned_transducer_stateless7_ctc/exp/cpu_jit.pt \
+ --model-filename ./pruned_transducer_stateless7_ctc/exp/cpu_jit.pt \
--HLG data/lang_bpe_500/HLG.pt \
--words-file data/lang_bpe_500/words.txt \
--method 1best \
@@ -51,7 +51,7 @@ Usage of this script:
(3) nbest-rescoring
./pruned_transducer_stateless7_ctc/jit_pretrained_ctc.py \
- --nn-model-filename ./pruned_transducer_stateless7_ctc/exp/cpu_jit.pt \
+ --model-filename ./pruned_transducer_stateless7_ctc/exp/cpu_jit.pt \
--HLG data/lang_bpe_500/HLG.pt \
--words-file data/lang_bpe_500/words.txt \
--G data/lm/G_4_gram.pt \
@@ -63,7 +63,7 @@ Usage of this script:
(4) whole-lattice-rescoring
./pruned_transducer_stateless7_ctc/jit_pretrained_ctc.py \
- --nn-model-filename ./pruned_transducer_stateless7_ctc/exp/cpu_jit.pt \
+ --model-filename ./pruned_transducer_stateless7_ctc/exp/cpu_jit.pt \
--HLG data/lang_bpe_500/HLG.pt \
--words-file data/lang_bpe_500/words.txt \
--G data/lm/G_4_gram.pt \
@@ -304,7 +304,10 @@ def main():
batch_size = nnet_output.shape[0]
supervision_segments = torch.tensor(
- [[i, 0, nnet_output.shape[1]] for i in range(batch_size)],
+ [
+ [i, 0, feature_lengths[i] // params.subsampling_factor]
+ for i in range(batch_size)
+ ],
dtype=torch.int32,
)
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/pretrained_ctc.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/pretrained_ctc.py
index 74aef1bc7..5d460edb5 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/pretrained_ctc.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/pretrained_ctc.py
@@ -322,7 +322,10 @@ def main():
batch_size = nnet_output.shape[0]
supervision_segments = torch.tensor(
- [[i, 0, nnet_output.shape[1]] for i in range(batch_size)],
+ [
+ [i, 0, feature_lengths[i] // params.subsampling_factor]
+ for i in range(batch_size)
+ ],
dtype=torch.int32,
)
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/train.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/train.py
index 162ad8412..5a05e1836 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/train.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc/train.py
@@ -1086,7 +1086,33 @@ def run(rank, world_size, args):
# You should use ../local/display_manifest_statistics.py to get
# an utterance duration distribution for your dataset to select
# the threshold
- return 1.0 <= c.duration <= 20.0
+ if c.duration < 1.0 or c.duration > 20.0:
+ logging.warning(
+ f"Exclude cut with ID {c.id} from training. Duration: {c.duration}"
+ )
+ return False
+
+ # In pruned RNN-T, we require that T >= S
+ # where T is the number of feature frames after subsampling
+ # and S is the number of tokens in the utterance
+
+ # In ./zipformer.py, the conv module uses the following expression
+ # for subsampling
+ T = ((c.num_frames - 7) // 2 + 1) // 2
+ tokens = sp.encode(c.supervisions[0].text, out_type=str)
+
+ if T < len(tokens):
+ logging.warning(
+ f"Exclude cut with ID {c.id} from training. "
+ f"Number of frames (before subsampling): {c.num_frames}. "
+ f"Number of frames (after subsampling): {T}. "
+ f"Text: {c.supervisions[0].text}. "
+ f"Tokens: {tokens}. "
+ f"Number of tokens: {len(tokens)}"
+ )
+ return False
+
+ return True
train_cuts = train_cuts.filter(remove_short_and_long_utt)
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/__init__.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/__init__.py
new file mode 100755
index 000000000..e69de29bb
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/asr_datamodule.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/asr_datamodule.py
new file mode 120000
index 000000000..a074d6085
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/asr_datamodule.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/asr_datamodule.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/beam_search.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/beam_search.py
new file mode 120000
index 000000000..8554e44cc
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/beam_search.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/beam_search.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/ctc_decode.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/ctc_decode.py
new file mode 100755
index 000000000..f137485b2
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/ctc_decode.py
@@ -0,0 +1,809 @@
+#!/usr/bin/env python3
+#
+# Copyright 2021-2022 Xiaomi Corporation (Author: Fangjun Kuang,
+# Liyong Guo,
+# Quandong Wang,
+# Zengwei Yao)
+#
+# 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.
+"""
+Usage:
+(1) ctc-decoding
+./pruned_transducer_stateless7_ctc_bs/ctc_decode.py \
+ --epoch 30 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method ctc-decoding
+(2) 1best
+./pruned_transducer_stateless7_ctc_bs/ctc_decode.py \
+ --epoch 30 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --hlg-scale 0.8 \
+ --decoding-method 1best
+(3) nbest
+./pruned_transducer_stateless7_ctc_bs/ctc_decode.py \
+ --epoch 30 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --hlg-scale 0.8 \
+ --decoding-method nbest
+(4) nbest-rescoring
+./pruned_transducer_stateless7_ctc_bs/ctc_decode.py \
+ --epoch 30 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --hlg-scale 0.8 \
+ --lm-dir data/lm \
+ --decoding-method nbest-rescoring
+(5) whole-lattice-rescoring
+./pruned_transducer_stateless7_ctc_bs/ctc_decode.py \
+ --epoch 30 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --hlg-scale 0.8 \
+ --lm-dir data/lm \
+ --decoding-method whole-lattice-rescoring
+"""
+
+
+import argparse
+import logging
+import math
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import k2
+import sentencepiece as spm
+import torch
+import torch.nn as nn
+from asr_datamodule import LibriSpeechAsrDataModule
+from train import add_model_arguments, get_params, get_transducer_model
+
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from icefall.decode import (
+ get_lattice,
+ nbest_decoding,
+ nbest_oracle,
+ 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,
+ str2bool,
+ write_error_stats,
+)
+
+LOG_EPS = math.log(1e-10)
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=30,
+ help="""It specifies the checkpoint to use for decoding.
+ Note: Epoch counts from 1.
+ You can specify --avg to use more checkpoints for model averaging.""",
+ )
+
+ parser.add_argument(
+ "--iter",
+ type=int,
+ default=0,
+ help="""If positive, --epoch is ignored and it
+ will use the checkpoint exp_dir/checkpoint-iter.pt.
+ You can specify --avg to use more checkpoints for model averaging.
+ """,
+ )
+
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=15,
+ help="Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch' and '--iter'",
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=True,
+ help="Whether to load averaged model. Currently it only supports "
+ "using --epoch. If True, it would decode with the averaged model "
+ "over the epoch range from `epoch-avg` (excluded) to `epoch`."
+ "Actually only the models with epoch number of `epoch-avg` and "
+ "`epoch` are loaded for averaging. ",
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="pruned_transducer_stateless7_ctc_bs/exp",
+ help="The experiment dir",
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ default="data/lang_bpe_500/bpe.model",
+ help="Path to the BPE model",
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=Path,
+ default="data/lang_bpe_500",
+ help="The lang dir containing word table and LG graph",
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram, 2 means tri-gram",
+ )
+
+ parser.add_argument(
+ "--decoding-method",
+ type=str,
+ default="ctc-decoding",
+ help="""Decoding method.
+ Supported values are:
+ - (1) ctc-decoding. Use CTC decoding. It uses a sentence piece
+ model, i.e., lang_dir/bpe.model, to convert word pieces to words.
+ It needs neither a lexicon nor an n-gram LM.
+ - (2) 1best. Extract the best path from the decoding lattice as the
+ decoding result.
+ - (3) nbest. Extract n paths from the decoding lattice; the path
+ with the highest score is the decoding result.
+ - (4) nbest-rescoring. Extract n paths from the decoding lattice,
+ rescore them with an n-gram LM (e.g., a 4-gram LM), the path with
+ the highest score is the decoding result.
+ - (5) whole-lattice-rescoring. Rescore the decoding lattice with an
+ n-gram LM (e.g., a 4-gram LM), the best path of rescored lattice
+ is the decoding result.
+ you have trained an RNN LM using ./rnn_lm/train.py
+ - (6) nbest-oracle. Its WER is the lower bound of any n-best
+ rescoring method can achieve. Useful for debugging n-best
+ rescoring method.
+ """,
+ )
+
+ parser.add_argument(
+ "--num-paths",
+ type=int,
+ default=100,
+ help="""Number of paths for n-best based decoding method.
+ Used only when "method" is one of the following values:
+ nbest, nbest-rescoring, and nbest-oracle
+ """,
+ )
+
+ parser.add_argument(
+ "--nbest-scale",
+ type=float,
+ default=0.5,
+ help="""The scale to be applied to `lattice.scores`.
+ It's needed if you use any kinds of n-best based rescoring.
+ Used only when "method" is one of the following values:
+ nbest, nbest-rescoring, and nbest-oracle
+ A smaller value results in more unique paths.
+ """,
+ )
+
+ parser.add_argument(
+ "--hlg-scale",
+ type=float,
+ default=0.8,
+ help="""The scale to be applied to `hlg.scores`.
+ """,
+ )
+
+ parser.add_argument(
+ "--lm-dir",
+ type=str,
+ default="data/lm",
+ help="""The n-gram LM dir.
+ It should contain either G_4_gram.pt or G_4_gram.fst.txt
+ """,
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def get_decoding_params() -> AttributeDict:
+ """Parameters for decoding."""
+ params = AttributeDict(
+ {
+ "frame_shift_ms": 10,
+ "search_beam": 20,
+ "output_beam": 8,
+ "min_active_states": 30,
+ "max_active_states": 10000,
+ "use_double_scores": True,
+ }
+ )
+ return params
+
+
+def decode_one_batch(
+ params: AttributeDict,
+ model: nn.Module,
+ HLG: Optional[k2.Fsa],
+ H: Optional[k2.Fsa],
+ bpe_model: Optional[spm.SentencePieceProcessor],
+ batch: dict,
+ word_table: k2.SymbolTable,
+ G: Optional[k2.Fsa] = None,
+) -> Dict[str, List[List[str]]]:
+ """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.decoding_method is "1best", it uses 1best decoding without LM rescoring.
+ - params.decoding_method is "nbest", it uses nbest decoding without LM rescoring.
+ - params.decoding_method is "nbest-rescoring", it uses nbest LM rescoring.
+ - params.decoding_method is "whole-lattice-rescoring", it uses whole lattice LM
+ rescoring.
+ model:
+ The neural model.
+ HLG:
+ The decoding graph. Used only when params.decoding_method is NOT ctc-decoding.
+ H:
+ The ctc topo. Used only when params.decoding_method is ctc-decoding.
+ bpe_model:
+ The BPE model. Used only when params.decoding_method is ctc-decoding.
+ batch:
+ It is the return value from iterating
+ `lhotse.dataset.K2SpeechRecognitionDataset`. See its documentation
+ for the format of the `batch`.
+ word_table:
+ The word symbol table.
+ G:
+ An LM. It is not None when params.decoding_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. Note: If it decodes to nothing, then return None.
+ """
+ if HLG is not None:
+ device = HLG.device
+ else:
+ device = H.device
+ feature = batch["inputs"]
+ assert feature.ndim == 3
+ feature = feature.to(device)
+ # at entry, feature is (N, T, C)
+
+ supervisions = batch["supervisions"]
+ feature_lens = supervisions["num_frames"].to(device)
+
+ encoder_out, encoder_out_lens = model.encoder(feature, feature_lens)
+ nnet_output = model.ctc_output(encoder_out)
+ # 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)
+
+ if H is None:
+ assert HLG is not None
+ decoding_graph = HLG
+ else:
+ assert HLG is None
+ assert bpe_model is not None
+ decoding_graph = H
+
+ lattice = get_lattice(
+ nnet_output=nnet_output,
+ decoding_graph=decoding_graph,
+ 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.decoding_method == "ctc-decoding":
+ best_path = one_best_decoding(
+ lattice=lattice, use_double_scores=params.use_double_scores
+ )
+ # Note: `best_path.aux_labels` contains token IDs, not word IDs
+ # since we are using H, not HLG here.
+ #
+ # token_ids is a lit-of-list of IDs
+ token_ids = get_texts(best_path)
+
+ # hyps is a list of str, e.g., ['xxx yyy zzz', ...]
+ hyps = bpe_model.decode(token_ids)
+
+ # hyps is a list of list of str, e.g., [['xxx', 'yyy', 'zzz'], ... ]
+ hyps = [s.split() for s in hyps]
+ key = "ctc-decoding"
+ return {key: hyps}
+
+ if params.decoding_method == "nbest-oracle":
+ # Note: You can also pass rescored lattices to it.
+ # We choose the HLG decoded lattice for speed reasons
+ # as HLG decoding is faster and the oracle WER
+ # is only slightly worse than that of rescored lattices.
+ best_path = nbest_oracle(
+ lattice=lattice,
+ num_paths=params.num_paths,
+ ref_texts=supervisions["text"],
+ word_table=word_table,
+ nbest_scale=params.nbest_scale,
+ oov="",
+ )
+ hyps = get_texts(best_path)
+ hyps = [[word_table[i] for i in ids] for ids in hyps]
+ key = f"oracle_{params.num_paths}_nbest_scale_{params.nbest_scale}" # noqa
+ return {key: hyps}
+
+ if params.decoding_method in ["1best", "nbest"]:
+ if params.decoding_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,
+ nbest_scale=params.nbest_scale,
+ )
+ key = f"no_rescore-nbest-scale-{params.nbest_scale}-{params.num_paths}" # noqa
+
+ hyps = get_texts(best_path)
+ hyps = [[word_table[i] for i in ids] for ids in hyps]
+ return {key: hyps}
+
+ assert params.decoding_method in [
+ "nbest-rescoring",
+ "whole-lattice-rescoring",
+ ]
+
+ lm_scale_list = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]
+ 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.decoding_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,
+ nbest_scale=params.nbest_scale,
+ )
+ elif params.decoding_method == "whole-lattice-rescoring":
+ best_path_dict = rescore_with_whole_lattice(
+ lattice=lattice,
+ G_with_epsilon_loops=G,
+ lm_scale_list=lm_scale_list,
+ )
+ else:
+ assert False, f"Unsupported decoding method: {params.decoding_method}"
+
+ ans = dict()
+ if best_path_dict is not None:
+ for lm_scale_str, best_path in best_path_dict.items():
+ hyps = get_texts(best_path)
+ hyps = [[word_table[i] for i in ids] for ids in hyps]
+ ans[lm_scale_str] = hyps
+ else:
+ ans = None
+ return ans
+
+
+def decode_dataset(
+ dl: torch.utils.data.DataLoader,
+ params: AttributeDict,
+ model: nn.Module,
+ HLG: Optional[k2.Fsa],
+ H: Optional[k2.Fsa],
+ bpe_model: Optional[spm.SentencePieceProcessor],
+ word_table: k2.SymbolTable,
+ G: Optional[k2.Fsa] = None,
+) -> Dict[str, List[Tuple[str, List[str], List[str]]]]:
+ """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. Used only when params.decoding_method is NOT ctc-decoding.
+ H:
+ The ctc topo. Used only when params.decoding_method is ctc-decoding.
+ bpe_model:
+ The BPE model. Used only when params.decoding_method is ctc-decoding.
+ word_table:
+ It is the word symbol table.
+ G:
+ An LM. It is not None when params.decoding_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.
+ """
+ num_cuts = 0
+
+ try:
+ num_batches = len(dl)
+ except TypeError:
+ num_batches = "?"
+
+ results = defaultdict(list)
+ for batch_idx, batch in enumerate(dl):
+ texts = batch["supervisions"]["text"]
+ cut_ids = [cut.id for cut in batch["supervisions"]["cut"]]
+
+ hyps_dict = decode_one_batch(
+ params=params,
+ model=model,
+ HLG=HLG,
+ H=H,
+ bpe_model=bpe_model,
+ batch=batch,
+ word_table=word_table,
+ G=G,
+ )
+
+ for name, hyps in hyps_dict.items():
+ this_batch = []
+ assert len(hyps) == len(texts)
+ for cut_id, hyp_words, ref_text in zip(cut_ids, hyps, texts):
+ ref_words = ref_text.split()
+ this_batch.append((cut_id, ref_words, hyp_words))
+
+ results[name].extend(this_batch)
+
+ num_cuts += len(texts)
+
+ if batch_idx % 100 == 0:
+ batch_str = f"{batch_idx}/{num_batches}"
+ logging.info(f"batch {batch_str}, cuts processed until now is {num_cuts}")
+ return results
+
+
+def save_results(
+ params: AttributeDict,
+ test_set_name: str,
+ results_dict: Dict[str, List[Tuple[str, List[str], List[str]]]],
+):
+ test_set_wers = dict()
+ for key, results in results_dict.items():
+ recog_path = (
+ params.res_dir / f"recogs-{test_set_name}-{key}-{params.suffix}.txt"
+ )
+ results = sorted(results)
+ 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.res_dir / f"errs-{test_set_name}-{key}-{params.suffix}.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.res_dir / f"wer-summary-{test_set_name}-{key}-{params.suffix}.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()
+ args.exp_dir = Path(args.exp_dir)
+ args.lang_dir = Path(args.lang_dir)
+ args.lm_dir = Path(args.lm_dir)
+
+ params = get_params()
+ # add decoding params
+ params.update(get_decoding_params())
+ params.update(vars(args))
+
+ assert params.decoding_method in (
+ "ctc-decoding",
+ "1best",
+ "nbest",
+ "nbest-rescoring",
+ "whole-lattice-rescoring",
+ "nbest-oracle",
+ )
+ params.res_dir = params.exp_dir / params.decoding_method
+
+ if params.iter > 0:
+ params.suffix = f"iter-{params.iter}-avg-{params.avg}"
+ else:
+ params.suffix = f"epoch-{params.epoch}-avg-{params.avg}"
+
+ if params.use_averaged_model:
+ params.suffix += "-use-averaged-model"
+
+ setup_logger(f"{params.res_dir}/log-decode-{params.suffix}")
+ logging.info("Decoding started")
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"Device: {device}")
+ logging.info(params)
+
+ lexicon = Lexicon(params.lang_dir)
+ max_token_id = max(lexicon.tokens)
+ num_classes = max_token_id + 1 # +1 for the blank
+
+ params.vocab_size = num_classes
+ # and are defined in local/train_bpe_model.py
+ params.blank_id = 0
+
+ if params.decoding_method == "ctc-decoding":
+ HLG = None
+ H = k2.ctc_topo(
+ max_token=max_token_id,
+ modified=False,
+ device=device,
+ )
+ bpe_model = spm.SentencePieceProcessor()
+ bpe_model.load(str(params.lang_dir / "bpe.model"))
+ else:
+ H = None
+ bpe_model = None
+ HLG = k2.Fsa.from_dict(
+ torch.load(f"{params.lang_dir}/HLG.pt", map_location=device)
+ )
+ assert HLG.requires_grad is False
+
+ HLG.scores *= params.hlg_scale
+ if not hasattr(HLG, "lm_scores"):
+ HLG.lm_scores = HLG.scores.clone()
+
+ if params.decoding_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.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
+ # See https://github.com/k2-fsa/k2/issues/874
+ # for why we need to set G.properties to None
+ G.__dict__["_properties"] = None
+ G = k2.Fsa.from_fsas([G]).to(device)
+ G = k2.arc_sort(G)
+ # Save a dummy value so that it can be loaded in C++.
+ # See https://github.com/pytorch/pytorch/issues/67902
+ # for why we need to do this.
+ G.dummy = 1
+
+ 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", map_location=device)
+ G = k2.Fsa.from_dict(d)
+
+ if params.decoding_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
+
+ logging.info("About to create model")
+ model = get_transducer_model(params)
+
+ if not params.use_averaged_model:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ elif 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 i >= 1:
+ filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg + 1
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg + 1:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ filename_start = filenames[-1]
+ filename_end = filenames[0]
+ logging.info(
+ "Calculating the averaged model over iteration checkpoints"
+ f" from {filename_start} (excluded) to {filename_end}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+ else:
+ assert params.avg > 0, params.avg
+ start = params.epoch - params.avg
+ assert start >= 1, start
+ filename_start = f"{params.exp_dir}/epoch-{start}.pt"
+ filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
+ logging.info(
+ f"Calculating the averaged model over epoch range from "
+ f"{start} (excluded) to {params.epoch}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+
+ model.to(device)
+ model.eval()
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ # we need cut ids to display recognition results.
+ args.return_cuts = True
+ librispeech = LibriSpeechAsrDataModule(args)
+
+ test_clean_cuts = librispeech.test_clean_cuts()
+ test_other_cuts = librispeech.test_other_cuts()
+
+ test_clean_dl = librispeech.test_dataloaders(test_clean_cuts)
+ test_other_dl = librispeech.test_dataloaders(test_other_cuts)
+
+ test_sets = ["test-clean", "test-other"]
+ test_dl = [test_clean_dl, test_other_dl]
+
+ for test_set, test_dl in zip(test_sets, test_dl):
+ results_dict = decode_dataset(
+ dl=test_dl,
+ params=params,
+ model=model,
+ HLG=HLG,
+ H=H,
+ bpe_model=bpe_model,
+ word_table=lexicon.word_table,
+ 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/pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py
new file mode 100755
index 000000000..9c2166aaf
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py
@@ -0,0 +1,857 @@
+#!/usr/bin/env python3
+#
+# Copyright 2021-2022 Xiaomi Corporation (Author: Fangjun Kuang,
+# Zengwei Yao,
+# Yifan Yang,)
+#
+#
+# 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.
+"""
+Usage:
+(1) greedy search
+./pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method greedy_search
+
+(2) beam search (not recommended)
+./pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method beam_search \
+ --beam-size 4
+
+(3) modified beam search
+./pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method modified_beam_search \
+ --beam-size 4
+
+(4) fast beam search (one best)
+./pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method fast_beam_search \
+ --beam 20.0 \
+ --max-contexts 8 \
+ --max-states 64
+
+(5) fast beam search (nbest)
+./pruned_transducer_stateless7_ctc/ctc_guild_decode_bs.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc/exp \
+ --max-duration 600 \
+ --decoding-method fast_beam_search_nbest \
+ --beam 20.0 \
+ --max-contexts 8 \
+ --max-states 64 \
+ --num-paths 200 \
+ --nbest-scale 0.5
+
+(6) fast beam search (nbest oracle WER)
+./pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method fast_beam_search_nbest_oracle \
+ --beam 20.0 \
+ --max-contexts 8 \
+ --max-states 64 \
+ --num-paths 200 \
+ --nbest-scale 0.5
+
+(7) fast beam search (with LG)
+./pruned_transducer_stateless7_ctc_bs/ctc_guild_decode_bs.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method fast_beam_search_nbest_LG \
+ --beam 20.0 \
+ --max-contexts 8 \
+ --max-states 64
+"""
+
+
+import argparse
+import logging
+import math
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import k2
+import sentencepiece as spm
+import torch
+import torch.nn as nn
+from asr_datamodule import LibriSpeechAsrDataModule
+from beam_search import (
+ beam_search,
+ fast_beam_search_nbest,
+ fast_beam_search_nbest_LG,
+ fast_beam_search_nbest_oracle,
+ fast_beam_search_one_best,
+ greedy_search,
+ greedy_search_batch,
+ modified_beam_search,
+)
+from train import add_model_arguments, get_params, get_transducer_model
+from torch.nn.utils.rnn import pad_sequence
+
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from icefall.lexicon import Lexicon
+from icefall.utils import (
+ AttributeDict,
+ make_pad_mask,
+ setup_logger,
+ store_transcripts,
+ str2bool,
+ write_error_stats,
+)
+
+LOG_EPS = math.log(1e-10)
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=30,
+ help="""It specifies the checkpoint to use for decoding.
+ Note: Epoch counts from 1.
+ You can specify --avg to use more checkpoints for model averaging.""",
+ )
+
+ parser.add_argument(
+ "--iter",
+ type=int,
+ default=0,
+ help="""If positive, --epoch is ignored and it
+ will use the checkpoint exp_dir/checkpoint-iter.pt.
+ You can specify --avg to use more checkpoints for model averaging.
+ """,
+ )
+
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=9,
+ help="Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch' and '--iter'",
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=True,
+ help="Whether to load averaged model. Currently it only supports "
+ "using --epoch. If True, it would decode with the averaged model "
+ "over the epoch range from `epoch-avg` (excluded) to `epoch`."
+ "Actually only the models with epoch number of `epoch-avg` and "
+ "`epoch` are loaded for averaging. ",
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="pruned_transducer_stateless7_ctc_bs/exp",
+ help="The experiment dir",
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ default="data/lang_bpe_500/bpe.model",
+ help="Path to the BPE model",
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=Path,
+ default="data/lang_bpe_500",
+ help="The lang dir containing word table and LG graph",
+ )
+
+ parser.add_argument(
+ "--decoding-method",
+ type=str,
+ default="greedy_search",
+ help="""Possible values are:
+ - greedy_search
+ - beam_search
+ - modified_beam_search
+ - fast_beam_search
+ - fast_beam_search_nbest
+ - fast_beam_search_nbest_oracle
+ - fast_beam_search_nbest_LG
+ If you use fast_beam_search_nbest_LG, you have to specify
+ `--lang-dir`, which should contain `LG.pt`.
+ """,
+ )
+
+ parser.add_argument(
+ "--beam-size",
+ type=int,
+ default=4,
+ help="""An integer indicating how many candidates we will keep for each
+ frame. Used only when --decoding-method is beam_search or
+ modified_beam_search.""",
+ )
+
+ parser.add_argument(
+ "--beam",
+ type=float,
+ default=20.0,
+ help="""A floating point value to calculate the cutoff score during beam
+ search (i.e., `cutoff = max-score - beam`), which is the same as the
+ `beam` in Kaldi.
+ Used only when --decoding-method is fast_beam_search,
+ fast_beam_search_nbest, fast_beam_search_nbest_LG,
+ and fast_beam_search_nbest_oracle
+ """,
+ )
+
+ parser.add_argument(
+ "--ngram-lm-scale",
+ type=float,
+ default=0.01,
+ help="""
+ Used only when --decoding_method is fast_beam_search_nbest_LG.
+ It specifies the scale for n-gram LM scores.
+ """,
+ )
+
+ parser.add_argument(
+ "--max-contexts",
+ type=int,
+ default=8,
+ help="""Used only when --decoding-method is
+ fast_beam_search, fast_beam_search_nbest, fast_beam_search_nbest_LG,
+ and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--max-states",
+ type=int,
+ default=64,
+ help="""Used only when --decoding-method is
+ fast_beam_search, fast_beam_search_nbest, fast_beam_search_nbest_LG,
+ and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram, 2 means tri-gram",
+ )
+ parser.add_argument(
+ "--max-sym-per-frame",
+ type=int,
+ default=1,
+ help="""Maximum number of symbols per frame.
+ Used only when --decoding_method is greedy_search""",
+ )
+
+ parser.add_argument(
+ "--num-paths",
+ type=int,
+ default=200,
+ help="""Number of paths for nbest decoding.
+ Used only when the decoding method is fast_beam_search_nbest,
+ fast_beam_search_nbest_LG, and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--nbest-scale",
+ type=float,
+ default=0.5,
+ help="""Scale applied to lattice scores when computing nbest paths.
+ Used only when the decoding method is fast_beam_search_nbest,
+ fast_beam_search_nbest_LG, and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--simulate-streaming",
+ type=str2bool,
+ default=False,
+ help="""Whether to simulate streaming in decoding, this is a good way to
+ test a streaming model.
+ """,
+ )
+
+ parser.add_argument(
+ "--decode-chunk-size",
+ type=int,
+ default=16,
+ help="The chunk size for decoding (in frames after subsampling)",
+ )
+
+ parser.add_argument(
+ "--left-context",
+ type=int,
+ default=64,
+ help="left context can be seen during decoding (in frames after subsampling)",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def decode_one_batch(
+ params: AttributeDict,
+ model: nn.Module,
+ sp: spm.SentencePieceProcessor,
+ batch: dict,
+ word_table: Optional[k2.SymbolTable] = None,
+ decoding_graph: Optional[k2.Fsa] = None,
+) -> Dict[str, List[List[str]]]:
+ """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 greedy_search is used, it would be "greedy_search"
+ If beam search with a beam size of 7 is used, it would be
+ "beam_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`.
+ model:
+ The neural model.
+ sp:
+ The BPE model.
+ batch:
+ It is the return value from iterating
+ `lhotse.dataset.K2SpeechRecognitionDataset`. See its documentation
+ for the format of the `batch`.
+ word_table:
+ The word symbol table.
+ decoding_graph:
+ The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+ only when --decoding_method is fast_beam_search, fast_beam_search_nbest,
+ fast_beam_search_nbest_oracle, and fast_beam_search_nbest_LG.
+ Returns:
+ Return the decoding result. See above description for the format of
+ the returned dict.
+ """
+ device = next(model.parameters()).device
+ feature = batch["inputs"]
+ assert feature.ndim == 3
+
+ feature = feature.to(device)
+ # at entry, feature is (N, T, C)
+
+ supervisions = batch["supervisions"]
+ feature_lens = supervisions["num_frames"].to(device)
+
+ if params.simulate_streaming:
+ feature_lens += params.left_context
+ feature = torch.nn.functional.pad(
+ feature,
+ pad=(0, 0, 0, params.left_context),
+ value=LOG_EPS,
+ )
+ encoder_out, encoder_out_lens, _ = model.encoder.streaming_forward(
+ x=feature,
+ x_lens=feature_lens,
+ chunk_size=params.decode_chunk_size,
+ left_context=params.left_context,
+ simulate_streaming=True,
+ )
+ else:
+ encoder_out, encoder_out_lens = model.encoder(x=feature, x_lens=feature_lens)
+
+ # filter out blank frames using ctc outputs
+ ctc_output = model.ctc_output(encoder_out)
+ encoder_out = model.lconv(
+ x=encoder_out,
+ src_key_padding_mask=make_pad_mask(encoder_out_lens),
+ )
+ encoder_out, encoder_out_lens = model.frame_reducer(
+ x=encoder_out,
+ x_lens=encoder_out_lens,
+ ctc_output=ctc_output,
+ blank_id=0,
+ )
+
+ hyps = []
+
+ if params.decoding_method == "fast_beam_search":
+ hyp_tokens = fast_beam_search_one_best(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.decoding_method == "fast_beam_search_nbest_LG":
+ hyp_tokens = fast_beam_search_nbest_LG(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ num_paths=params.num_paths,
+ nbest_scale=params.nbest_scale,
+ )
+ for hyp in hyp_tokens:
+ hyps.append([word_table[i] for i in hyp])
+ elif params.decoding_method == "fast_beam_search_nbest":
+ hyp_tokens = fast_beam_search_nbest(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ num_paths=params.num_paths,
+ nbest_scale=params.nbest_scale,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.decoding_method == "fast_beam_search_nbest_oracle":
+ hyp_tokens = fast_beam_search_nbest_oracle(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ num_paths=params.num_paths,
+ ref_texts=sp.encode(supervisions["text"]),
+ nbest_scale=params.nbest_scale,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.decoding_method == "greedy_search" and params.max_sym_per_frame == 1:
+ hyp_tokens = greedy_search_batch(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.decoding_method == "modified_beam_search":
+ hyp_tokens = modified_beam_search(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam_size,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ else:
+ batch_size = encoder_out.size(0)
+
+ for i in range(batch_size):
+ # fmt: off
+ encoder_out_i = encoder_out[i:i+1, :encoder_out_lens[i]]
+ # fmt: on
+ if params.decoding_method == "greedy_search":
+ hyp = greedy_search(
+ model=model,
+ encoder_out=encoder_out_i,
+ max_sym_per_frame=params.max_sym_per_frame,
+ )
+ elif params.decoding_method == "beam_search":
+ hyp = beam_search(
+ model=model,
+ encoder_out=encoder_out_i,
+ beam=params.beam_size,
+ )
+ else:
+ raise ValueError(
+ f"Unsupported decoding method: {params.decoding_method}"
+ )
+ hyps.append(sp.decode(hyp).split())
+
+ if params.decoding_method == "greedy_search":
+ return {"greedy_search": hyps}
+ elif "fast_beam_search" in params.decoding_method:
+ key = f"beam_{params.beam}_"
+ key += f"max_contexts_{params.max_contexts}_"
+ key += f"max_states_{params.max_states}"
+ if "nbest" in params.decoding_method:
+ key += f"_num_paths_{params.num_paths}_"
+ key += f"nbest_scale_{params.nbest_scale}"
+ if "LG" in params.decoding_method:
+ key += f"_ngram_lm_scale_{params.ngram_lm_scale}"
+
+ return {key: hyps}
+ else:
+ return {f"beam_size_{params.beam_size}": hyps}
+
+
+def decode_dataset(
+ dl: torch.utils.data.DataLoader,
+ params: AttributeDict,
+ model: nn.Module,
+ sp: spm.SentencePieceProcessor,
+ word_table: Optional[k2.SymbolTable] = None,
+ decoding_graph: Optional[k2.Fsa] = None,
+) -> Dict[str, List[Tuple[str, List[str], List[str]]]]:
+ """Decode dataset.
+
+ Args:
+ dl:
+ PyTorch's dataloader containing the dataset to decode.
+ params:
+ It is returned by :func:`get_params`.
+ model:
+ The neural model.
+ sp:
+ The BPE model.
+ word_table:
+ The word symbol table.
+ decoding_graph:
+ The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+ only when --decoding_method is fast_beam_search, fast_beam_search_nbest,
+ fast_beam_search_nbest_oracle, and fast_beam_search_nbest_LG.
+ Returns:
+ Return a dict, whose key may be "greedy_search" if greedy search
+ is used, or it may be "beam_7" if beam size of 7 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.
+ """
+ num_cuts = 0
+
+ try:
+ num_batches = len(dl)
+ except TypeError:
+ num_batches = "?"
+
+ if params.decoding_method == "greedy_search":
+ log_interval = 50
+ else:
+ log_interval = 20
+
+ results = defaultdict(list)
+ for batch_idx, batch in enumerate(dl):
+ texts = batch["supervisions"]["text"]
+ cut_ids = [cut.id for cut in batch["supervisions"]["cut"]]
+
+ hyps_dict = decode_one_batch(
+ params=params,
+ model=model,
+ sp=sp,
+ decoding_graph=decoding_graph,
+ word_table=word_table,
+ batch=batch,
+ )
+
+ for name, hyps in hyps_dict.items():
+ this_batch = []
+ assert len(hyps) == len(texts)
+ for cut_id, hyp_words, ref_text in zip(cut_ids, hyps, texts):
+ ref_words = ref_text.split()
+ this_batch.append((cut_id, ref_words, hyp_words))
+
+ results[name].extend(this_batch)
+
+ num_cuts += len(texts)
+
+ if batch_idx % log_interval == 0:
+ batch_str = f"{batch_idx}/{num_batches}"
+ logging.info(f"batch {batch_str}, cuts processed until now is {num_cuts}")
+ return results
+
+
+def save_results(
+ params: AttributeDict,
+ test_set_name: str,
+ results_dict: Dict[str, List[Tuple[str, List[str], List[str]]]],
+):
+ test_set_wers = dict()
+ for key, results in results_dict.items():
+ recog_path = (
+ params.res_dir / f"recogs-{test_set_name}-{key}-{params.suffix}.txt"
+ )
+ results = sorted(results)
+ 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.res_dir / f"errs-{test_set_name}-{key}-{params.suffix}.txt"
+ )
+ with open(errs_filename, "w") as f:
+ wer = write_error_stats(
+ f, f"{test_set_name}-{key}", results, enable_log=True
+ )
+ 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.res_dir / f"wer-summary-{test_set_name}-{key}-{params.suffix}.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()
+ args.exp_dir = Path(args.exp_dir)
+
+ params = get_params()
+ params.update(vars(args))
+
+ assert params.decoding_method in (
+ "greedy_search",
+ "beam_search",
+ "fast_beam_search",
+ "fast_beam_search_nbest",
+ "fast_beam_search_nbest_LG",
+ "fast_beam_search_nbest_oracle",
+ "modified_beam_search",
+ )
+ params.res_dir = params.exp_dir / params.decoding_method
+
+ if params.iter > 0:
+ params.suffix = f"iter-{params.iter}-avg-{params.avg}"
+ else:
+ params.suffix = f"epoch-{params.epoch}-avg-{params.avg}"
+
+ if params.simulate_streaming:
+ params.suffix += f"-streaming-chunk-size-{params.decode_chunk_size}"
+ params.suffix += f"-left-context-{params.left_context}"
+
+ if "fast_beam_search" in params.decoding_method:
+ params.suffix += f"-beam-{params.beam}"
+ params.suffix += f"-max-contexts-{params.max_contexts}"
+ params.suffix += f"-max-states-{params.max_states}"
+ if "nbest" in params.decoding_method:
+ params.suffix += f"-nbest-scale-{params.nbest_scale}"
+ params.suffix += f"-num-paths-{params.num_paths}"
+ if "LG" in params.decoding_method:
+ params.suffix += f"-ngram-lm-scale-{params.ngram_lm_scale}"
+ elif "beam_search" in params.decoding_method:
+ params.suffix += f"-{params.decoding_method}-beam-size-{params.beam_size}"
+ else:
+ params.suffix += f"-context-{params.context_size}"
+ params.suffix += f"-max-sym-per-frame-{params.max_sym_per_frame}"
+
+ if params.use_averaged_model:
+ params.suffix += "-use-averaged-model"
+
+ setup_logger(f"{params.res_dir}/log-decode-{params.suffix}")
+ logging.info("Decoding started")
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"Device: {device}")
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(params.bpe_model)
+
+ # and are defined in local/train_bpe_model.py
+ params.blank_id = sp.piece_to_id("")
+ params.unk_id = sp.piece_to_id("")
+ params.vocab_size = sp.get_piece_size()
+
+ if params.simulate_streaming:
+ assert (
+ params.causal_convolution
+ ), "Decoding in streaming requires causal convolution"
+
+ logging.info(params)
+
+ logging.info("About to create model")
+ model = get_transducer_model(params)
+
+ if not params.use_averaged_model:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ elif 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 i >= 1:
+ filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg + 1
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg + 1:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ filename_start = filenames[-1]
+ filename_end = filenames[0]
+ logging.info(
+ "Calculating the averaged model over iteration checkpoints"
+ f" from {filename_start} (excluded) to {filename_end}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+ else:
+ assert params.avg > 0, params.avg
+ start = params.epoch - params.avg
+ assert start >= 1, start
+ filename_start = f"{params.exp_dir}/epoch-{start}.pt"
+ filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
+ logging.info(
+ f"Calculating the averaged model over epoch range from "
+ f"{start} (excluded) to {params.epoch}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+
+ model.to(device)
+ model.eval()
+
+ if "fast_beam_search" in params.decoding_method:
+ if params.decoding_method == "fast_beam_search_nbest_LG":
+ lexicon = Lexicon(params.lang_dir)
+ word_table = lexicon.word_table
+ lg_filename = params.lang_dir / "LG.pt"
+ logging.info(f"Loading {lg_filename}")
+ decoding_graph = k2.Fsa.from_dict(
+ torch.load(lg_filename, map_location=device)
+ )
+ decoding_graph.scores *= params.ngram_lm_scale
+ else:
+ word_table = None
+ decoding_graph = k2.trivial_graph(params.vocab_size - 1, device=device)
+ else:
+ decoding_graph = None
+ word_table = None
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ # we need cut ids to display recognition results.
+ args.return_cuts = True
+ librispeech = LibriSpeechAsrDataModule(args)
+
+ test_clean_cuts = librispeech.test_clean_cuts()
+ test_other_cuts = librispeech.test_other_cuts()
+
+ test_clean_dl = librispeech.test_dataloaders(test_clean_cuts)
+ test_other_dl = librispeech.test_dataloaders(test_other_cuts)
+
+ test_sets = ["test-clean", "test-other"]
+ test_dl = [test_clean_dl, test_other_dl]
+
+ for test_set, test_dl in zip(test_sets, test_dl):
+ results_dict = decode_dataset(
+ dl=test_dl,
+ params=params,
+ model=model,
+ sp=sp,
+ word_table=word_table,
+ decoding_graph=decoding_graph,
+ )
+
+ 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/pruned_transducer_stateless7_ctc_bs/decode.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/decode.py
new file mode 100755
index 000000000..ce45a4beb
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/decode.py
@@ -0,0 +1,841 @@
+#!/usr/bin/env python3
+#
+# Copyright 2021-2022 Xiaomi Corporation (Author: Fangjun Kuang,
+# Zengwei Yao)
+#
+# 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.
+"""
+Usage:
+(1) greedy search
+./pruned_transducer_stateless7_ctc_bs/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method greedy_search
+
+(2) beam search (not recommended)
+./pruned_transducer_stateless7_ctc_bs/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method beam_search \
+ --beam-size 4
+
+(3) modified beam search
+./pruned_transducer_stateless7_ctc_bs/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method modified_beam_search \
+ --beam-size 4
+
+(4) fast beam search (one best)
+./pruned_transducer_stateless7_ctc_bs/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method fast_beam_search \
+ --beam 20.0 \
+ --max-contexts 8 \
+ --max-states 64
+
+(5) fast beam search (nbest)
+./pruned_transducer_stateless7_ctc_bs/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method fast_beam_search_nbest \
+ --beam 20.0 \
+ --max-contexts 8 \
+ --max-states 64 \
+ --num-paths 200 \
+ --nbest-scale 0.5
+
+(6) fast beam search (nbest oracle WER)
+./pruned_transducer_stateless7_ctc_bs/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method fast_beam_search_nbest_oracle \
+ --beam 20.0 \
+ --max-contexts 8 \
+ --max-states 64 \
+ --num-paths 200 \
+ --nbest-scale 0.5
+
+(7) fast beam search (with LG)
+./pruned_transducer_stateless7_ctc_bs/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --max-duration 600 \
+ --decoding-method fast_beam_search_nbest_LG \
+ --beam 20.0 \
+ --max-contexts 8 \
+ --max-states 64
+"""
+
+
+import argparse
+import logging
+import math
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import k2
+import sentencepiece as spm
+import torch
+import torch.nn as nn
+from asr_datamodule import LibriSpeechAsrDataModule
+from beam_search import (
+ beam_search,
+ fast_beam_search_nbest,
+ fast_beam_search_nbest_LG,
+ fast_beam_search_nbest_oracle,
+ fast_beam_search_one_best,
+ greedy_search,
+ greedy_search_batch,
+ modified_beam_search,
+)
+from train import add_model_arguments, get_params, get_transducer_model
+
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from icefall.lexicon import Lexicon
+from icefall.utils import (
+ AttributeDict,
+ setup_logger,
+ store_transcripts,
+ str2bool,
+ write_error_stats,
+)
+
+LOG_EPS = math.log(1e-10)
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=30,
+ help="""It specifies the checkpoint to use for decoding.
+ Note: Epoch counts from 1.
+ You can specify --avg to use more checkpoints for model averaging.""",
+ )
+
+ parser.add_argument(
+ "--iter",
+ type=int,
+ default=0,
+ help="""If positive, --epoch is ignored and it
+ will use the checkpoint exp_dir/checkpoint-iter.pt.
+ You can specify --avg to use more checkpoints for model averaging.
+ """,
+ )
+
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=9,
+ help="Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch' and '--iter'",
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=True,
+ help="Whether to load averaged model. Currently it only supports "
+ "using --epoch. If True, it would decode with the averaged model "
+ "over the epoch range from `epoch-avg` (excluded) to `epoch`."
+ "Actually only the models with epoch number of `epoch-avg` and "
+ "`epoch` are loaded for averaging. ",
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="pruned_transducer_stateless7_ctc_bs/exp",
+ help="The experiment dir",
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ default="data/lang_bpe_500/bpe.model",
+ help="Path to the BPE model",
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=Path,
+ default="data/lang_bpe_500",
+ help="The lang dir containing word table and LG graph",
+ )
+
+ parser.add_argument(
+ "--decoding-method",
+ type=str,
+ default="greedy_search",
+ help="""Possible values are:
+ - greedy_search
+ - beam_search
+ - modified_beam_search
+ - fast_beam_search
+ - fast_beam_search_nbest
+ - fast_beam_search_nbest_oracle
+ - fast_beam_search_nbest_LG
+ If you use fast_beam_search_nbest_LG, you have to specify
+ `--lang-dir`, which should contain `LG.pt`.
+ """,
+ )
+
+ parser.add_argument(
+ "--beam-size",
+ type=int,
+ default=4,
+ help="""An integer indicating how many candidates we will keep for each
+ frame. Used only when --decoding-method is beam_search or
+ modified_beam_search.""",
+ )
+
+ parser.add_argument(
+ "--beam",
+ type=float,
+ default=20.0,
+ help="""A floating point value to calculate the cutoff score during beam
+ search (i.e., `cutoff = max-score - beam`), which is the same as the
+ `beam` in Kaldi.
+ Used only when --decoding-method is fast_beam_search,
+ fast_beam_search_nbest, fast_beam_search_nbest_LG,
+ and fast_beam_search_nbest_oracle
+ """,
+ )
+
+ parser.add_argument(
+ "--ngram-lm-scale",
+ type=float,
+ default=0.01,
+ help="""
+ Used only when --decoding_method is fast_beam_search_nbest_LG.
+ It specifies the scale for n-gram LM scores.
+ """,
+ )
+
+ parser.add_argument(
+ "--max-contexts",
+ type=int,
+ default=8,
+ help="""Used only when --decoding-method is
+ fast_beam_search, fast_beam_search_nbest, fast_beam_search_nbest_LG,
+ and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--max-states",
+ type=int,
+ default=64,
+ help="""Used only when --decoding-method is
+ fast_beam_search, fast_beam_search_nbest, fast_beam_search_nbest_LG,
+ and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram, 2 means tri-gram",
+ )
+ parser.add_argument(
+ "--max-sym-per-frame",
+ type=int,
+ default=1,
+ help="""Maximum number of symbols per frame.
+ Used only when --decoding_method is greedy_search""",
+ )
+
+ parser.add_argument(
+ "--num-paths",
+ type=int,
+ default=200,
+ help="""Number of paths for nbest decoding.
+ Used only when the decoding method is fast_beam_search_nbest,
+ fast_beam_search_nbest_LG, and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--nbest-scale",
+ type=float,
+ default=0.5,
+ help="""Scale applied to lattice scores when computing nbest paths.
+ Used only when the decoding method is fast_beam_search_nbest,
+ fast_beam_search_nbest_LG, and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--simulate-streaming",
+ type=str2bool,
+ default=False,
+ help="""Whether to simulate streaming in decoding, this is a good way to
+ test a streaming model.
+ """,
+ )
+
+ parser.add_argument(
+ "--decode-chunk-size",
+ type=int,
+ default=16,
+ help="The chunk size for decoding (in frames after subsampling)",
+ )
+
+ parser.add_argument(
+ "--left-context",
+ type=int,
+ default=64,
+ help="left context can be seen during decoding (in frames after subsampling)",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def decode_one_batch(
+ params: AttributeDict,
+ model: nn.Module,
+ sp: spm.SentencePieceProcessor,
+ batch: dict,
+ word_table: Optional[k2.SymbolTable] = None,
+ decoding_graph: Optional[k2.Fsa] = None,
+) -> Dict[str, List[List[str]]]:
+ """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 greedy_search is used, it would be "greedy_search"
+ If beam search with a beam size of 7 is used, it would be
+ "beam_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`.
+ model:
+ The neural model.
+ sp:
+ The BPE model.
+ batch:
+ It is the return value from iterating
+ `lhotse.dataset.K2SpeechRecognitionDataset`. See its documentation
+ for the format of the `batch`.
+ word_table:
+ The word symbol table.
+ decoding_graph:
+ The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+ only when --decoding_method is fast_beam_search, fast_beam_search_nbest,
+ fast_beam_search_nbest_oracle, and fast_beam_search_nbest_LG.
+ Returns:
+ Return the decoding result. See above description for the format of
+ the returned dict.
+ """
+ device = next(model.parameters()).device
+ feature = batch["inputs"]
+ assert feature.ndim == 3
+
+ feature = feature.to(device)
+ # at entry, feature is (N, T, C)
+
+ supervisions = batch["supervisions"]
+ feature_lens = supervisions["num_frames"].to(device)
+
+ if params.simulate_streaming:
+ feature_lens += params.left_context
+ feature = torch.nn.functional.pad(
+ feature,
+ pad=(0, 0, 0, params.left_context),
+ value=LOG_EPS,
+ )
+ encoder_out, encoder_out_lens, _ = model.encoder.streaming_forward(
+ x=feature,
+ x_lens=feature_lens,
+ chunk_size=params.decode_chunk_size,
+ left_context=params.left_context,
+ simulate_streaming=True,
+ )
+ else:
+ encoder_out, encoder_out_lens = model.encoder(x=feature, x_lens=feature_lens)
+
+ hyps = []
+
+ if params.decoding_method == "fast_beam_search":
+ hyp_tokens = fast_beam_search_one_best(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.decoding_method == "fast_beam_search_nbest_LG":
+ hyp_tokens = fast_beam_search_nbest_LG(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ num_paths=params.num_paths,
+ nbest_scale=params.nbest_scale,
+ )
+ for hyp in hyp_tokens:
+ hyps.append([word_table[i] for i in hyp])
+ elif params.decoding_method == "fast_beam_search_nbest":
+ hyp_tokens = fast_beam_search_nbest(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ num_paths=params.num_paths,
+ nbest_scale=params.nbest_scale,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.decoding_method == "fast_beam_search_nbest_oracle":
+ hyp_tokens = fast_beam_search_nbest_oracle(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ num_paths=params.num_paths,
+ ref_texts=sp.encode(supervisions["text"]),
+ nbest_scale=params.nbest_scale,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.decoding_method == "greedy_search" and params.max_sym_per_frame == 1:
+ hyp_tokens = greedy_search_batch(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.decoding_method == "modified_beam_search":
+ hyp_tokens = modified_beam_search(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam_size,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ else:
+ batch_size = encoder_out.size(0)
+
+ for i in range(batch_size):
+ # fmt: off
+ encoder_out_i = encoder_out[i:i+1, :encoder_out_lens[i]]
+ # fmt: on
+ if params.decoding_method == "greedy_search":
+ hyp = greedy_search(
+ model=model,
+ encoder_out=encoder_out_i,
+ max_sym_per_frame=params.max_sym_per_frame,
+ )
+ elif params.decoding_method == "beam_search":
+ hyp = beam_search(
+ model=model,
+ encoder_out=encoder_out_i,
+ beam=params.beam_size,
+ )
+ else:
+ raise ValueError(
+ f"Unsupported decoding method: {params.decoding_method}"
+ )
+ hyps.append(sp.decode(hyp).split())
+
+ if params.decoding_method == "greedy_search":
+ return {"greedy_search": hyps}
+ elif "fast_beam_search" in params.decoding_method:
+ key = f"beam_{params.beam}_"
+ key += f"max_contexts_{params.max_contexts}_"
+ key += f"max_states_{params.max_states}"
+ if "nbest" in params.decoding_method:
+ key += f"_num_paths_{params.num_paths}_"
+ key += f"nbest_scale_{params.nbest_scale}"
+ if "LG" in params.decoding_method:
+ key += f"_ngram_lm_scale_{params.ngram_lm_scale}"
+
+ return {key: hyps}
+ else:
+ return {f"beam_size_{params.beam_size}": hyps}
+
+
+def decode_dataset(
+ dl: torch.utils.data.DataLoader,
+ params: AttributeDict,
+ model: nn.Module,
+ sp: spm.SentencePieceProcessor,
+ word_table: Optional[k2.SymbolTable] = None,
+ decoding_graph: Optional[k2.Fsa] = None,
+) -> Dict[str, List[Tuple[str, List[str], List[str]]]]:
+ """Decode dataset.
+
+ Args:
+ dl:
+ PyTorch's dataloader containing the dataset to decode.
+ params:
+ It is returned by :func:`get_params`.
+ model:
+ The neural model.
+ sp:
+ The BPE model.
+ word_table:
+ The word symbol table.
+ decoding_graph:
+ The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+ only when --decoding_method is fast_beam_search, fast_beam_search_nbest,
+ fast_beam_search_nbest_oracle, and fast_beam_search_nbest_LG.
+ Returns:
+ Return a dict, whose key may be "greedy_search" if greedy search
+ is used, or it may be "beam_7" if beam size of 7 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.
+ """
+ num_cuts = 0
+
+ try:
+ num_batches = len(dl)
+ except TypeError:
+ num_batches = "?"
+
+ if params.decoding_method == "greedy_search":
+ log_interval = 50
+ else:
+ log_interval = 20
+
+ results = defaultdict(list)
+ for batch_idx, batch in enumerate(dl):
+ texts = batch["supervisions"]["text"]
+ cut_ids = [cut.id for cut in batch["supervisions"]["cut"]]
+
+ hyps_dict = decode_one_batch(
+ params=params,
+ model=model,
+ sp=sp,
+ decoding_graph=decoding_graph,
+ word_table=word_table,
+ batch=batch,
+ )
+
+ for name, hyps in hyps_dict.items():
+ this_batch = []
+ assert len(hyps) == len(texts)
+ for cut_id, hyp_words, ref_text in zip(cut_ids, hyps, texts):
+ ref_words = ref_text.split()
+ this_batch.append((cut_id, ref_words, hyp_words))
+
+ results[name].extend(this_batch)
+
+ num_cuts += len(texts)
+
+ if batch_idx % log_interval == 0:
+ batch_str = f"{batch_idx}/{num_batches}"
+
+ logging.info(f"batch {batch_str}, cuts processed until now is {num_cuts}")
+ return results
+
+
+def save_results(
+ params: AttributeDict,
+ test_set_name: str,
+ results_dict: Dict[str, List[Tuple[str, List[str], List[str]]]],
+):
+ test_set_wers = dict()
+ for key, results in results_dict.items():
+ recog_path = (
+ params.res_dir / f"recogs-{test_set_name}-{key}-{params.suffix}.txt"
+ )
+ results = sorted(results)
+ 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.res_dir / f"errs-{test_set_name}-{key}-{params.suffix}.txt"
+ )
+ with open(errs_filename, "w") as f:
+ wer = write_error_stats(
+ f, f"{test_set_name}-{key}", results, enable_log=True
+ )
+ 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.res_dir / f"wer-summary-{test_set_name}-{key}-{params.suffix}.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()
+ args.exp_dir = Path(args.exp_dir)
+
+ params = get_params()
+ params.update(vars(args))
+
+ assert params.decoding_method in (
+ "greedy_search",
+ "beam_search",
+ "fast_beam_search",
+ "fast_beam_search_nbest",
+ "fast_beam_search_nbest_LG",
+ "fast_beam_search_nbest_oracle",
+ "modified_beam_search",
+ )
+ params.res_dir = params.exp_dir / params.decoding_method
+
+ if params.iter > 0:
+ params.suffix = f"iter-{params.iter}-avg-{params.avg}"
+ else:
+ params.suffix = f"epoch-{params.epoch}-avg-{params.avg}"
+
+ if params.simulate_streaming:
+ params.suffix += f"-streaming-chunk-size-{params.decode_chunk_size}"
+ params.suffix += f"-left-context-{params.left_context}"
+
+ if "fast_beam_search" in params.decoding_method:
+ params.suffix += f"-beam-{params.beam}"
+ params.suffix += f"-max-contexts-{params.max_contexts}"
+ params.suffix += f"-max-states-{params.max_states}"
+ if "nbest" in params.decoding_method:
+ params.suffix += f"-nbest-scale-{params.nbest_scale}"
+ params.suffix += f"-num-paths-{params.num_paths}"
+ if "LG" in params.decoding_method:
+ params.suffix += f"-ngram-lm-scale-{params.ngram_lm_scale}"
+ elif "beam_search" in params.decoding_method:
+ params.suffix += f"-{params.decoding_method}-beam-size-{params.beam_size}"
+ else:
+ params.suffix += f"-context-{params.context_size}"
+ params.suffix += f"-max-sym-per-frame-{params.max_sym_per_frame}"
+
+ if params.use_averaged_model:
+ params.suffix += "-use-averaged-model"
+
+ setup_logger(f"{params.res_dir}/log-decode-{params.suffix}")
+ logging.info("Decoding started")
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"Device: {device}")
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(params.bpe_model)
+
+ # and are defined in local/train_bpe_model.py
+ params.blank_id = sp.piece_to_id("")
+ params.unk_id = sp.piece_to_id("")
+ params.vocab_size = sp.get_piece_size()
+
+ if params.simulate_streaming:
+ assert (
+ params.causal_convolution
+ ), "Decoding in streaming requires causal convolution"
+
+ logging.info(params)
+
+ logging.info("About to create model")
+ model = get_transducer_model(params)
+
+ if not params.use_averaged_model:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ elif 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 i >= 1:
+ filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg + 1
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg + 1:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ filename_start = filenames[-1]
+ filename_end = filenames[0]
+ logging.info(
+ "Calculating the averaged model over iteration checkpoints"
+ f" from {filename_start} (excluded) to {filename_end}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+ else:
+ assert params.avg > 0, params.avg
+ start = params.epoch - params.avg
+ assert start >= 1, start
+ filename_start = f"{params.exp_dir}/epoch-{start}.pt"
+ filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
+ logging.info(
+ f"Calculating the averaged model over epoch range from "
+ f"{start} (excluded) to {params.epoch}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+
+ model.to(device)
+ model.eval()
+
+ if "fast_beam_search" in params.decoding_method:
+ if params.decoding_method == "fast_beam_search_nbest_LG":
+ lexicon = Lexicon(params.lang_dir)
+ word_table = lexicon.word_table
+ lg_filename = params.lang_dir / "LG.pt"
+ logging.info(f"Loading {lg_filename}")
+ decoding_graph = k2.Fsa.from_dict(
+ torch.load(lg_filename, map_location=device)
+ )
+ decoding_graph.scores *= params.ngram_lm_scale
+ else:
+ word_table = None
+ decoding_graph = k2.trivial_graph(params.vocab_size - 1, device=device)
+ else:
+ decoding_graph = None
+ word_table = None
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ # we need cut ids to display recognition results.
+ args.return_cuts = True
+ librispeech = LibriSpeechAsrDataModule(args)
+
+ test_clean_cuts = librispeech.test_clean_cuts()
+ test_other_cuts = librispeech.test_other_cuts()
+
+ test_clean_dl = librispeech.test_dataloaders(test_clean_cuts)
+ test_other_dl = librispeech.test_dataloaders(test_other_cuts)
+
+ test_sets = ["test-clean", "test-other"]
+ test_dl = [test_clean_dl, test_other_dl]
+
+ for test_set, test_dl in zip(test_sets, test_dl):
+ results_dict = decode_dataset(
+ dl=test_dl,
+ params=params,
+ model=model,
+ sp=sp,
+ word_table=word_table,
+ decoding_graph=decoding_graph,
+ )
+
+ 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/pruned_transducer_stateless7_ctc_bs/decoder.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/decoder.py
new file mode 120000
index 000000000..33944d0d2
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/decoder.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/decoder.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/encoder_interface.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/encoder_interface.py
new file mode 120000
index 000000000..b9aa0ae08
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/encoder_interface.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/encoder_interface.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/export.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/export.py
new file mode 100755
index 000000000..96d316604
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/export.py
@@ -0,0 +1,319 @@
+#!/usr/bin/env python3
+#
+# Copyright 2021 Xiaomi Corporation (Author: Fangjun Kuang)
+#
+# 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.
+
+# This script converts several saved checkpoints
+# to a single one using model averaging.
+"""
+
+Usage:
+
+(1) Export to torchscript model using torch.jit.script()
+
+./pruned_transducer_stateless7_ctc_bs/export.py \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 13 \
+ --jit 1
+
+It will generate a file `cpu_jit.pt` in the given `exp_dir`. You can later
+load it by `torch.jit.load("cpu_jit.pt")`.
+
+Note `cpu` in the name `cpu_jit.pt` means the parameters when loaded into Python
+are on CPU. You can use `to("cuda")` to move them to a CUDA device.
+
+Check
+https://github.com/k2-fsa/sherpa
+for how to use the exported models outside of icefall.
+
+(2) Export `model.state_dict()`
+
+./pruned_transducer_stateless7_ctc_bs/export.py \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 13
+
+It will generate a file `pretrained.pt` in the given `exp_dir`. You can later
+load it by `icefall.checkpoint.load_checkpoint()`.
+
+To use the generated file with `pruned_transducer_stateless7_ctc_bs/decode.py`,
+you can do:
+
+ cd /path/to/exp_dir
+ ln -s pretrained.pt epoch-9999.pt
+
+ cd /path/to/egs/librispeech/ASR
+ ./pruned_transducer_stateless7_ctc_bs/decode.py \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --epoch 9999 \
+ --avg 1 \
+ --max-duration 600 \
+ --decoding-method greedy_search \
+ --bpe-model data/lang_bpe_500/bpe.model
+
+Check ./pretrained.py for its usage.
+
+Note: If you don't want to train a model from scratch, we have
+provided one for you. You can get it at
+
+https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless7-2022-11-11
+
+with the following commands:
+
+ sudo apt-get install git-lfs
+ git lfs install
+ git clone https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless7-2022-11-11
+ # You will find the pre-trained model in icefall-asr-librispeech-pruned-transducer-stateless7-2022-11-11/exp
+"""
+
+import argparse
+import logging
+from pathlib import Path
+
+import sentencepiece as spm
+import torch
+from scaling_converter import convert_scaled_to_non_scaled
+from train import add_model_arguments, get_params, get_transducer_model
+
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from icefall.utils import str2bool
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=30,
+ help="""It specifies the checkpoint to use for decoding.
+ Note: Epoch counts from 1.
+ You can specify --avg to use more checkpoints for model averaging.""",
+ )
+
+ parser.add_argument(
+ "--iter",
+ type=int,
+ default=0,
+ help="""If positive, --epoch is ignored and it
+ will use the checkpoint exp_dir/checkpoint-iter.pt.
+ You can specify --avg to use more checkpoints for model averaging.
+ """,
+ )
+
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=9,
+ help="Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch' and '--iter'",
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=True,
+ help="Whether to load averaged model. Currently it only supports "
+ "using --epoch. If True, it would decode with the averaged model "
+ "over the epoch range from `epoch-avg` (excluded) to `epoch`."
+ "Actually only the models with epoch number of `epoch-avg` and "
+ "`epoch` are loaded for averaging. ",
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="pruned_transducer_stateless7/exp",
+ help="""It specifies the directory where all training related
+ files, e.g., checkpoints, log, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ default="data/lang_bpe_500/bpe.model",
+ help="Path to the BPE model",
+ )
+
+ parser.add_argument(
+ "--jit",
+ type=str2bool,
+ default=False,
+ help="""True to save a model after applying torch.jit.script.
+ It will generate a file named cpu_jit.pt
+
+ Check ./jit_pretrained.py for how to use it.
+ """,
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram, 2 means tri-gram",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+@torch.no_grad()
+def main():
+ args = get_parser().parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ params = get_params()
+ params.update(vars(args))
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"device: {device}")
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(params.bpe_model)
+
+ # is defined in local/train_bpe_model.py
+ params.blank_id = sp.piece_to_id("")
+ params.vocab_size = sp.get_piece_size()
+
+ logging.info(params)
+
+ logging.info("About to create model")
+ model = get_transducer_model(params)
+
+ model.to(device)
+
+ if not params.use_averaged_model:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ elif 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 i >= 1:
+ filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg + 1
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg + 1:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ filename_start = filenames[-1]
+ filename_end = filenames[0]
+ logging.info(
+ "Calculating the averaged model over iteration checkpoints"
+ f" from {filename_start} (excluded) to {filename_end}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+ else:
+ assert params.avg > 0, params.avg
+ start = params.epoch - params.avg
+ assert start >= 1, start
+ filename_start = f"{params.exp_dir}/epoch-{start}.pt"
+ filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
+ logging.info(
+ f"Calculating the averaged model over epoch range from "
+ f"{start} (excluded) to {params.epoch}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+
+ model.to("cpu")
+ model.eval()
+
+ if params.jit is True:
+ convert_scaled_to_non_scaled(model, inplace=True)
+ logging.info("Using torch.jit.script()")
+ # We won't use the forward() method of the model in C++, so just ignore
+ # it here.
+ # Otherwise, one of its arguments is a ragged tensor and is not
+ # torch scriptabe.
+ model.__class__.forward = torch.jit.ignore(model.__class__.forward)
+ logging.info("Using torch.jit.script")
+ model = torch.jit.script(model)
+ filename = params.exp_dir / "cpu_jit.pt"
+ model.save(str(filename))
+ logging.info(f"Saved to {filename}")
+ else:
+ logging.info("Not using torchscript. Export model.state_dict()")
+ # Save it using a format so that it can be loaded
+ # by :func:`load_checkpoint`
+ filename = params.exp_dir / "pretrained.pt"
+ torch.save({"model": model.state_dict()}, str(filename))
+ logging.info(f"Saved to {filename}")
+
+
+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/pruned_transducer_stateless7_ctc_bs/frame_reducer.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/frame_reducer.py
new file mode 100755
index 000000000..9fe88929d
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/frame_reducer.py
@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022 Xiaomi Corp. (authors: Yifan Yang,
+# Zengwei Yao)
+#
+# 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.
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from torch.nn.utils.rnn import pad_sequence
+from icefall.utils import make_pad_mask
+
+
+class FrameReducer(nn.Module):
+ """The encoder output is first used to calculate
+ the CTC posterior probability; then for each output frame,
+ if its blank posterior is bigger than some thresholds,
+ it will be simply discarded from the encoder output.
+ """
+
+ def __init__(
+ self,
+ ):
+ super().__init__()
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ x_lens: torch.Tensor,
+ ctc_output: torch.Tensor,
+ blank_id: int = 0,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Args:
+ x:
+ The shared encoder output with shape [N, T, C].
+ x_lens:
+ A tensor of shape (batch_size,) containing the number of frames in
+ `x` before padding.
+ ctc_output:
+ The CTC output with shape [N, T, vocab_size].
+ blank_id:
+ The ID of the blank symbol.
+ Returns:
+ x_fr:
+ The frame reduced encoder output with shape [N, T', C].
+ x_lens_fr:
+ A tensor of shape (batch_size,) containing the number of frames in
+ `x_fr` before padding.
+ """
+
+ padding_mask = make_pad_mask(x_lens)
+ non_blank_mask = (ctc_output[:, :, blank_id] < math.log(0.9)) * (~padding_mask)
+
+ frames_list: List[torch.Tensor] = []
+ lens_list: List[int] = []
+ for i in range(x.shape[0]):
+ frames = x[i][non_blank_mask[i]]
+ frames_list.append(frames)
+ lens_list.append(frames.shape[0])
+ x_fr = pad_sequence(frames_list, batch_first=True)
+ x_lens_fr = torch.tensor(lens_list).to(device=x.device)
+
+ return x_fr, x_lens_fr
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/jit_pretrained.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/jit_pretrained.py
new file mode 100755
index 000000000..da2c6a39a
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/jit_pretrained.py
@@ -0,0 +1,271 @@
+#!/usr/bin/env python3
+# Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang)
+#
+# 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.
+"""
+This script loads torchscript models, exported by `torch.jit.script()`
+and uses them to decode waves.
+You can use the following command to get the exported models:
+
+./pruned_transducer_stateless7_ctc_bs/export.py \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 20 \
+ --avg 10 \
+ --jit 1
+
+Usage of this script:
+
+./pruned_transducer_stateless7_ctc_bs/jit_pretrained.py \
+ --nn-model-filename ./pruned_transducer_stateless7_ctc_bs/exp/cpu_jit.pt \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+"""
+
+import argparse
+import logging
+import math
+from typing import List
+
+import kaldifeat
+import sentencepiece as spm
+import torch
+import torchaudio
+from torch.nn.utils.rnn import pad_sequence
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--nn-model-filename",
+ type=str,
+ required=True,
+ help="Path to the torchscript model cpu_jit.pt",
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ help="""Path to bpe.model.""",
+ )
+
+ parser.add_argument(
+ "sound_files",
+ type=str,
+ nargs="+",
+ help="The input sound file(s) to transcribe. "
+ "Supported formats are those supported by torchaudio.load(). "
+ "For example, wav and flac are supported. "
+ "The sample rate has to be 16kHz.",
+ )
+
+ return parser
+
+
+def read_sound_files(
+ filenames: List[str], expected_sample_rate: float = 16000
+) -> List[torch.Tensor]:
+ """Read a list of sound files into a list 1-D float32 torch tensors.
+ Args:
+ filenames:
+ A list of sound filenames.
+ expected_sample_rate:
+ The expected sample rate of the sound files.
+ Returns:
+ Return a list of 1-D float32 torch tensors.
+ """
+ ans = []
+ for f in filenames:
+ wave, sample_rate = torchaudio.load(f)
+ assert (
+ sample_rate == expected_sample_rate
+ ), f"Expected sample rate: {expected_sample_rate}. Given: {sample_rate}"
+ # We use only the first channel
+ ans.append(wave[0])
+ return ans
+
+
+def greedy_search(
+ model: torch.jit.ScriptModule,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+) -> List[List[int]]:
+ """Greedy search in batch mode. It hardcodes --max-sym-per-frame=1.
+ Args:
+ model:
+ The transducer model.
+ encoder_out:
+ A 3-D tensor of shape (N, T, C)
+ encoder_out_lens:
+ A 1-D tensor of shape (N,).
+ Returns:
+ Return the decoded results for each utterance.
+ """
+ assert encoder_out.ndim == 3
+ assert encoder_out.size(0) >= 1, encoder_out.size(0)
+
+ packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
+ input=encoder_out,
+ lengths=encoder_out_lens.cpu(),
+ batch_first=True,
+ enforce_sorted=False,
+ )
+
+ device = encoder_out.device
+ blank_id = 0 # hard-code to 0
+
+ batch_size_list = packed_encoder_out.batch_sizes.tolist()
+ N = encoder_out.size(0)
+
+ assert torch.all(encoder_out_lens > 0), encoder_out_lens
+ assert N == batch_size_list[0], (N, batch_size_list)
+
+ context_size = model.decoder.context_size
+ hyps = [[blank_id] * context_size for _ in range(N)]
+
+ decoder_input = torch.tensor(
+ hyps,
+ device=device,
+ dtype=torch.int64,
+ ) # (N, context_size)
+
+ decoder_out = model.decoder(
+ decoder_input,
+ need_pad=torch.tensor([False]),
+ ).squeeze(1)
+
+ offset = 0
+ for batch_size in batch_size_list:
+ start = offset
+ end = offset + batch_size
+ current_encoder_out = packed_encoder_out.data[start:end]
+ current_encoder_out = current_encoder_out
+ # current_encoder_out's shape: (batch_size, encoder_out_dim)
+ offset = end
+
+ decoder_out = decoder_out[:batch_size]
+
+ logits = model.joiner(
+ current_encoder_out,
+ decoder_out,
+ )
+ # logits'shape (batch_size, vocab_size)
+
+ assert logits.ndim == 2, logits.shape
+ y = logits.argmax(dim=1).tolist()
+ emitted = False
+ for i, v in enumerate(y):
+ if v != blank_id:
+ hyps[i].append(v)
+ emitted = True
+ if emitted:
+ # update decoder output
+ decoder_input = [h[-context_size:] for h in hyps[:batch_size]]
+ decoder_input = torch.tensor(
+ decoder_input,
+ device=device,
+ dtype=torch.int64,
+ )
+ decoder_out = model.decoder(
+ decoder_input,
+ need_pad=torch.tensor([False]),
+ )
+ decoder_out = decoder_out.squeeze(1)
+
+ sorted_ans = [h[context_size:] for h in hyps]
+ ans = []
+ unsorted_indices = packed_encoder_out.unsorted_indices.tolist()
+ for i in range(N):
+ ans.append(sorted_ans[unsorted_indices[i]])
+
+ return ans
+
+
+@torch.no_grad()
+def main():
+ parser = get_parser()
+ args = parser.parse_args()
+ logging.info(vars(args))
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"device: {device}")
+
+ model = torch.jit.load(args.nn_model_filename)
+
+ model.eval()
+
+ model.to(device)
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(args.bpe_model)
+
+ logging.info("Constructing Fbank computer")
+ opts = kaldifeat.FbankOptions()
+ opts.device = device
+ opts.frame_opts.dither = 0
+ opts.frame_opts.snip_edges = False
+ opts.frame_opts.samp_freq = 16000
+ opts.mel_opts.num_bins = 80
+
+ fbank = kaldifeat.Fbank(opts)
+
+ logging.info(f"Reading sound files: {args.sound_files}")
+ waves = read_sound_files(
+ filenames=args.sound_files,
+ )
+ waves = [w.to(device) for w in waves]
+
+ logging.info("Decoding started")
+ features = fbank(waves)
+ feature_lengths = [f.size(0) for f in features]
+
+ features = pad_sequence(
+ features,
+ batch_first=True,
+ padding_value=math.log(1e-10),
+ )
+
+ feature_lengths = torch.tensor(feature_lengths, device=device)
+
+ encoder_out, encoder_out_lens = model.encoder(
+ x=features,
+ x_lens=feature_lengths,
+ )
+
+ hyps = greedy_search(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ )
+ s = "\n"
+ for filename, hyp in zip(args.sound_files, hyps):
+ words = sp.decode(hyp)
+ s += f"{filename}:\n{words}\n\n"
+ logging.info(s)
+
+ logging.info("Decoding Done")
+
+
+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/pruned_transducer_stateless7_ctc_bs/jit_pretrained_ctc.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/jit_pretrained_ctc.py
new file mode 100755
index 000000000..653c25e06
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/jit_pretrained_ctc.py
@@ -0,0 +1,426 @@
+#!/usr/bin/env python3
+# Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang,
+# Zengwei Yao)
+#
+# 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.
+"""
+This script loads torchscript models, exported by `torch.jit.script()`
+and uses them to decode waves.
+You can use the following command to get the exported models:
+
+./pruned_transducer_stateless7_ctc_bs/export.py \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 13 \
+ --jit 1
+
+Usage of this script:
+
+(1) ctc-decoding
+./pruned_transducer_stateless7_ctc_bs/jit_pretrained_ctc.py \
+ --model-filename ./pruned_transducer_stateless7_ctc_bs/exp/cpu_jit.pt \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --method ctc-decoding \
+ --sample-rate 16000 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+(2) 1best
+./pruned_transducer_stateless7_ctc_bs/jit_pretrained_ctc.py \
+ --model-filename ./pruned_transducer_stateless7_ctc_bs/exp/cpu_jit.pt \
+ --HLG data/lang_bpe_500/HLG.pt \
+ --words-file data/lang_bpe_500/words.txt \
+ --method 1best \
+ --sample-rate 16000 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+
+(3) nbest-rescoring
+./pruned_transducer_stateless7_ctc_bs/jit_pretrained_ctc.py \
+ --model-filename ./pruned_transducer_stateless7_ctc_bs/exp/cpu_jit.pt \
+ --HLG data/lang_bpe_500/HLG.pt \
+ --words-file data/lang_bpe_500/words.txt \
+ --G data/lm/G_4_gram.pt \
+ --method nbest-rescoring \
+ --sample-rate 16000 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+
+(4) whole-lattice-rescoring
+./pruned_transducer_stateless7_ctc_bs/jit_pretrained_ctc.py \
+ --model-filename ./pruned_transducer_stateless7_ctc_bs/exp/cpu_jit.pt \
+ --HLG data/lang_bpe_500/HLG.pt \
+ --words-file data/lang_bpe_500/words.txt \
+ --G data/lm/G_4_gram.pt \
+ --method whole-lattice-rescoring \
+ --sample-rate 16000 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+"""
+
+import argparse
+import logging
+import math
+from typing import List
+
+import k2
+import kaldifeat
+import sentencepiece as spm
+import torch
+import torchaudio
+from ctc_decode import get_decoding_params
+from torch.nn.utils.rnn import pad_sequence
+from train import get_params
+
+from icefall.decode import (
+ get_lattice,
+ one_best_decoding,
+ rescore_with_n_best_list,
+ rescore_with_whole_lattice,
+)
+from icefall.utils import get_texts
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--model-filename",
+ type=str,
+ required=True,
+ help="Path to the torchscript model.",
+ )
+
+ parser.add_argument(
+ "--words-file",
+ type=str,
+ help="""Path to words.txt.
+ Used only when method is not ctc-decoding.
+ """,
+ )
+
+ parser.add_argument(
+ "--HLG",
+ type=str,
+ help="""Path to HLG.pt.
+ Used only when method is not ctc-decoding.
+ """,
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ help="""Path to bpe.model.
+ Used only when method is ctc-decoding.
+ """,
+ )
+
+ parser.add_argument(
+ "--method",
+ type=str,
+ default="1best",
+ help="""Decoding method.
+ Possible values are:
+ (0) ctc-decoding - Use CTC decoding. It uses a sentence
+ piece model, i.e., lang_dir/bpe.model, to convert
+ word pieces to words. It needs neither a lexicon
+ nor an n-gram LM.
+ (1) 1best - Use the best path as decoding output. Only
+ the transformer encoder output is used for decoding.
+ We call it HLG decoding.
+ (2) nbest-rescoring. Extract n paths from the decoding lattice,
+ rescore them with an LM, the path with
+ the highest score is the decoding result.
+ We call it HLG decoding + n-gram LM rescoring.
+ (3) whole-lattice-rescoring - Use an LM to rescore the
+ decoding lattice and then use 1best to decode the
+ rescored lattice.
+ We call it HLG decoding + n-gram LM rescoring.
+ """,
+ )
+
+ parser.add_argument(
+ "--G",
+ type=str,
+ help="""An LM for rescoring.
+ Used only when method is
+ whole-lattice-rescoring or nbest-rescoring.
+ It's usually a 4-gram LM.
+ """,
+ )
+
+ parser.add_argument(
+ "--num-paths",
+ type=int,
+ default=100,
+ help="""
+ Used only when method is attention-decoder.
+ It specifies the size of n-best list.""",
+ )
+
+ parser.add_argument(
+ "--ngram-lm-scale",
+ type=float,
+ default=1.3,
+ help="""
+ Used only when method is whole-lattice-rescoring and nbest-rescoring.
+ It specifies the scale for n-gram LM scores.
+ (Note: You need to tune it on a dataset.)
+ """,
+ )
+
+ parser.add_argument(
+ "--nbest-scale",
+ type=float,
+ default=0.5,
+ help="""
+ Used only when method is nbest-rescoring.
+ It specifies the scale for lattice.scores when
+ extracting n-best lists. A smaller value results in
+ more unique number of paths with the risk of missing
+ the best path.
+ """,
+ )
+
+ parser.add_argument(
+ "--num-classes",
+ type=int,
+ default=500,
+ help="""
+ Vocab size in the BPE model.
+ """,
+ )
+
+ parser.add_argument(
+ "--sample-rate",
+ type=int,
+ default=16000,
+ help="The sample rate of the input sound file",
+ )
+
+ parser.add_argument(
+ "sound_files",
+ type=str,
+ nargs="+",
+ help="The input sound file(s) to transcribe. "
+ "Supported formats are those supported by torchaudio.load(). "
+ "For example, wav and flac are supported. "
+ "The sample rate has to be 16kHz.",
+ )
+
+ return parser
+
+
+def read_sound_files(
+ filenames: List[str], expected_sample_rate: float = 16000
+) -> List[torch.Tensor]:
+ """Read a list of sound files into a list 1-D float32 torch tensors.
+ Args:
+ filenames:
+ A list of sound filenames.
+ expected_sample_rate:
+ The expected sample rate of the sound files.
+ Returns:
+ Return a list of 1-D float32 torch tensors.
+ """
+ ans = []
+ for f in filenames:
+ wave, sample_rate = torchaudio.load(f)
+ assert (
+ sample_rate == expected_sample_rate
+ ), f"Expected sample rate: {expected_sample_rate}. Given: {sample_rate}"
+ # We use only the first channel
+ ans.append(wave[0])
+ return ans
+
+
+@torch.no_grad()
+def main():
+ parser = get_parser()
+ args = parser.parse_args()
+
+ params = get_params()
+ # add decoding params
+ params.update(get_decoding_params())
+ params.update(vars(args))
+
+ logging.info(f"{params}")
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"device: {device}")
+
+ model = torch.jit.load(args.model_filename)
+ model.to(device)
+ model.eval()
+
+ logging.info("Constructing Fbank computer")
+ opts = kaldifeat.FbankOptions()
+ opts.device = device
+ opts.frame_opts.dither = 0
+ opts.frame_opts.snip_edges = False
+ opts.frame_opts.samp_freq = params.sample_rate
+ opts.mel_opts.num_bins = params.feature_dim
+
+ fbank = kaldifeat.Fbank(opts)
+
+ logging.info(f"Reading sound files: {params.sound_files}")
+ waves = read_sound_files(
+ filenames=params.sound_files, expected_sample_rate=params.sample_rate
+ )
+ waves = [w.to(device) for w in waves]
+
+ logging.info("Decoding started")
+ features = fbank(waves)
+ feature_lengths = [f.size(0) for f in features]
+
+ features = pad_sequence(features, batch_first=True, padding_value=math.log(1e-10))
+ feature_lengths = torch.tensor(feature_lengths, device=device)
+
+ encoder_out, encoder_out_lens = model.encoder(
+ x=features,
+ x_lens=feature_lengths,
+ )
+ nnet_output = model.ctc_output(encoder_out)
+
+ batch_size = nnet_output.shape[0]
+ supervision_segments = torch.tensor(
+ [
+ [i, 0, feature_lengths[i] // params.subsampling_factor]
+ for i in range(batch_size)
+ ],
+ dtype=torch.int32,
+ )
+
+ if params.method == "ctc-decoding":
+ logging.info("Use CTC decoding")
+ bpe_model = spm.SentencePieceProcessor()
+ bpe_model.load(params.bpe_model)
+ max_token_id = params.num_classes - 1
+
+ H = k2.ctc_topo(
+ max_token=max_token_id,
+ modified=False,
+ device=device,
+ )
+
+ lattice = get_lattice(
+ nnet_output=nnet_output,
+ decoding_graph=H,
+ 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,
+ )
+
+ best_path = one_best_decoding(
+ lattice=lattice, use_double_scores=params.use_double_scores
+ )
+ token_ids = get_texts(best_path)
+ hyps = bpe_model.decode(token_ids)
+ hyps = [s.split() for s in hyps]
+ elif params.method in [
+ "1best",
+ "nbest-rescoring",
+ "whole-lattice-rescoring",
+ ]:
+ logging.info(f"Loading HLG from {params.HLG}")
+ HLG = k2.Fsa.from_dict(torch.load(params.HLG, map_location="cpu"))
+ HLG = HLG.to(device)
+ if not hasattr(HLG, "lm_scores"):
+ # For whole-lattice-rescoring and attention-decoder
+ HLG.lm_scores = HLG.scores.clone()
+
+ if params.method in [
+ "nbest-rescoring",
+ "whole-lattice-rescoring",
+ ]:
+ logging.info(f"Loading G from {params.G}")
+ G = k2.Fsa.from_dict(torch.load(params.G, map_location="cpu"))
+ G = G.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.lm_scores is used to replace HLG.lm_scores during
+ # LM rescoring.
+ G.lm_scores = G.scores.clone()
+
+ lattice = get_lattice(
+ nnet_output=nnet_output,
+ decoding_graph=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 == "1best":
+ logging.info("Use HLG decoding")
+ best_path = one_best_decoding(
+ lattice=lattice, use_double_scores=params.use_double_scores
+ )
+ if params.method == "nbest-rescoring":
+ logging.info("Use HLG decoding + LM rescoring")
+ best_path_dict = rescore_with_n_best_list(
+ lattice=lattice,
+ G=G,
+ num_paths=params.num_paths,
+ lm_scale_list=[params.ngram_lm_scale],
+ nbest_scale=params.nbest_scale,
+ )
+ best_path = next(iter(best_path_dict.values()))
+ elif params.method == "whole-lattice-rescoring":
+ logging.info("Use HLG decoding + LM rescoring")
+ best_path_dict = rescore_with_whole_lattice(
+ lattice=lattice,
+ G_with_epsilon_loops=G,
+ lm_scale_list=[params.ngram_lm_scale],
+ )
+ best_path = next(iter(best_path_dict.values()))
+
+ hyps = get_texts(best_path)
+ word_sym_table = k2.SymbolTable.from_file(params.words_file)
+ hyps = [[word_sym_table[i] for i in ids] for ids in hyps]
+ else:
+ raise ValueError(f"Unsupported decoding method: {params.method}")
+
+ s = "\n"
+ for filename, hyp in zip(params.sound_files, hyps):
+ words = " ".join(hyp)
+ s += f"{filename}:\n{words}\n\n"
+ logging.info(s)
+
+ logging.info("Decoding Done")
+
+
+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/pruned_transducer_stateless7_ctc_bs/joiner.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/joiner.py
new file mode 120000
index 000000000..ecfb6dd8a
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/joiner.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/joiner.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/lconv.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/lconv.py
new file mode 100755
index 000000000..bfd49d533
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/lconv.py
@@ -0,0 +1,114 @@
+# Copyright 2022 Xiaomi Corp. (authors: Yifan Yang)
+#
+# 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.
+
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from scaling import (
+ ActivationBalancer,
+ ScaledConv1d,
+)
+
+
+class LConv(nn.Module):
+ """A convolution module to prevent information loss."""
+
+ def __init__(
+ self,
+ channels: int,
+ kernel_size: int = 7,
+ bias: bool = True,
+ ):
+ """
+ Args:
+ channels:
+ Dimension of the input embedding, and of the lconv output.
+ """
+ super().__init__()
+ self.pointwise_conv1 = nn.Conv1d(
+ channels,
+ 2 * channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ bias=bias,
+ )
+
+ self.deriv_balancer1 = ActivationBalancer(
+ 2 * channels,
+ channel_dim=1,
+ max_abs=10.0,
+ min_positive=0.05,
+ max_positive=1.0,
+ )
+
+ self.depthwise_conv = nn.Conv1d(
+ 2 * channels,
+ 2 * channels,
+ kernel_size=kernel_size,
+ stride=1,
+ padding=(kernel_size - 1) // 2,
+ groups=channels,
+ bias=bias,
+ )
+
+ self.deriv_balancer2 = ActivationBalancer(
+ 2 * channels,
+ channel_dim=1,
+ min_positive=0.05,
+ max_positive=1.0,
+ max_abs=20.0,
+ )
+
+ self.pointwise_conv2 = ScaledConv1d(
+ 2 * channels,
+ channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ bias=bias,
+ initial_scale=0.05,
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ src_key_padding_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """
+ Args:
+ x: A 3-D tensor of shape (N, T, C).
+ Returns:
+ Return a tensor of shape (N, T, C).
+ """
+ # exchange the temporal dimension and the feature dimension
+ x = x.permute(0, 2, 1) # (#batch, channels, time).
+
+ x = self.pointwise_conv1(x) # (batch, 2*channels, time)
+
+ x = self.deriv_balancer1(x)
+
+ if src_key_padding_mask is not None:
+ x = x.masked_fill(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
+
+ x = self.depthwise_conv(x)
+
+ x = self.deriv_balancer2(x)
+
+ x = self.pointwise_conv2(x) # (batch, channels, time)
+
+ return x.permute(0, 2, 1)
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/model.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/model.py
new file mode 100755
index 000000000..86acc5a10
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/model.py
@@ -0,0 +1,224 @@
+# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang, Wei Kang)
+#
+# 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.
+
+
+from typing import Tuple
+
+import k2
+import torch
+import torch.nn as nn
+from encoder_interface import EncoderInterface
+
+from icefall.utils import add_sos, make_pad_mask
+
+
+class Transducer(nn.Module):
+ """It implements https://arxiv.org/pdf/1211.3711.pdf
+ "Sequence Transduction with Recurrent Neural Networks"
+ """
+
+ def __init__(
+ self,
+ encoder: EncoderInterface,
+ decoder: nn.Module,
+ joiner: nn.Module,
+ lconv: nn.Module,
+ frame_reducer: nn.Module,
+ encoder_dim: int,
+ decoder_dim: int,
+ joiner_dim: int,
+ vocab_size: int,
+ ):
+ """
+ Args:
+ encoder:
+ It is the transcription network in the paper. Its accepts
+ two inputs: `x` of (N, T, encoder_dim) and `x_lens` of shape (N,).
+ It returns two tensors: `logits` of shape (N, T, encoder_dm) and
+ `logit_lens` of shape (N,).
+ decoder:
+ It is the prediction network in the paper. Its input shape
+ is (N, U) and its output shape is (N, U, decoder_dim).
+ It should contain one attribute: `blank_id`.
+ joiner:
+ It has two inputs with shapes: (N, T, encoder_dim) and (N, U, decoder_dim).
+ Its output shape is (N, T, U, vocab_size). Note that its output contains
+ unnormalized probs, i.e., not processed by log-softmax.
+ """
+ super().__init__()
+ assert isinstance(encoder, EncoderInterface), type(encoder)
+ assert hasattr(decoder, "blank_id")
+
+ self.encoder = encoder
+ self.decoder = decoder
+ self.joiner = joiner
+ self.lconv = lconv
+ self.frame_reducer = frame_reducer
+
+ self.simple_am_proj = nn.Linear(
+ encoder_dim,
+ vocab_size,
+ )
+ self.simple_lm_proj = nn.Linear(decoder_dim, vocab_size)
+
+ self.ctc_output = nn.Sequential(
+ nn.Dropout(p=0.1),
+ nn.Linear(encoder_dim, vocab_size),
+ nn.LogSoftmax(dim=-1),
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ x_lens: torch.Tensor,
+ y: k2.RaggedTensor,
+ prune_range: int = 5,
+ am_scale: float = 0.0,
+ lm_scale: float = 0.0,
+ warmup: float = 1.0,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Args:
+ x:
+ A 3-D tensor of shape (N, T, C).
+ x_lens:
+ A 1-D tensor of shape (N,). It contains the number of frames in `x`
+ before padding.
+ y:
+ A ragged tensor with 2 axes [utt][label]. It contains labels of each
+ utterance.
+ prune_range:
+ The prune range for rnnt loss, it means how many symbols(context)
+ we are considering for each frame to compute the loss.
+ am_scale:
+ The scale to smooth the loss with am (output of encoder network)
+ part
+ lm_scale:
+ The scale to smooth the loss with lm (output of predictor network)
+ part
+ warmup:
+ A floating point value which decides whether to do blank skip.
+ Returns:
+ Return a tuple containing simple loss, pruned loss, and ctc-output.
+ Note:
+ Regarding am_scale & lm_scale, it will make the loss-function one of
+ the form:
+ lm_scale * lm_probs + am_scale * am_probs +
+ (1-lm_scale-am_scale) * combined_probs
+ """
+ assert x.ndim == 3, x.shape
+ assert x_lens.ndim == 1, x_lens.shape
+ assert y.num_axes == 2, y.num_axes
+
+ assert x.size(0) == x_lens.size(0) == y.dim0
+
+ encoder_out, x_lens = self.encoder(x, x_lens)
+ assert torch.all(x_lens > 0)
+
+ # compute ctc log-probs
+ ctc_output = self.ctc_output(encoder_out)
+
+ # blank skip
+ blank_id = self.decoder.blank_id
+
+ if warmup >= 2.0:
+ # lconv
+ encoder_out = self.lconv(
+ x=encoder_out,
+ src_key_padding_mask=make_pad_mask(x_lens),
+ )
+
+ # frame reduce
+ encoder_out_fr, x_lens_fr = self.frame_reducer(
+ encoder_out,
+ x_lens,
+ ctc_output,
+ blank_id,
+ )
+ else:
+ encoder_out_fr = encoder_out
+ x_lens_fr = x_lens
+
+ # Now for the decoder, i.e., the prediction network
+ row_splits = y.shape.row_splits(1)
+ y_lens = row_splits[1:] - row_splits[:-1]
+
+ sos_y = add_sos(y, sos_id=blank_id)
+
+ # sos_y_padded: [B, S + 1], start with SOS.
+ sos_y_padded = sos_y.pad(mode="constant", padding_value=blank_id)
+
+ # decoder_out: [B, S + 1, decoder_dim]
+ decoder_out = self.decoder(sos_y_padded)
+
+ # Note: y does not start with SOS
+ # y_padded : [B, S]
+ y_padded = y.pad(mode="constant", padding_value=0)
+
+ y_padded = y_padded.to(torch.int64)
+ boundary = torch.zeros((x.size(0), 4), dtype=torch.int64, device=x.device)
+ boundary[:, 2] = y_lens
+ boundary[:, 3] = x_lens_fr
+
+ am = self.simple_am_proj(encoder_out_fr)
+ lm = self.simple_lm_proj(decoder_out)
+
+ with torch.cuda.amp.autocast(enabled=False):
+ simple_loss, (px_grad, py_grad) = k2.rnnt_loss_smoothed(
+ lm=lm.float(),
+ am=am.float(),
+ symbols=y_padded,
+ termination_symbol=blank_id,
+ lm_only_scale=lm_scale,
+ am_only_scale=am_scale,
+ boundary=boundary,
+ reduction="sum",
+ return_grad=True,
+ )
+
+ # ranges : [B, T, prune_range]
+ ranges = k2.get_rnnt_prune_ranges(
+ px_grad=px_grad,
+ py_grad=py_grad,
+ boundary=boundary,
+ s_range=prune_range,
+ )
+
+ # am_pruned : [B, T, prune_range, encoder_dim]
+ # lm_pruned : [B, T, prune_range, decoder_dim]
+ am_pruned, lm_pruned = k2.do_rnnt_pruning(
+ am=self.joiner.encoder_proj(encoder_out_fr),
+ lm=self.joiner.decoder_proj(decoder_out),
+ ranges=ranges,
+ )
+
+ # logits : [B, T, prune_range, vocab_size]
+
+ # project_input=False since we applied the decoder's input projections
+ # prior to do_rnnt_pruning (this is an optimization for speed).
+ logits = self.joiner(am_pruned, lm_pruned, project_input=False)
+
+ with torch.cuda.amp.autocast(enabled=False):
+ pruned_loss = k2.rnnt_loss_pruned(
+ logits=logits.float(),
+ symbols=y_padded,
+ ranges=ranges,
+ termination_symbol=blank_id,
+ boundary=boundary,
+ reduction="sum",
+ )
+
+ return (simple_loss, pruned_loss, ctc_output)
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/optim.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/optim.py
new file mode 120000
index 000000000..81ac4a89a
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/optim.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/optim.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/pretrained.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/pretrained.py
new file mode 100755
index 000000000..ea0fe9164
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/pretrained.py
@@ -0,0 +1,352 @@
+#!/usr/bin/env python3
+# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang)
+#
+# 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.
+"""
+This script loads a checkpoint and uses it to decode waves.
+You can generate the checkpoint with the following command:
+
+./pruned_transducer_stateless7_ctc_bs/export.py \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 13
+
+Usage of this script:
+
+(1) greedy search
+./pruned_transducer_stateless7_ctc_bs/pretrained.py \
+ --checkpoint ./pruned_transducer_stateless7_ctc_bs/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method greedy_search \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+(2) beam search
+./pruned_transducer_stateless7_ctc_bs/pretrained.py \
+ --checkpoint ./pruned_transducer_stateless7_ctc_bs/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method beam_search \
+ --beam-size 4 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+(3) modified beam search
+./pruned_transducer_stateless7_ctc_bs/pretrained.py \
+ --checkpoint ./pruned_transducer_stateless7_ctc_bs/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method modified_beam_search \
+ --beam-size 4 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+(4) fast beam search
+./pruned_transducer_stateless7_ctc_bs/pretrained.py \
+ --checkpoint ./pruned_transducer_stateless7_ctc_bs/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method fast_beam_search \
+ --beam-size 4 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+You can also use `./pruned_transducer_stateless7_ctc_bs/exp/epoch-xx.pt`.
+
+Note: ./pruned_transducer_stateless7_ctc_bs/exp/pretrained.pt is generated by
+./pruned_transducer_stateless7_ctc_bs/export.py
+"""
+
+
+import argparse
+import logging
+import math
+from typing import List
+
+import k2
+import kaldifeat
+import sentencepiece as spm
+import torch
+import torchaudio
+from beam_search import (
+ beam_search,
+ fast_beam_search_one_best,
+ greedy_search,
+ greedy_search_batch,
+ modified_beam_search,
+)
+from torch.nn.utils.rnn import pad_sequence
+from train import add_model_arguments, get_params, get_transducer_model
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--checkpoint",
+ type=str,
+ required=True,
+ help="Path to the checkpoint. "
+ "The checkpoint is assumed to be saved by "
+ "icefall.checkpoint.save_checkpoint().",
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ help="""Path to bpe.model.""",
+ )
+
+ parser.add_argument(
+ "--method",
+ type=str,
+ default="greedy_search",
+ help="""Possible values are:
+ - greedy_search
+ - beam_search
+ - modified_beam_search
+ - fast_beam_search
+ """,
+ )
+
+ parser.add_argument(
+ "sound_files",
+ type=str,
+ nargs="+",
+ help="The input sound file(s) to transcribe. "
+ "Supported formats are those supported by torchaudio.load(). "
+ "For example, wav and flac are supported. "
+ "The sample rate has to be 16kHz.",
+ )
+
+ parser.add_argument(
+ "--sample-rate",
+ type=int,
+ default=16000,
+ help="The sample rate of the input sound file",
+ )
+
+ parser.add_argument(
+ "--beam-size",
+ type=int,
+ default=4,
+ help="""An integer indicating how many candidates we will keep for each
+ frame. Used only when --method is beam_search or
+ modified_beam_search.""",
+ )
+
+ parser.add_argument(
+ "--beam",
+ type=float,
+ default=4,
+ help="""A floating point value to calculate the cutoff score during beam
+ search (i.e., `cutoff = max-score - beam`), which is the same as the
+ `beam` in Kaldi.
+ Used only when --method is fast_beam_search""",
+ )
+
+ parser.add_argument(
+ "--max-contexts",
+ type=int,
+ default=4,
+ help="""Used only when --method is fast_beam_search""",
+ )
+
+ parser.add_argument(
+ "--max-states",
+ type=int,
+ default=8,
+ help="""Used only when --method is fast_beam_search""",
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram, 2 means tri-gram",
+ )
+ parser.add_argument(
+ "--max-sym-per-frame",
+ type=int,
+ default=1,
+ help="""Maximum number of symbols per frame. Used only when
+ --method is greedy_search.
+ """,
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def read_sound_files(
+ filenames: List[str], expected_sample_rate: float
+) -> List[torch.Tensor]:
+ """Read a list of sound files into a list 1-D float32 torch tensors.
+ Args:
+ filenames:
+ A list of sound filenames.
+ expected_sample_rate:
+ The expected sample rate of the sound files.
+ Returns:
+ Return a list of 1-D float32 torch tensors.
+ """
+ ans = []
+ for f in filenames:
+ wave, sample_rate = torchaudio.load(f)
+ assert sample_rate == expected_sample_rate, (
+ f"expected sample rate: {expected_sample_rate}. " f"Given: {sample_rate}"
+ )
+ # We use only the first channel
+ ans.append(wave[0])
+ return ans
+
+
+@torch.no_grad()
+def main():
+ parser = get_parser()
+ args = parser.parse_args()
+
+ params = get_params()
+
+ params.update(vars(args))
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(params.bpe_model)
+
+ # is defined in local/train_bpe_model.py
+ params.blank_id = sp.piece_to_id("")
+ params.unk_id = sp.piece_to_id("")
+ params.vocab_size = sp.get_piece_size()
+
+ logging.info(f"{params}")
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"device: {device}")
+
+ logging.info("Creating model")
+ model = get_transducer_model(params)
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ checkpoint = torch.load(args.checkpoint, map_location="cpu")
+ model.load_state_dict(checkpoint["model"], strict=False)
+ model.to(device)
+ model.eval()
+ model.device = device
+
+ logging.info("Constructing Fbank computer")
+ opts = kaldifeat.FbankOptions()
+ opts.device = device
+ opts.frame_opts.dither = 0
+ opts.frame_opts.snip_edges = False
+ opts.frame_opts.samp_freq = params.sample_rate
+ opts.mel_opts.num_bins = params.feature_dim
+
+ fbank = kaldifeat.Fbank(opts)
+
+ logging.info(f"Reading sound files: {params.sound_files}")
+ waves = read_sound_files(
+ filenames=params.sound_files, expected_sample_rate=params.sample_rate
+ )
+ waves = [w.to(device) for w in waves]
+
+ logging.info("Decoding started")
+ features = fbank(waves)
+ feature_lengths = [f.size(0) for f in features]
+
+ features = pad_sequence(features, batch_first=True, padding_value=math.log(1e-10))
+
+ feature_lengths = torch.tensor(feature_lengths, device=device)
+
+ encoder_out, encoder_out_lens = model.encoder(x=features, x_lens=feature_lengths)
+
+ num_waves = encoder_out.size(0)
+ hyps = []
+ msg = f"Using {params.method}"
+ if params.method == "beam_search":
+ msg += f" with beam size {params.beam_size}"
+ logging.info(msg)
+
+ if params.method == "fast_beam_search":
+ decoding_graph = k2.trivial_graph(params.vocab_size - 1, device=device)
+ hyp_tokens = fast_beam_search_one_best(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.method == "modified_beam_search":
+ hyp_tokens = modified_beam_search(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam_size,
+ )
+
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.method == "greedy_search" and params.max_sym_per_frame == 1:
+ hyp_tokens = greedy_search_batch(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ else:
+ for i in range(num_waves):
+ # fmt: off
+ encoder_out_i = encoder_out[i:i+1, :encoder_out_lens[i]]
+ # fmt: on
+ if params.method == "greedy_search":
+ hyp = greedy_search(
+ model=model,
+ encoder_out=encoder_out_i,
+ max_sym_per_frame=params.max_sym_per_frame,
+ )
+ elif params.method == "beam_search":
+ hyp = beam_search(
+ model=model,
+ encoder_out=encoder_out_i,
+ beam=params.beam_size,
+ )
+ else:
+ raise ValueError(f"Unsupported method: {params.method}")
+
+ hyps.append(sp.decode(hyp).split())
+
+ s = "\n"
+ for filename, hyp in zip(params.sound_files, hyps):
+ words = " ".join(hyp)
+ s += f"{filename}:\n{words}\n\n"
+ logging.info(s)
+
+ logging.info("Decoding Done")
+
+
+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/pruned_transducer_stateless7_ctc_bs/pretrained_ctc.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/pretrained_ctc.py
new file mode 100755
index 000000000..412631ba1
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/pretrained_ctc.py
@@ -0,0 +1,440 @@
+#!/usr/bin/env python3
+# Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang,
+# Zengwei Yao)
+#
+# 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.
+"""
+This script loads torchscript models, exported by `torch.jit.script()`
+and uses them to decode waves.
+You can use the following command to get the exported models:
+
+./pruned_transducer_stateless7_ctc_bs/export.py \
+ --exp-dir ./pruned_transducer_stateless7_ctc_bs/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 20 \
+ --avg 10
+
+Usage of this script:
+
+(1) ctc-decoding
+./pruned_transducer_stateless7_ctc_bs/jit_pretrained_ctc.py \
+ --checkpoint ./pruned_transducer_stateless7_ctc_bs/exp/pretrained.pt \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --method ctc-decoding \
+ --sample-rate 16000 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+(2) 1best
+./pruned_transducer_stateless7_ctc_bs/jit_pretrained_ctc.py \
+ --checkpoint ./pruned_transducer_stateless7_ctc_bs/exp/pretrained.pt \
+ --HLG data/lang_bpe_500/HLG.pt \
+ --words-file data/lang_bpe_500/words.txt \
+ --method 1best \
+ --sample-rate 16000 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+(3) nbest-rescoring
+./bruned_transducer_stateless7_ctc/jit_pretrained_ctc.py \
+ --checkpoint ./pruned_transducer_stateless7_ctc_bs/exp/pretrained.pt \
+ --HLG data/lang_bpe_500/HLG.pt \
+ --words-file data/lang_bpe_500/words.txt \
+ --G data/lm/G_4_gram.pt \
+ --method nbest-rescoring \
+ --sample-rate 16000 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+
+(4) whole-lattice-rescoring
+./pruned_transducer_stateless7_ctc_bs/jit_pretrained_ctc.py \
+ --checkpoint ./pruned_transducer_stateless7_ctc_bs/exp/pretrained.pt \
+ --HLG data/lang_bpe_500/HLG.pt \
+ --words-file data/lang_bpe_500/words.txt \
+ --G data/lm/G_4_gram.pt \
+ --method whole-lattice-rescoring \
+ --sample-rate 16000 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+"""
+
+import argparse
+import logging
+import math
+from typing import List
+
+import k2
+import kaldifeat
+import sentencepiece as spm
+import torch
+import torchaudio
+from ctc_decode import get_decoding_params
+from torch.nn.utils.rnn import pad_sequence
+from train import add_model_arguments, get_params, get_transducer_model
+
+from icefall.decode import (
+ get_lattice,
+ one_best_decoding,
+ rescore_with_n_best_list,
+ rescore_with_whole_lattice,
+)
+from icefall.utils import get_texts
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--checkpoint",
+ type=str,
+ required=True,
+ help="Path to the checkpoint. "
+ "The checkpoint is assumed to be saved by "
+ "icefall.checkpoint.save_checkpoint().",
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram, 2 means tri-gram",
+ )
+
+ parser.add_argument(
+ "--words-file",
+ type=str,
+ help="""Path to words.txt.
+ Used only when method is not ctc-decoding.
+ """,
+ )
+
+ parser.add_argument(
+ "--HLG",
+ type=str,
+ help="""Path to HLG.pt.
+ Used only when method is not ctc-decoding.
+ """,
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ help="""Path to bpe.model.
+ Used only when method is ctc-decoding.
+ """,
+ )
+
+ parser.add_argument(
+ "--method",
+ type=str,
+ default="1best",
+ help="""Decoding method.
+ Possible values are:
+ (0) ctc-decoding - Use CTC decoding. It uses a sentence
+ piece model, i.e., lang_dir/bpe.model, to convert
+ word pieces to words. It needs neither a lexicon
+ nor an n-gram LM.
+ (1) 1best - Use the best path as decoding output. Only
+ the transformer encoder output is used for decoding.
+ We call it HLG decoding.
+ (2) nbest-rescoring. Extract n paths from the decoding lattice,
+ rescore them with an LM, the path with
+ the highest score is the decoding result.
+ We call it HLG decoding + n-gram LM rescoring.
+ (3) whole-lattice-rescoring - Use an LM to rescore the
+ decoding lattice and then use 1best to decode the
+ rescored lattice.
+ We call it HLG decoding + n-gram LM rescoring.
+ """,
+ )
+
+ parser.add_argument(
+ "--G",
+ type=str,
+ help="""An LM for rescoring.
+ Used only when method is
+ whole-lattice-rescoring or nbest-rescoring.
+ It's usually a 4-gram LM.
+ """,
+ )
+
+ parser.add_argument(
+ "--num-paths",
+ type=int,
+ default=100,
+ help="""
+ Used only when method is attention-decoder.
+ It specifies the size of n-best list.""",
+ )
+
+ parser.add_argument(
+ "--ngram-lm-scale",
+ type=float,
+ default=1.3,
+ help="""
+ Used only when method is whole-lattice-rescoring and nbest-rescoring.
+ It specifies the scale for n-gram LM scores.
+ (Note: You need to tune it on a dataset.)
+ """,
+ )
+
+ parser.add_argument(
+ "--nbest-scale",
+ type=float,
+ default=0.5,
+ help="""
+ Used only when method is nbest-rescoring.
+ It specifies the scale for lattice.scores when
+ extracting n-best lists. A smaller value results in
+ more unique number of paths with the risk of missing
+ the best path.
+ """,
+ )
+
+ parser.add_argument(
+ "--num-classes",
+ type=int,
+ default=500,
+ help="""
+ Vocab size in the BPE model.
+ """,
+ )
+
+ parser.add_argument(
+ "--sample-rate",
+ type=int,
+ default=16000,
+ help="The sample rate of the input sound file",
+ )
+
+ parser.add_argument(
+ "sound_files",
+ type=str,
+ nargs="+",
+ help="The input sound file(s) to transcribe. "
+ "Supported formats are those supported by torchaudio.load(). "
+ "For example, wav and flac are supported. "
+ "The sample rate has to be 16kHz.",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def read_sound_files(
+ filenames: List[str], expected_sample_rate: float = 16000
+) -> List[torch.Tensor]:
+ """Read a list of sound files into a list 1-D float32 torch tensors.
+ Args:
+ filenames:
+ A list of sound filenames.
+ expected_sample_rate:
+ The expected sample rate of the sound files.
+ Returns:
+ Return a list of 1-D float32 torch tensors.
+ """
+ ans = []
+ for f in filenames:
+ wave, sample_rate = torchaudio.load(f)
+ assert sample_rate == expected_sample_rate, (
+ f"expected sample rate: {expected_sample_rate}. " f"Given: {sample_rate}"
+ )
+ # We use only the first channel
+ ans.append(wave[0])
+ return ans
+
+
+@torch.no_grad()
+def main():
+ parser = get_parser()
+ args = parser.parse_args()
+
+ params = get_params()
+ # add decoding params
+ params.update(get_decoding_params())
+ params.update(vars(args))
+ params.vocab_size = params.num_classes
+ params.blank_id = 0
+
+ logging.info(f"{params}")
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"device: {device}")
+
+ logging.info("Creating model")
+ model = get_transducer_model(params)
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ checkpoint = torch.load(args.checkpoint, map_location="cpu")
+ model.load_state_dict(checkpoint["model"], strict=False)
+ model.to(device)
+ model.eval()
+
+ logging.info("Constructing Fbank computer")
+ opts = kaldifeat.FbankOptions()
+ opts.device = device
+ opts.frame_opts.dither = 0
+ opts.frame_opts.snip_edges = False
+ opts.frame_opts.samp_freq = params.sample_rate
+ opts.mel_opts.num_bins = params.feature_dim
+
+ fbank = kaldifeat.Fbank(opts)
+
+ logging.info(f"Reading sound files: {params.sound_files}")
+ waves = read_sound_files(
+ filenames=params.sound_files, expected_sample_rate=params.sample_rate
+ )
+ waves = [w.to(device) for w in waves]
+
+ logging.info("Decoding started")
+ features = fbank(waves)
+ feature_lengths = [f.size(0) for f in features]
+
+ features = pad_sequence(features, batch_first=True, padding_value=math.log(1e-10))
+ feature_lengths = torch.tensor(feature_lengths, device=device)
+
+ encoder_out, encoder_out_lens = model.encoder(
+ x=features,
+ x_lens=feature_lengths,
+ )
+ nnet_output = model.ctc_output(encoder_out)
+
+ batch_size = nnet_output.shape[0]
+ supervision_segments = torch.tensor(
+ [[i, 0, nnet_output.shape[1]] for i in range(batch_size)],
+ dtype=torch.int32,
+ )
+
+ if params.method == "ctc-decoding":
+ logging.info("Use CTC decoding")
+ bpe_model = spm.SentencePieceProcessor()
+ bpe_model.load(params.bpe_model)
+ max_token_id = params.num_classes - 1
+
+ H = k2.ctc_topo(
+ max_token=max_token_id,
+ modified=False,
+ device=device,
+ )
+
+ lattice = get_lattice(
+ nnet_output=nnet_output,
+ decoding_graph=H,
+ 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,
+ )
+
+ best_path = one_best_decoding(
+ lattice=lattice, use_double_scores=params.use_double_scores
+ )
+ token_ids = get_texts(best_path)
+ hyps = bpe_model.decode(token_ids)
+ hyps = [s.split() for s in hyps]
+ elif params.method in [
+ "1best",
+ "nbest-rescoring",
+ "whole-lattice-rescoring",
+ ]:
+ logging.info(f"Loading HLG from {params.HLG}")
+ HLG = k2.Fsa.from_dict(torch.load(params.HLG, map_location="cpu"))
+ HLG = HLG.to(device)
+ if not hasattr(HLG, "lm_scores"):
+ # For whole-lattice-rescoring and attention-decoder
+ HLG.lm_scores = HLG.scores.clone()
+
+ if params.method in [
+ "nbest-rescoring",
+ "whole-lattice-rescoring",
+ ]:
+ logging.info(f"Loading G from {params.G}")
+ G = k2.Fsa.from_dict(torch.load(params.G, map_location="cpu"))
+ G = G.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.lm_scores is used to replace HLG.lm_scores during
+ # LM rescoring.
+ G.lm_scores = G.scores.clone()
+
+ lattice = get_lattice(
+ nnet_output=nnet_output,
+ decoding_graph=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 == "1best":
+ logging.info("Use HLG decoding")
+ best_path = one_best_decoding(
+ lattice=lattice, use_double_scores=params.use_double_scores
+ )
+ if params.method == "nbest-rescoring":
+ logging.info("Use HLG decoding + LM rescoring")
+ best_path_dict = rescore_with_n_best_list(
+ lattice=lattice,
+ G=G,
+ num_paths=params.num_paths,
+ lm_scale_list=[params.ngram_lm_scale],
+ nbest_scale=params.nbest_scale,
+ )
+ best_path = next(iter(best_path_dict.values()))
+ elif params.method == "whole-lattice-rescoring":
+ logging.info("Use HLG decoding + LM rescoring")
+ best_path_dict = rescore_with_whole_lattice(
+ lattice=lattice,
+ G_with_epsilon_loops=G,
+ lm_scale_list=[params.ngram_lm_scale],
+ )
+ best_path = next(iter(best_path_dict.values()))
+
+ hyps = get_texts(best_path)
+ word_sym_table = k2.SymbolTable.from_file(params.words_file)
+ hyps = [[word_sym_table[i] for i in ids] for ids in hyps]
+ else:
+ raise ValueError(f"Unsupported decoding method: {params.method}")
+
+ s = "\n"
+ for filename, hyp in zip(params.sound_files, hyps):
+ words = " ".join(hyp)
+ s += f"{filename}:\n{words}\n\n"
+ logging.info(s)
+
+ logging.info("Decoding Done")
+
+
+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/pruned_transducer_stateless7_ctc_bs/scaling.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/scaling.py
new file mode 120000
index 000000000..2428b74b9
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/scaling.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/scaling.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/scaling_converter.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/scaling_converter.py
new file mode 120000
index 000000000..b8b8ba432
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/scaling_converter.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/scaling_converter.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/test_model.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/test_model.py
new file mode 100755
index 000000000..7f0893985
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/test_model.py
@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+# Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang)
+#
+# 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.
+
+
+"""
+To run this file, do:
+
+ cd icefall/egs/librispeech/ASR
+ python ./pruned_transducer_stateless7_ctc_bs/test_model.py
+"""
+
+from train import get_params, get_transducer_model
+
+
+def test_model_1():
+ params = get_params()
+ params.vocab_size = 500
+ params.blank_id = 0
+ params.context_size = 2
+ params.num_encoder_layers = "2,4,3,2,4"
+ params.feedforward_dims = "1024,1024,2048,2048,1024"
+ params.nhead = "8,8,8,8,8"
+ params.encoder_dims = "384,384,384,384,384"
+ params.attention_dims = "192,192,192,192,192"
+ params.encoder_unmasked_dims = "256,256,256,256,256"
+ params.zipformer_downsampling_factors = "1,2,4,8,2"
+ params.cnn_module_kernels = "31,31,31,31,31"
+ params.decoder_dim = 512
+ params.joiner_dim = 512
+ model = get_transducer_model(params)
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ print(f"Number of model parameters: {num_param}")
+
+
+def main():
+ test_model_1()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/train.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/train.py
new file mode 100755
index 000000000..522ecc974
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/train.py
@@ -0,0 +1,1277 @@
+#!/usr/bin/env python3
+# Copyright 2021-2022 Xiaomi Corp. (authors: Fangjun Kuang,
+# Wei Kang,
+# Mingshuang Luo,
+# Zengwei Yao)
+#
+# 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.
+"""
+Usage:
+export CUDA_VISIBLE_DEVICES="0,1,2,3"
+./pruned_transducer_stateless7_ctc_bs/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --exp-dir pruned_transducer_stateless7_ctc_bs/exp \
+ --full-libri 1 \
+ --max-duration 300
+# For mix precision training:
+./pruned_transducer_stateless7_ctc_bs/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --use-fp16 1 \
+ --exp-dir pruned_transducer_stateless7_ctc_bs/exp \
+ --full-libri 1 \
+ --max-duration 550
+"""
+
+
+import argparse
+import copy
+import logging
+import warnings
+from pathlib import Path
+from shutil import copyfile
+from typing import Any, Dict, Optional, Tuple, Union
+
+import k2
+import optim
+import sentencepiece as spm
+import torch
+import torch.multiprocessing as mp
+import torch.nn as nn
+from asr_datamodule import LibriSpeechAsrDataModule
+from decoder import Decoder
+from joiner import Joiner
+from lconv import LConv
+from frame_reducer import FrameReducer
+from lhotse.cut import Cut
+from lhotse.dataset.sampling.base import CutSampler
+from lhotse.utils import fix_random_seed
+from model import Transducer
+from optim import Eden, ScaledAdam
+from torch import Tensor
+from torch.cuda.amp import GradScaler
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.tensorboard import SummaryWriter
+from zipformer import Zipformer
+
+from icefall import diagnostics
+from icefall.checkpoint import load_checkpoint, remove_checkpoints
+from icefall.checkpoint import save_checkpoint as save_checkpoint_impl
+from icefall.checkpoint import (
+ save_checkpoint_with_global_batch_idx,
+ update_averaged_model,
+)
+from icefall.dist import cleanup_dist, setup_dist
+from icefall.env import get_env_info
+from icefall.hooks import register_inf_check_hooks
+from icefall.utils import (
+ AttributeDict,
+ MetricsTracker,
+ encode_supervisions,
+ setup_logger,
+ str2bool,
+)
+
+LRSchedulerType = Union[torch.optim.lr_scheduler._LRScheduler, optim.LRScheduler]
+
+
+def set_batch_count(model: Union[nn.Module, DDP], batch_count: float) -> None:
+ if isinstance(model, DDP):
+ # get underlying nn.Module
+ model = model.module
+ for module in model.modules():
+ if hasattr(module, "batch_count"):
+ module.batch_count = batch_count
+
+
+def add_model_arguments(parser: argparse.ArgumentParser):
+ parser.add_argument(
+ "--num-encoder-layers",
+ type=str,
+ default="2,4,3,2,4",
+ help="Number of zipformer encoder layers, comma separated.",
+ )
+
+ parser.add_argument(
+ "--feedforward-dims",
+ type=str,
+ default="1024,1024,2048,2048,1024",
+ help="Feedforward dimension of the zipformer encoder layers, comma separated.",
+ )
+
+ parser.add_argument(
+ "--nhead",
+ type=str,
+ default="8,8,8,8,8",
+ help="Number of attention heads in the zipformer encoder layers.",
+ )
+
+ parser.add_argument(
+ "--encoder-dims",
+ type=str,
+ default="384,384,384,384,384",
+ help="Embedding dimension in the 2 blocks of zipformer encoder layers, comma separated",
+ )
+
+ parser.add_argument(
+ "--attention-dims",
+ type=str,
+ default="192,192,192,192,192",
+ help="""Attention dimension in the 2 blocks of zipformer encoder layers, comma separated;
+ not the same as embedding dimension.""",
+ )
+
+ parser.add_argument(
+ "--encoder-unmasked-dims",
+ type=str,
+ default="256,256,256,256,256",
+ help="Unmasked dimensions in the encoders, relates to augmentation during training. "
+ "Must be <= each of encoder_dims. Empirically, less than 256 seems to make performance "
+ " worse.",
+ )
+
+ parser.add_argument(
+ "--zipformer-downsampling-factors",
+ type=str,
+ default="1,2,4,8,2",
+ help="Downsampling factor for each stack of encoder layers.",
+ )
+
+ parser.add_argument(
+ "--cnn-module-kernels",
+ type=str,
+ default="31,31,31,31,31",
+ help="Sizes of kernels in convolution modules",
+ )
+
+ parser.add_argument(
+ "--decoder-dim",
+ type=int,
+ default=512,
+ help="Embedding dimension in the decoder model.",
+ )
+
+ parser.add_argument(
+ "--joiner-dim",
+ type=int,
+ default=512,
+ help="""Dimension used in the joiner model.
+ Outputs from the encoder and decoder model are projected
+ to this dimension before adding.
+ """,
+ )
+
+
+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.",
+ )
+
+ parser.add_argument(
+ "--num-epochs",
+ type=int,
+ default=30,
+ help="Number of epochs to train.",
+ )
+
+ parser.add_argument(
+ "--start-epoch",
+ type=int,
+ default=1,
+ help="""Resume training from this epoch. It should be positive.
+ If larger than 1, it will load checkpoint from
+ exp-dir/epoch-{start_epoch-1}.pt
+ """,
+ )
+
+ parser.add_argument(
+ "--start-batch",
+ type=int,
+ default=0,
+ help="""If positive, --start-epoch is ignored and
+ it loads the checkpoint from exp-dir/checkpoint-{start_batch}.pt
+ """,
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="pruned_transducer_stateless7_ctc_bs/exp",
+ help="""The experiment dir.
+ It specifies the directory where all training related
+ files, e.g., checkpoints, log, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ default="data/lang_bpe_500/bpe.model",
+ help="Path to the BPE model",
+ )
+
+ parser.add_argument(
+ "--base-lr", type=float, default=0.05, help="The base learning rate."
+ )
+
+ parser.add_argument(
+ "--lr-batches",
+ type=float,
+ default=5000,
+ help="""Number of steps that affects how rapidly the learning rate
+ decreases. We suggest not to change this.""",
+ )
+
+ parser.add_argument(
+ "--lr-epochs",
+ type=float,
+ default=3.5,
+ help="""Number of epochs that affects how rapidly the learning rate decreases.
+ """,
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram, 2 means tri-gram",
+ )
+
+ parser.add_argument(
+ "--prune-range",
+ type=int,
+ default=5,
+ help="The prune range for rnnt loss, it means how many symbols(context)"
+ "we are using to compute the loss",
+ )
+
+ parser.add_argument(
+ "--lm-scale",
+ type=float,
+ default=0.25,
+ help="The scale to smooth the loss with lm "
+ "(output of prediction network) part.",
+ )
+
+ parser.add_argument(
+ "--am-scale",
+ type=float,
+ default=0.0,
+ help="The scale to smooth the loss with am (output of encoder network) part.",
+ )
+
+ parser.add_argument(
+ "--simple-loss-scale",
+ type=float,
+ default=0.5,
+ help="To get pruning ranges, we will calculate a simple version"
+ "loss(joiner is just addition), this simple loss also uses for"
+ "training (as a regularization item). We will scale the simple loss"
+ "with this parameter before adding to the final loss.",
+ )
+
+ parser.add_argument(
+ "--ctc-loss-scale",
+ type=float,
+ default=0.5,
+ help="Scale for CTC loss.",
+ )
+
+ parser.add_argument(
+ "--seed",
+ type=int,
+ default=42,
+ help="The seed for random generators intended for reproducibility",
+ )
+
+ parser.add_argument(
+ "--print-diagnostics",
+ type=str2bool,
+ default=False,
+ help="Accumulate stats on activations, print them and exit.",
+ )
+
+ parser.add_argument(
+ "--inf-check",
+ type=str2bool,
+ default=False,
+ help="Add hooks to check for infinite module outputs and gradients.",
+ )
+
+ parser.add_argument(
+ "--save-every-n",
+ type=int,
+ default=2000,
+ help="""Save checkpoint after processing this number of batches"
+ periodically. We save checkpoint to exp-dir/ whenever
+ params.batch_idx_train % save_every_n == 0. The checkpoint filename
+ has the form: f'exp-dir/checkpoint-{params.batch_idx_train}.pt'
+ Note: It also saves checkpoint to `exp-dir/epoch-xxx.pt` at the
+ end of each epoch where `xxx` is the epoch number counting from 0.
+ """,
+ )
+
+ parser.add_argument(
+ "--keep-last-k",
+ type=int,
+ default=30,
+ help="""Only keep this number of checkpoints on disk.
+ For instance, if it is 3, there are only 3 checkpoints
+ in the exp-dir with filenames `checkpoint-xxx.pt`.
+ It does not affect checkpoints with name `epoch-xxx.pt`.
+ """,
+ )
+
+ parser.add_argument(
+ "--average-period",
+ type=int,
+ default=200,
+ help="""Update the averaged model, namely `model_avg`, after processing
+ this number of batches. `model_avg` is a separate version of model,
+ in which each floating-point parameter is the average of all the
+ parameters from the start of training. Each time we take the average,
+ we do: `model_avg = model * (average_period / batch_idx_train) +
+ model_avg * ((batch_idx_train - average_period) / batch_idx_train)`.
+ """,
+ )
+
+ parser.add_argument(
+ "--use-fp16",
+ type=str2bool,
+ default=False,
+ help="Whether to use half precision training.",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def get_params() -> AttributeDict:
+ """Return a dict containing training parameters.
+ All training related parameters that are not passed from the commandline
+ are 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`:
+ - 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
+ - reset_interval: Reset statistics if batch_idx % reset_interval is 0
+ - valid_interval: Run validation if batch_idx % valid_interval is 0
+ - feature_dim: The model input dim. It has to match the one used
+ in computing features.
+ - subsampling_factor: The subsampling factor for the model.
+ - encoder_dim: Hidden dim for multi-head attention model.
+ - num_decoder_layers: Number of decoder layer of transformer decoder.
+ - warm_step: The warmup period that dictates the decay of the
+ scale on "simple" (un-pruned) loss.
+ """
+ params = AttributeDict(
+ {
+ "best_train_loss": float("inf"),
+ "best_valid_loss": float("inf"),
+ "best_train_epoch": -1,
+ "best_valid_epoch": -1,
+ "batch_idx_train": 0,
+ "log_interval": 50,
+ "reset_interval": 200,
+ "valid_interval": 3000, # For the 100h subset, use 800
+ # parameters for zipformer
+ "feature_dim": 80,
+ "subsampling_factor": 4, # not passed in, this is fixed.
+ # parameters for ctc loss
+ "beam_size": 10,
+ "use_double_scores": True,
+ "warm_step": 2000,
+ "env_info": get_env_info(),
+ }
+ )
+
+ return params
+
+
+def get_encoder_model(params: AttributeDict) -> nn.Module:
+ # TODO: We can add an option to switch between Zipformer and Transformer
+ def to_int_tuple(s: str):
+ return tuple(map(int, s.split(",")))
+
+ encoder = Zipformer(
+ num_features=params.feature_dim,
+ output_downsampling_factor=2,
+ zipformer_downsampling_factors=to_int_tuple(
+ params.zipformer_downsampling_factors
+ ),
+ encoder_dims=to_int_tuple(params.encoder_dims),
+ attention_dim=to_int_tuple(params.attention_dims),
+ encoder_unmasked_dims=to_int_tuple(params.encoder_unmasked_dims),
+ nhead=to_int_tuple(params.nhead),
+ feedforward_dim=to_int_tuple(params.feedforward_dims),
+ cnn_module_kernels=to_int_tuple(params.cnn_module_kernels),
+ num_encoder_layers=to_int_tuple(params.num_encoder_layers),
+ )
+ return encoder
+
+
+def get_decoder_model(params: AttributeDict) -> nn.Module:
+ decoder = Decoder(
+ vocab_size=params.vocab_size,
+ decoder_dim=params.decoder_dim,
+ blank_id=params.blank_id,
+ context_size=params.context_size,
+ )
+ return decoder
+
+
+def get_joiner_model(params: AttributeDict) -> nn.Module:
+ joiner = Joiner(
+ encoder_dim=int(params.encoder_dims.split(",")[-1]),
+ decoder_dim=params.decoder_dim,
+ joiner_dim=params.joiner_dim,
+ vocab_size=params.vocab_size,
+ )
+ return joiner
+
+
+def get_lconv(params: AttributeDict) -> nn.Module:
+ lconv = LConv(
+ channels=int(params.encoder_dims.split(",")[-1]),
+ )
+ return lconv
+
+
+def get_frame_reducer(params: AttributeDict) -> nn.Module:
+ frame_reducer = FrameReducer()
+ return frame_reducer
+
+
+def get_transducer_model(params: AttributeDict) -> nn.Module:
+ encoder = get_encoder_model(params)
+ decoder = get_decoder_model(params)
+ joiner = get_joiner_model(params)
+ lconv = get_lconv(params)
+ frame_reducer = get_frame_reducer(params)
+
+ model = Transducer(
+ encoder=encoder,
+ decoder=decoder,
+ joiner=joiner,
+ lconv=lconv,
+ frame_reducer=frame_reducer,
+ encoder_dim=int(params.encoder_dims.split(",")[-1]),
+ decoder_dim=params.decoder_dim,
+ joiner_dim=params.joiner_dim,
+ vocab_size=params.vocab_size,
+ )
+ return model
+
+
+def load_checkpoint_if_available(
+ params: AttributeDict,
+ model: nn.Module,
+ model_avg: nn.Module = None,
+ optimizer: Optional[torch.optim.Optimizer] = None,
+ scheduler: Optional[LRSchedulerType] = None,
+) -> Optional[Dict[str, Any]]:
+ """Load checkpoint from file.
+ If params.start_batch is positive, it will load the checkpoint from
+ `params.exp_dir/checkpoint-{params.start_batch}.pt`. Otherwise, if
+ params.start_epoch is larger than 1, it will load the checkpoint from
+ `params.start_epoch - 1`.
+ Apart from loading state dict for `model` and `optimizer` 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.
+ model_avg:
+ The stored model averaged from the start of training.
+ optimizer:
+ The optimizer that we are using.
+ scheduler:
+ The scheduler that we are using.
+ Returns:
+ Return a dict containing previously saved training info.
+ """
+ if params.start_batch > 0:
+ filename = params.exp_dir / f"checkpoint-{params.start_batch}.pt"
+ elif params.start_epoch > 1:
+ filename = params.exp_dir / f"epoch-{params.start_epoch-1}.pt"
+ else:
+ return None
+
+ assert filename.is_file(), f"{filename} does not exist!"
+
+ saved_params = load_checkpoint(
+ filename,
+ model=model,
+ model_avg=model_avg,
+ 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]
+
+ if params.start_batch > 0:
+ if "cur_epoch" in saved_params:
+ params["start_epoch"] = saved_params["cur_epoch"]
+
+ if "cur_batch_idx" in saved_params:
+ params["cur_batch_idx"] = saved_params["cur_batch_idx"]
+
+ return saved_params
+
+
+def save_checkpoint(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ model_avg: Optional[nn.Module] = None,
+ optimizer: Optional[torch.optim.Optimizer] = None,
+ scheduler: Optional[LRSchedulerType] = None,
+ sampler: Optional[CutSampler] = None,
+ scaler: Optional[GradScaler] = 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.
+ model_avg:
+ The stored model averaged from the start of training.
+ optimizer:
+ The optimizer used in the training.
+ sampler:
+ The sampler for the training dataset.
+ scaler:
+ The scaler used for mix precision training.
+ """
+ if rank != 0:
+ return
+ filename = params.exp_dir / f"epoch-{params.cur_epoch}.pt"
+ save_checkpoint_impl(
+ filename=filename,
+ model=model,
+ model_avg=model_avg,
+ params=params,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=sampler,
+ scaler=scaler,
+ 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: Union[nn.Module, DDP],
+ sp: spm.SentencePieceProcessor,
+ batch: dict,
+ is_training: bool,
+) -> Tuple[Tensor, MetricsTracker]:
+ """
+ Compute transducer 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 Zipformer in our case.
+ batch:
+ A batch of data. See `lhotse.dataset.K2SpeechRecognitionDataset()`
+ for the content in it.
+ 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 = model.device if isinstance(model, DDP) else next(model.parameters()).device
+ feature = batch["inputs"]
+ # at entry, feature is (N, T, C)
+ assert feature.ndim == 3
+ feature = feature.to(device)
+
+ supervisions = batch["supervisions"]
+ feature_lens = supervisions["num_frames"].to(device)
+
+ batch_idx_train = params.batch_idx_train
+ warm_step = params.warm_step
+ warmup = batch_idx_train / warm_step
+
+ texts = batch["supervisions"]["text"]
+ token_ids = sp.encode(texts, out_type=int)
+ y = k2.RaggedTensor(token_ids).to(device)
+
+ with torch.set_grad_enabled(is_training):
+ simple_loss, pruned_loss, ctc_output = model(
+ x=feature,
+ x_lens=feature_lens,
+ y=y,
+ prune_range=params.prune_range,
+ am_scale=params.am_scale,
+ lm_scale=params.lm_scale,
+ warmup=warmup,
+ )
+
+ s = params.simple_loss_scale
+ # take down the scale on the simple loss from 1.0 at the start
+ # to params.simple_loss scale by warm_step.
+ simple_loss_scale = (
+ s
+ if batch_idx_train >= warm_step
+ else 1.0 - (batch_idx_train / warm_step) * (1.0 - s)
+ )
+ pruned_loss_scale = (
+ 1.0
+ if batch_idx_train >= warm_step
+ else 0.1 + 0.9 * (batch_idx_train / warm_step)
+ )
+
+ loss = simple_loss_scale * simple_loss + pruned_loss_scale * pruned_loss
+
+ # Compute ctc loss
+
+ # NOTE: We need `encode_supervisions` to sort sequences with
+ # different duration in decreasing order, required by
+ # `k2.intersect_dense` called in `k2.ctc_loss`
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ supervision_segments, token_ids = encode_supervisions(
+ supervisions,
+ subsampling_factor=params.subsampling_factor,
+ token_ids=token_ids,
+ )
+
+ # Works with a BPE model
+ decoding_graph = k2.ctc_graph(token_ids, modified=False, device=device)
+ dense_fsa_vec = k2.DenseFsaVec(
+ ctc_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="sum",
+ use_double_scores=params.use_double_scores,
+ )
+ assert ctc_loss.requires_grad == is_training
+ loss += params.ctc_loss_scale * ctc_loss
+
+ assert loss.requires_grad == is_training
+
+ info = MetricsTracker()
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ info["frames"] = (feature_lens // params.subsampling_factor).sum().item()
+
+ # Note: We use reduction=sum while computing the loss.
+ info["loss"] = loss.detach().cpu().item()
+ info["simple_loss"] = simple_loss.detach().cpu().item()
+ info["pruned_loss"] = pruned_loss.detach().cpu().item()
+ info["ctc_loss"] = ctc_loss.detach().cpu().item()
+
+ return loss, info
+
+
+def compute_validation_loss(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ sp: spm.SentencePieceProcessor,
+ valid_dl: torch.utils.data.DataLoader,
+ world_size: int = 1,
+) -> MetricsTracker:
+ """Run the validation process."""
+ model.eval()
+
+ tot_loss = MetricsTracker()
+
+ for batch_idx, batch in enumerate(valid_dl):
+ loss, loss_info = compute_loss(
+ params=params,
+ model=model,
+ sp=sp,
+ batch=batch,
+ is_training=False,
+ )
+ assert loss.requires_grad is False
+ tot_loss = tot_loss + loss_info
+
+ if world_size > 1:
+ tot_loss.reduce(loss.device)
+
+ loss_value = tot_loss["loss"] / tot_loss["frames"]
+ if loss_value < params.best_valid_loss:
+ params.best_valid_epoch = params.cur_epoch
+ params.best_valid_loss = loss_value
+
+ return tot_loss
+
+
+def train_one_epoch(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ optimizer: torch.optim.Optimizer,
+ scheduler: LRSchedulerType,
+ sp: spm.SentencePieceProcessor,
+ train_dl: torch.utils.data.DataLoader,
+ valid_dl: torch.utils.data.DataLoader,
+ scaler: GradScaler,
+ model_avg: Optional[nn.Module] = None,
+ tb_writer: Optional[SummaryWriter] = None,
+ world_size: int = 1,
+ rank: int = 0,
+) -> 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.
+ scheduler:
+ The learning rate scheduler, we call step() every step.
+ train_dl:
+ Dataloader for the training dataset.
+ valid_dl:
+ Dataloader for the validation dataset.
+ scaler:
+ The scaler used for mix precision training.
+ model_avg:
+ The stored model averaged from the start of training.
+ tb_writer:
+ Writer to write log messages to tensorboard.
+ world_size:
+ Number of nodes in DDP training. If it is 1, DDP is disabled.
+ rank:
+ The rank of the node in DDP training. If no DDP is used, it should
+ be set to 0.
+ """
+ model.train()
+
+ tot_loss = MetricsTracker()
+
+ cur_batch_idx = params.get("cur_batch_idx", 0)
+
+ for batch_idx, batch in enumerate(train_dl):
+ if batch_idx < cur_batch_idx:
+ continue
+ cur_batch_idx = batch_idx
+
+ params.batch_idx_train += 1
+ batch_size = len(batch["supervisions"]["text"])
+
+ try:
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, loss_info = compute_loss(
+ params=params,
+ model=model,
+ sp=sp,
+ batch=batch,
+ is_training=True,
+ )
+ # summary stats
+ tot_loss = (tot_loss * (1 - 1 / params.reset_interval)) + loss_info
+
+ # NOTE: We use reduction==sum and loss is computed over utterances
+ # in the batch and there is no normalization to it so far.
+ scaler.scale(loss).backward()
+ set_batch_count(model, params.batch_idx_train)
+ scheduler.step_batch(params.batch_idx_train)
+
+ scaler.step(optimizer)
+ scaler.update()
+ optimizer.zero_grad()
+ except: # noqa
+ display_and_save_batch(batch, params=params, sp=sp)
+ raise
+
+ if params.print_diagnostics and batch_idx == 5:
+ return
+
+ if (
+ rank == 0
+ and params.batch_idx_train > 0
+ and params.batch_idx_train % params.average_period == 0
+ ):
+ update_averaged_model(
+ params=params,
+ model_cur=model,
+ model_avg=model_avg,
+ )
+
+ if (
+ params.batch_idx_train > 0
+ and params.batch_idx_train % params.save_every_n == 0
+ ):
+ params.cur_batch_idx = batch_idx
+ save_checkpoint_with_global_batch_idx(
+ out_dir=params.exp_dir,
+ global_batch_idx=params.batch_idx_train,
+ model=model,
+ model_avg=model_avg,
+ params=params,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=train_dl.sampler,
+ scaler=scaler,
+ rank=rank,
+ )
+ del params.cur_batch_idx
+ remove_checkpoints(
+ out_dir=params.exp_dir,
+ topk=params.keep_last_k,
+ rank=rank,
+ )
+
+ if batch_idx % 100 == 0 and params.use_fp16:
+ # If the grad scale was less than 1, try increasing it. The _growth_interval
+ # of the grad scaler is configurable, but we can't configure it to have different
+ # behavior depending on the current grad scale.
+ cur_grad_scale = scaler._scale.item()
+ if cur_grad_scale < 1.0 or (cur_grad_scale < 8.0 and batch_idx % 400 == 0):
+ scaler.update(cur_grad_scale * 2.0)
+ if cur_grad_scale < 0.01:
+ logging.warning(f"Grad scale is small: {cur_grad_scale}")
+ if cur_grad_scale < 1.0e-05:
+ raise RuntimeError(
+ f"grad_scale is too small, exiting: {cur_grad_scale}"
+ )
+
+ if batch_idx % params.log_interval == 0:
+ cur_lr = scheduler.get_last_lr()[0]
+ cur_grad_scale = scaler._scale.item() if params.use_fp16 else 1.0
+
+ logging.info(
+ f"Epoch {params.cur_epoch}, "
+ f"batch {batch_idx}, loss[{loss_info}], "
+ f"tot_loss[{tot_loss}], batch size: {batch_size}, "
+ f"lr: {cur_lr:.2e}, "
+ + (f"grad_scale: {scaler._scale.item()}" if params.use_fp16 else "")
+ )
+
+ if tb_writer is not None:
+ tb_writer.add_scalar(
+ "train/learning_rate", cur_lr, params.batch_idx_train
+ )
+
+ loss_info.write_summary(
+ tb_writer, "train/current_", params.batch_idx_train
+ )
+ tot_loss.write_summary(tb_writer, "train/tot_", params.batch_idx_train)
+ if params.use_fp16:
+ tb_writer.add_scalar(
+ "train/grad_scale",
+ cur_grad_scale,
+ params.batch_idx_train,
+ )
+
+ if batch_idx % params.valid_interval == 0 and not params.print_diagnostics:
+ logging.info("Computing validation loss")
+ valid_info = compute_validation_loss(
+ params=params,
+ model=model,
+ sp=sp,
+ valid_dl=valid_dl,
+ world_size=world_size,
+ )
+ model.train()
+ logging.info(f"Epoch {params.cur_epoch}, validation: {valid_info}")
+ logging.info(
+ f"Maximum memory allocated so far is {torch.cuda.max_memory_allocated()//1000000}MB"
+ )
+ if tb_writer is not None:
+ valid_info.write_summary(
+ tb_writer, "train/valid_", params.batch_idx_train
+ )
+
+ loss_value = tot_loss["loss"] / tot_loss["frames"]
+ params.train_loss = loss_value
+ 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))
+ if params.full_libri is False:
+ params.valid_interval = 1600
+
+ fix_random_seed(params.seed)
+ 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")
+
+ if args.tensorboard and rank == 0:
+ tb_writer = SummaryWriter(log_dir=f"{params.exp_dir}/tensorboard")
+ else:
+ tb_writer = None
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", rank)
+ logging.info(f"Device: {device}")
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(params.bpe_model)
+
+ # is defined in local/train_bpe_model.py
+ params.blank_id = sp.piece_to_id("")
+ params.vocab_size = sp.get_piece_size()
+
+ logging.info(params)
+
+ logging.info("About to create model")
+ model = get_transducer_model(params)
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ assert params.save_every_n >= params.average_period
+ model_avg: Optional[nn.Module] = None
+ if rank == 0:
+ # model_avg is only used with rank 0
+ model_avg = copy.deepcopy(model).to(torch.float64)
+
+ assert params.start_epoch > 0, params.start_epoch
+ checkpoints = load_checkpoint_if_available(
+ params=params, model=model, model_avg=model_avg
+ )
+
+ model.to(device)
+ if world_size > 1:
+ logging.info("Using DDP")
+ model = DDP(model, device_ids=[rank], find_unused_parameters=True)
+
+ parameters_names = []
+ parameters_names.append(
+ [name_param_pair[0] for name_param_pair in model.named_parameters()]
+ )
+
+ optimizer = ScaledAdam(
+ model.parameters(),
+ lr=params.base_lr,
+ clipping_scale=2.0,
+ parameters_names=parameters_names,
+ )
+
+ scheduler = Eden(optimizer, params.lr_batches, params.lr_epochs)
+
+ if checkpoints and "optimizer" in checkpoints:
+ logging.info("Loading optimizer state dict")
+ optimizer.load_state_dict(checkpoints["optimizer"])
+
+ if (
+ checkpoints
+ and "scheduler" in checkpoints
+ and checkpoints["scheduler"] is not None
+ ):
+ logging.info("Loading scheduler state dict")
+ scheduler.load_state_dict(checkpoints["scheduler"])
+
+ if params.print_diagnostics:
+ opts = diagnostics.TensorDiagnosticOptions(
+ 2**22
+ ) # allow 4 megabytes per sub-module
+ diagnostic = diagnostics.attach_diagnostics(model, opts)
+
+ if params.inf_check:
+ register_inf_check_hooks(model)
+
+ librispeech = LibriSpeechAsrDataModule(args)
+
+ train_cuts = librispeech.train_clean_100_cuts()
+ if params.full_libri:
+ train_cuts += librispeech.train_clean_360_cuts()
+ train_cuts += librispeech.train_other_500_cuts()
+
+ def remove_short_and_long_utt(c: Cut):
+ # Keep only utterances with duration between 1 second and 20 seconds
+ #
+ # Caution: There is a reason to select 20.0 here. Please see
+ # ../local/display_manifest_statistics.py
+ #
+ # You should use ../local/display_manifest_statistics.py to get
+ # an utterance duration distribution for your dataset to select
+ # the threshold
+ if c.duration < 1.0 or c.duration > 20.0:
+ logging.warning(
+ f"Exclude cut with ID {c.id} from training. Duration: {c.duration}"
+ )
+ return False
+
+ # In pruned RNN-T, we require that T >= S
+ # where T is the number of feature frames after subsampling
+ # and S is the number of tokens in the utterance
+
+ # In ./zipformer.py, the conv module uses the following expression
+ # for subsampling
+ T = ((c.num_frames - 7) // 2 + 1) // 2
+ tokens = sp.encode(c.supervisions[0].text, out_type=str)
+
+ if T < len(tokens):
+ logging.warning(
+ f"Exclude cut with ID {c.id} from training. "
+ f"Number of frames (before subsampling): {c.num_frames}. "
+ f"Number of frames (after subsampling): {T}. "
+ f"Text: {c.supervisions[0].text}. "
+ f"Tokens: {tokens}. "
+ f"Number of tokens: {len(tokens)}"
+ )
+ return False
+
+ return True
+
+ train_cuts = train_cuts.filter(remove_short_and_long_utt)
+
+ if params.start_batch > 0 and checkpoints and "sampler" in checkpoints:
+ # We only load the sampler's state dict when it loads a checkpoint
+ # saved in the middle of an epoch
+ sampler_state_dict = checkpoints["sampler"]
+ else:
+ sampler_state_dict = None
+
+ train_dl = librispeech.train_dataloaders(
+ train_cuts, sampler_state_dict=sampler_state_dict
+ )
+
+ valid_cuts = librispeech.dev_clean_cuts()
+ valid_cuts += librispeech.dev_other_cuts()
+ valid_dl = librispeech.valid_dataloaders(valid_cuts)
+
+ if not params.print_diagnostics:
+ scan_pessimistic_batches_for_oom(
+ model=model,
+ train_dl=train_dl,
+ optimizer=optimizer,
+ sp=sp,
+ params=params,
+ )
+
+ scaler = GradScaler(enabled=params.use_fp16, init_scale=1.0)
+ if checkpoints and "grad_scaler" in checkpoints:
+ logging.info("Loading grad scaler state dict")
+ scaler.load_state_dict(checkpoints["grad_scaler"])
+
+ for epoch in range(params.start_epoch, params.num_epochs + 1):
+ scheduler.step_epoch(epoch - 1)
+ fix_random_seed(params.seed + epoch - 1)
+ train_dl.sampler.set_epoch(epoch - 1)
+
+ if tb_writer is not None:
+ tb_writer.add_scalar("train/epoch", epoch, params.batch_idx_train)
+
+ params.cur_epoch = epoch
+
+ train_one_epoch(
+ params=params,
+ model=model,
+ model_avg=model_avg,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sp=sp,
+ train_dl=train_dl,
+ valid_dl=valid_dl,
+ scaler=scaler,
+ tb_writer=tb_writer,
+ world_size=world_size,
+ rank=rank,
+ )
+
+ if params.print_diagnostics:
+ diagnostic.print_diagnostics()
+ break
+
+ save_checkpoint(
+ params=params,
+ model=model,
+ model_avg=model_avg,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=train_dl.sampler,
+ scaler=scaler,
+ rank=rank,
+ )
+
+ logging.info("Done!")
+
+ if world_size > 1:
+ torch.distributed.barrier()
+ cleanup_dist()
+
+
+def display_and_save_batch(
+ batch: dict,
+ params: AttributeDict,
+ sp: spm.SentencePieceProcessor,
+) -> None:
+ """Display the batch statistics and save the batch into disk.
+ Args:
+ batch:
+ A batch of data. See `lhotse.dataset.K2SpeechRecognitionDataset()`
+ for the content in it.
+ params:
+ Parameters for training. See :func:`get_params`.
+ sp:
+ The BPE model.
+ """
+ from lhotse.utils import uuid4
+
+ filename = f"{params.exp_dir}/batch-{uuid4()}.pt"
+ logging.info(f"Saving batch to {filename}")
+ torch.save(batch, filename)
+
+ supervisions = batch["supervisions"]
+ features = batch["inputs"]
+
+ logging.info(f"features shape: {features.shape}")
+
+ y = sp.encode(supervisions["text"], out_type=int)
+ num_tokens = sum(len(i) for i in y)
+ logging.info(f"num tokens: {num_tokens}")
+
+
+def scan_pessimistic_batches_for_oom(
+ model: Union[nn.Module, DDP],
+ train_dl: torch.utils.data.DataLoader,
+ optimizer: torch.optim.Optimizer,
+ sp: spm.SentencePieceProcessor,
+ params: AttributeDict,
+):
+ from lhotse.dataset import find_pessimistic_batches
+
+ logging.info(
+ "Sanity check -- see if any of the batches in epoch 1 would cause OOM."
+ )
+ batches, crit_values = find_pessimistic_batches(train_dl.sampler)
+ for criterion, cuts in batches.items():
+ batch = train_dl.dataset[cuts]
+ try:
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, _ = compute_loss(
+ params=params,
+ model=model,
+ sp=sp,
+ batch=batch,
+ is_training=True,
+ )
+ loss.backward()
+ optimizer.zero_grad()
+ except Exception as e:
+ if "CUDA out of memory" in str(e):
+ logging.error(
+ "Your GPU ran out of memory with the current "
+ "max_duration setting. We recommend decreasing "
+ "max_duration and trying again.\n"
+ f"Failing criterion: {criterion} "
+ f"(={crit_values[criterion]}) ..."
+ )
+ display_and_save_batch(batch, params=params, sp=sp)
+ raise
+ logging.info(
+ f"Maximum memory allocated so far is {torch.cuda.max_memory_allocated()//1000000}MB"
+ )
+
+
+def main():
+ parser = get_parser()
+ LibriSpeechAsrDataModule.add_arguments(parser)
+ args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ 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)
+
+
+torch.set_num_threads(1)
+torch.set_num_interop_threads(1)
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/zipformer.py b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/zipformer.py
new file mode 120000
index 000000000..79b076556
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_ctc_bs/zipformer.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/zipformer.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/README.md b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/README.md
new file mode 100644
index 000000000..6e461e196
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/README.md
@@ -0,0 +1,3 @@
+This recipe implements Streaming Zipformer-Transducer model.
+
+See https://k2-fsa.github.io/icefall/recipes/Streaming-ASR/librispeech/zipformer_transducer.html for detailed tutorials.
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/asr_datamodule.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/asr_datamodule.py
new file mode 120000
index 000000000..a074d6085
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/asr_datamodule.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/asr_datamodule.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/beam_search.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/beam_search.py
new file mode 120000
index 000000000..8554e44cc
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/beam_search.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/beam_search.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/decode.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/decode.py
new file mode 100755
index 000000000..aebe2b94b
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/decode.py
@@ -0,0 +1,813 @@
+#!/usr/bin/env python3
+#
+# Copyright 2021-2022 Xiaomi Corporation (Author: Fangjun Kuang,
+# Zengwei Yao)
+#
+# 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.
+"""
+Usage:
+(1) greedy search
+./pruned_transducer_stateless7_streaming/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --max-duration 600 \
+ --decode-chunk-len 32 \
+ --decoding-method greedy_search
+
+(2) beam search (not recommended)
+./pruned_transducer_stateless7_streaming/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --max-duration 600 \
+ --decode-chunk-len 32 \
+ --decoding-method beam_search \
+ --beam-size 4
+
+(3) modified beam search
+./pruned_transducer_stateless7_streaming/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --max-duration 600 \
+ --decode-chunk-len 32 \
+ --decoding-method modified_beam_search \
+ --beam-size 4
+
+(4) fast beam search (one best)
+./pruned_transducer_stateless7_streaming/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --max-duration 600 \
+ --decode-chunk-len 32 \
+ --decoding-method fast_beam_search \
+ --beam 20.0 \
+ --max-contexts 8 \
+ --max-states 64
+
+(5) fast beam search (nbest)
+./pruned_transducer_stateless7_streaming/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --max-duration 600 \
+ --decode-chunk-len 32 \
+ --decoding-method fast_beam_search_nbest \
+ --beam 20.0 \
+ --max-contexts 8 \
+ --max-states 64 \
+ --num-paths 200 \
+ --nbest-scale 0.5
+
+(6) fast beam search (nbest oracle WER)
+./pruned_transducer_stateless7_streaming/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --max-duration 600 \
+ --decode-chunk-len 32 \
+ --decoding-method fast_beam_search_nbest_oracle \
+ --beam 20.0 \
+ --max-contexts 8 \
+ --max-states 64 \
+ --num-paths 200 \
+ --nbest-scale 0.5
+
+(7) fast beam search (with LG)
+./pruned_transducer_stateless7_streaming/decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --max-duration 600 \
+ --decode-chunk-len 32 \
+ --decoding-method fast_beam_search_nbest_LG \
+ --beam 20.0 \
+ --max-contexts 8 \
+ --max-states 64
+"""
+
+
+import argparse
+import logging
+import math
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import k2
+import sentencepiece as spm
+import torch
+import torch.nn as nn
+from asr_datamodule import LibriSpeechAsrDataModule
+from beam_search import (
+ beam_search,
+ fast_beam_search_nbest,
+ fast_beam_search_nbest_LG,
+ fast_beam_search_nbest_oracle,
+ fast_beam_search_one_best,
+ greedy_search,
+ greedy_search_batch,
+ modified_beam_search,
+)
+from train import add_model_arguments, get_params, get_transducer_model
+
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from icefall.lexicon import Lexicon
+from icefall.utils import (
+ AttributeDict,
+ setup_logger,
+ store_transcripts,
+ str2bool,
+ write_error_stats,
+)
+
+LOG_EPS = math.log(1e-10)
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=30,
+ help="""It specifies the checkpoint to use for decoding.
+ Note: Epoch counts from 1.
+ You can specify --avg to use more checkpoints for model averaging.""",
+ )
+
+ parser.add_argument(
+ "--iter",
+ type=int,
+ default=0,
+ help="""If positive, --epoch is ignored and it
+ will use the checkpoint exp_dir/checkpoint-iter.pt.
+ You can specify --avg to use more checkpoints for model averaging.
+ """,
+ )
+
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=9,
+ help="Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch' and '--iter'",
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=True,
+ help="Whether to load averaged model. Currently it only supports "
+ "using --epoch. If True, it would decode with the averaged model "
+ "over the epoch range from `epoch-avg` (excluded) to `epoch`."
+ "Actually only the models with epoch number of `epoch-avg` and "
+ "`epoch` are loaded for averaging. ",
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="pruned_transducer_stateless7_streaming/exp",
+ help="The experiment dir",
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ default="data/lang_bpe_500/bpe.model",
+ help="Path to the BPE model",
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=Path,
+ default="data/lang_bpe_500",
+ help="The lang dir containing word table and LG graph",
+ )
+
+ parser.add_argument(
+ "--decoding-method",
+ type=str,
+ default="greedy_search",
+ help="""Possible values are:
+ - greedy_search
+ - beam_search
+ - modified_beam_search
+ - fast_beam_search
+ - fast_beam_search_nbest
+ - fast_beam_search_nbest_oracle
+ - fast_beam_search_nbest_LG
+ If you use fast_beam_search_nbest_LG, you have to specify
+ `--lang-dir`, which should contain `LG.pt`.
+ """,
+ )
+
+ parser.add_argument(
+ "--beam-size",
+ type=int,
+ default=4,
+ help="""An integer indicating how many candidates we will keep for each
+ frame. Used only when --decoding-method is beam_search or
+ modified_beam_search.""",
+ )
+
+ parser.add_argument(
+ "--beam",
+ type=float,
+ default=20.0,
+ help="""A floating point value to calculate the cutoff score during beam
+ search (i.e., `cutoff = max-score - beam`), which is the same as the
+ `beam` in Kaldi.
+ Used only when --decoding-method is fast_beam_search,
+ fast_beam_search_nbest, fast_beam_search_nbest_LG,
+ and fast_beam_search_nbest_oracle
+ """,
+ )
+
+ parser.add_argument(
+ "--ngram-lm-scale",
+ type=float,
+ default=0.01,
+ help="""
+ Used only when --decoding_method is fast_beam_search_nbest_LG.
+ It specifies the scale for n-gram LM scores.
+ """,
+ )
+
+ parser.add_argument(
+ "--max-contexts",
+ type=int,
+ default=8,
+ help="""Used only when --decoding-method is
+ fast_beam_search, fast_beam_search_nbest, fast_beam_search_nbest_LG,
+ and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--max-states",
+ type=int,
+ default=64,
+ help="""Used only when --decoding-method is
+ fast_beam_search, fast_beam_search_nbest, fast_beam_search_nbest_LG,
+ and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram; 2 means tri-gram",
+ )
+ parser.add_argument(
+ "--max-sym-per-frame",
+ type=int,
+ default=1,
+ help="""Maximum number of symbols per frame.
+ Used only when --decoding_method is greedy_search""",
+ )
+
+ parser.add_argument(
+ "--num-paths",
+ type=int,
+ default=200,
+ help="""Number of paths for nbest decoding.
+ Used only when the decoding method is fast_beam_search_nbest,
+ fast_beam_search_nbest_LG, and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--nbest-scale",
+ type=float,
+ default=0.5,
+ help="""Scale applied to lattice scores when computing nbest paths.
+ Used only when the decoding method is fast_beam_search_nbest,
+ fast_beam_search_nbest_LG, and fast_beam_search_nbest_oracle""",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def decode_one_batch(
+ params: AttributeDict,
+ model: nn.Module,
+ sp: spm.SentencePieceProcessor,
+ batch: dict,
+ word_table: Optional[k2.SymbolTable] = None,
+ decoding_graph: Optional[k2.Fsa] = None,
+) -> Dict[str, List[List[str]]]:
+ """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 greedy_search is used, it would be "greedy_search"
+ If beam search with a beam size of 7 is used, it would be
+ "beam_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`.
+ model:
+ The neural model.
+ sp:
+ The BPE model.
+ batch:
+ It is the return value from iterating
+ `lhotse.dataset.K2SpeechRecognitionDataset`. See its documentation
+ for the format of the `batch`.
+ word_table:
+ The word symbol table.
+ decoding_graph:
+ The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+ only when --decoding_method is fast_beam_search, fast_beam_search_nbest,
+ fast_beam_search_nbest_oracle, and fast_beam_search_nbest_LG.
+ Returns:
+ Return the decoding result. See above description for the format of
+ the returned dict.
+ """
+ device = next(model.parameters()).device
+ feature = batch["inputs"]
+ assert feature.ndim == 3
+
+ feature = feature.to(device)
+ # at entry, feature is (N, T, C)
+
+ supervisions = batch["supervisions"]
+ feature_lens = supervisions["num_frames"].to(device)
+
+ feature_lens += 30
+ feature = torch.nn.functional.pad(
+ feature,
+ pad=(0, 0, 0, 30),
+ value=LOG_EPS,
+ )
+ encoder_out, encoder_out_lens = model.encoder(x=feature, x_lens=feature_lens)
+
+ hyps = []
+
+ if params.decoding_method == "fast_beam_search":
+ hyp_tokens = fast_beam_search_one_best(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.decoding_method == "fast_beam_search_nbest_LG":
+ hyp_tokens = fast_beam_search_nbest_LG(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ num_paths=params.num_paths,
+ nbest_scale=params.nbest_scale,
+ )
+ for hyp in hyp_tokens:
+ hyps.append([word_table[i] for i in hyp])
+ elif params.decoding_method == "fast_beam_search_nbest":
+ hyp_tokens = fast_beam_search_nbest(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ num_paths=params.num_paths,
+ nbest_scale=params.nbest_scale,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.decoding_method == "fast_beam_search_nbest_oracle":
+ hyp_tokens = fast_beam_search_nbest_oracle(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ num_paths=params.num_paths,
+ ref_texts=sp.encode(supervisions["text"]),
+ nbest_scale=params.nbest_scale,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.decoding_method == "greedy_search" and params.max_sym_per_frame == 1:
+ hyp_tokens = greedy_search_batch(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.decoding_method == "modified_beam_search":
+ hyp_tokens = modified_beam_search(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam_size,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ else:
+ batch_size = encoder_out.size(0)
+
+ for i in range(batch_size):
+ # fmt: off
+ encoder_out_i = encoder_out[i:i+1, :encoder_out_lens[i]]
+ # fmt: on
+ if params.decoding_method == "greedy_search":
+ hyp = greedy_search(
+ model=model,
+ encoder_out=encoder_out_i,
+ max_sym_per_frame=params.max_sym_per_frame,
+ )
+ elif params.decoding_method == "beam_search":
+ hyp = beam_search(
+ model=model,
+ encoder_out=encoder_out_i,
+ beam=params.beam_size,
+ )
+ else:
+ raise ValueError(
+ f"Unsupported decoding method: {params.decoding_method}"
+ )
+ hyps.append(sp.decode(hyp).split())
+
+ if params.decoding_method == "greedy_search":
+ return {"greedy_search": hyps}
+ elif "fast_beam_search" in params.decoding_method:
+ key = f"beam_{params.beam}_"
+ key += f"max_contexts_{params.max_contexts}_"
+ key += f"max_states_{params.max_states}"
+ if "nbest" in params.decoding_method:
+ key += f"_num_paths_{params.num_paths}_"
+ key += f"nbest_scale_{params.nbest_scale}"
+ if "LG" in params.decoding_method:
+ key += f"_ngram_lm_scale_{params.ngram_lm_scale}"
+
+ return {key: hyps}
+ else:
+ return {f"beam_size_{params.beam_size}": hyps}
+
+
+def decode_dataset(
+ dl: torch.utils.data.DataLoader,
+ params: AttributeDict,
+ model: nn.Module,
+ sp: spm.SentencePieceProcessor,
+ word_table: Optional[k2.SymbolTable] = None,
+ decoding_graph: Optional[k2.Fsa] = None,
+) -> Dict[str, List[Tuple[str, List[str], List[str]]]]:
+ """Decode dataset.
+
+ Args:
+ dl:
+ PyTorch's dataloader containing the dataset to decode.
+ params:
+ It is returned by :func:`get_params`.
+ model:
+ The neural model.
+ sp:
+ The BPE model.
+ word_table:
+ The word symbol table.
+ decoding_graph:
+ The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+ only when --decoding_method is fast_beam_search, fast_beam_search_nbest,
+ fast_beam_search_nbest_oracle, and fast_beam_search_nbest_LG.
+ Returns:
+ Return a dict, whose key may be "greedy_search" if greedy search
+ is used, or it may be "beam_7" if beam size of 7 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.
+ """
+ num_cuts = 0
+
+ try:
+ num_batches = len(dl)
+ except TypeError:
+ num_batches = "?"
+
+ if params.decoding_method == "greedy_search":
+ log_interval = 50
+ else:
+ log_interval = 20
+
+ results = defaultdict(list)
+ for batch_idx, batch in enumerate(dl):
+ texts = batch["supervisions"]["text"]
+ cut_ids = [cut.id for cut in batch["supervisions"]["cut"]]
+
+ hyps_dict = decode_one_batch(
+ params=params,
+ model=model,
+ sp=sp,
+ decoding_graph=decoding_graph,
+ word_table=word_table,
+ batch=batch,
+ )
+
+ for name, hyps in hyps_dict.items():
+ this_batch = []
+ assert len(hyps) == len(texts)
+ for cut_id, hyp_words, ref_text in zip(cut_ids, hyps, texts):
+ ref_words = ref_text.split()
+ this_batch.append((cut_id, ref_words, hyp_words))
+
+ results[name].extend(this_batch)
+
+ num_cuts += len(texts)
+
+ if batch_idx % log_interval == 0:
+ batch_str = f"{batch_idx}/{num_batches}"
+
+ logging.info(f"batch {batch_str}, cuts processed until now is {num_cuts}")
+ return results
+
+
+def save_results(
+ params: AttributeDict,
+ test_set_name: str,
+ results_dict: Dict[str, List[Tuple[str, List[str], List[str]]]],
+):
+ test_set_wers = dict()
+ for key, results in results_dict.items():
+ recog_path = (
+ params.res_dir / f"recogs-{test_set_name}-{key}-{params.suffix}.txt"
+ )
+ results = sorted(results)
+ 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.res_dir / f"errs-{test_set_name}-{key}-{params.suffix}.txt"
+ )
+ with open(errs_filename, "w") as f:
+ wer = write_error_stats(
+ f, f"{test_set_name}-{key}", results, enable_log=True
+ )
+ 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.res_dir / f"wer-summary-{test_set_name}-{key}-{params.suffix}.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()
+ args.exp_dir = Path(args.exp_dir)
+
+ params = get_params()
+ params.update(vars(args))
+
+ assert params.decoding_method in (
+ "greedy_search",
+ "beam_search",
+ "fast_beam_search",
+ "fast_beam_search_nbest",
+ "fast_beam_search_nbest_LG",
+ "fast_beam_search_nbest_oracle",
+ "modified_beam_search",
+ )
+ params.res_dir = params.exp_dir / params.decoding_method
+
+ if params.iter > 0:
+ params.suffix = f"iter-{params.iter}-avg-{params.avg}"
+ else:
+ params.suffix = f"epoch-{params.epoch}-avg-{params.avg}"
+
+ params.suffix += f"-streaming-chunk-size-{params.decode_chunk_len}"
+
+ if "fast_beam_search" in params.decoding_method:
+ params.suffix += f"-beam-{params.beam}"
+ params.suffix += f"-max-contexts-{params.max_contexts}"
+ params.suffix += f"-max-states-{params.max_states}"
+ if "nbest" in params.decoding_method:
+ params.suffix += f"-nbest-scale-{params.nbest_scale}"
+ params.suffix += f"-num-paths-{params.num_paths}"
+ if "LG" in params.decoding_method:
+ params.suffix += f"-ngram-lm-scale-{params.ngram_lm_scale}"
+ elif "beam_search" in params.decoding_method:
+ params.suffix += f"-{params.decoding_method}-beam-size-{params.beam_size}"
+ else:
+ params.suffix += f"-context-{params.context_size}"
+ params.suffix += f"-max-sym-per-frame-{params.max_sym_per_frame}"
+
+ if params.use_averaged_model:
+ params.suffix += "-use-averaged-model"
+
+ setup_logger(f"{params.res_dir}/log-decode-{params.suffix}")
+ logging.info("Decoding started")
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"Device: {device}")
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(params.bpe_model)
+
+ # and are defined in local/train_bpe_model.py
+ params.blank_id = sp.piece_to_id("")
+ params.unk_id = sp.piece_to_id("")
+ params.vocab_size = sp.get_piece_size()
+
+ logging.info(params)
+
+ logging.info("About to create model")
+ model = get_transducer_model(params)
+ assert model.encoder.decode_chunk_size == params.decode_chunk_len // 2, (
+ model.encoder.decode_chunk_size,
+ params.decode_chunk_len,
+ )
+
+ if not params.use_averaged_model:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ elif 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 i >= 1:
+ filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg + 1
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg + 1:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ filename_start = filenames[-1]
+ filename_end = filenames[0]
+ logging.info(
+ "Calculating the averaged model over iteration checkpoints"
+ f" from {filename_start} (excluded) to {filename_end}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+ else:
+ assert params.avg > 0, params.avg
+ start = params.epoch - params.avg
+ assert start >= 1, start
+ filename_start = f"{params.exp_dir}/epoch-{start}.pt"
+ filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
+ logging.info(
+ f"Calculating the averaged model over epoch range from "
+ f"{start} (excluded) to {params.epoch}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+
+ model.to(device)
+ model.eval()
+
+ if "fast_beam_search" in params.decoding_method:
+ if params.decoding_method == "fast_beam_search_nbest_LG":
+ lexicon = Lexicon(params.lang_dir)
+ word_table = lexicon.word_table
+ lg_filename = params.lang_dir / "LG.pt"
+ logging.info(f"Loading {lg_filename}")
+ decoding_graph = k2.Fsa.from_dict(
+ torch.load(lg_filename, map_location=device)
+ )
+ decoding_graph.scores *= params.ngram_lm_scale
+ else:
+ word_table = None
+ decoding_graph = k2.trivial_graph(params.vocab_size - 1, device=device)
+ else:
+ decoding_graph = None
+ word_table = None
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ # we need cut ids to display recognition results.
+ args.return_cuts = True
+ librispeech = LibriSpeechAsrDataModule(args)
+
+ test_clean_cuts = librispeech.test_clean_cuts()
+ test_other_cuts = librispeech.test_other_cuts()
+
+ test_clean_dl = librispeech.test_dataloaders(test_clean_cuts)
+ test_other_dl = librispeech.test_dataloaders(test_other_cuts)
+
+ test_sets = ["test-clean", "test-other"]
+ test_dl = [test_clean_dl, test_other_dl]
+
+ for test_set, test_dl in zip(test_sets, test_dl):
+ results_dict = decode_dataset(
+ dl=test_dl,
+ params=params,
+ model=model,
+ sp=sp,
+ word_table=word_table,
+ decoding_graph=decoding_graph,
+ )
+
+ 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/pruned_transducer_stateless7_streaming/decode_stream.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/decode_stream.py
new file mode 100644
index 000000000..0d7e86fcf
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/decode_stream.py
@@ -0,0 +1,151 @@
+# Copyright 2022 Xiaomi Corp. (authors: Wei Kang,
+# Zengwei Yao)
+#
+# 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.
+
+import math
+from typing import List, Optional, Tuple
+
+import k2
+import torch
+from beam_search import Hypothesis, HypothesisList
+
+from icefall.utils import AttributeDict
+
+
+class DecodeStream(object):
+ def __init__(
+ self,
+ params: AttributeDict,
+ cut_id: str,
+ initial_states: List[torch.Tensor],
+ decoding_graph: Optional[k2.Fsa] = None,
+ device: torch.device = torch.device("cpu"),
+ ) -> None:
+ """
+ Args:
+ initial_states:
+ Initial decode states of the model, e.g. the return value of
+ `get_init_state` in conformer.py
+ decoding_graph:
+ Decoding graph used for decoding, may be a TrivialGraph or a HLG.
+ Used only when decoding_method is fast_beam_search.
+ device:
+ The device to run this stream.
+ """
+ if params.decoding_method == "fast_beam_search":
+ assert decoding_graph is not None
+ assert device == decoding_graph.device
+
+ self.params = params
+ self.cut_id = cut_id
+ self.LOG_EPS = math.log(1e-10)
+
+ self.states = initial_states
+
+ # It contains a 2-D tensors representing the feature frames.
+ self.features: torch.Tensor = None
+
+ self.num_frames: int = 0
+ # how many frames have been processed. (before subsampling).
+ # we only modify this value in `func:get_feature_frames`.
+ self.num_processed_frames: int = 0
+
+ self._done: bool = False
+
+ # The transcript of current utterance.
+ self.ground_truth: str = ""
+
+ # The decoding result (partial or final) of current utterance.
+ self.hyp: List = []
+
+ # how many frames have been processed, after subsampling (i.e. a
+ # cumulative sum of the second return value of
+ # encoder.streaming_forward
+ self.done_frames: int = 0
+
+ # It has two steps of feature subsampling in zipformer: out_lens=((x_lens-7)//2+1)//2
+ # 1) feature embedding: out_lens=(x_lens-7)//2
+ # 2) output subsampling: out_lens=(out_lens+1)//2
+ self.pad_length = 7
+
+ if params.decoding_method == "greedy_search":
+ self.hyp = [params.blank_id] * params.context_size
+ elif params.decoding_method == "modified_beam_search":
+ self.hyps = HypothesisList()
+ self.hyps.add(
+ Hypothesis(
+ ys=[params.blank_id] * params.context_size,
+ log_prob=torch.zeros(1, dtype=torch.float32, device=device),
+ )
+ )
+ elif params.decoding_method == "fast_beam_search":
+ # The rnnt_decoding_stream for fast_beam_search.
+ self.rnnt_decoding_stream: k2.RnntDecodingStream = k2.RnntDecodingStream(
+ decoding_graph
+ )
+ else:
+ raise ValueError(f"Unsupported decoding method: {params.decoding_method}")
+
+ @property
+ def done(self) -> bool:
+ """Return True if all the features are processed."""
+ return self._done
+
+ @property
+ def id(self) -> str:
+ return self.cut_id
+
+ def set_features(
+ self,
+ features: torch.Tensor,
+ tail_pad_len: int = 0,
+ ) -> None:
+ """Set features tensor of current utterance."""
+ assert features.dim() == 2, features.dim()
+ self.features = torch.nn.functional.pad(
+ features,
+ (0, 0, 0, self.pad_length + tail_pad_len),
+ mode="constant",
+ value=self.LOG_EPS,
+ )
+ self.num_frames = self.features.size(0)
+
+ def get_feature_frames(self, chunk_size: int) -> Tuple[torch.Tensor, int]:
+ """Consume chunk_size frames of features"""
+ chunk_length = chunk_size + self.pad_length
+
+ ret_length = min(self.num_frames - self.num_processed_frames, chunk_length)
+
+ ret_features = self.features[
+ self.num_processed_frames : self.num_processed_frames + ret_length # noqa
+ ]
+
+ self.num_processed_frames += chunk_size
+ if self.num_processed_frames >= self.num_frames:
+ self._done = True
+
+ return ret_features, ret_length
+
+ def decoding_result(self) -> List[int]:
+ """Obtain current decoding result."""
+ if self.params.decoding_method == "greedy_search":
+ return self.hyp[self.params.context_size :] # noqa
+ elif self.params.decoding_method == "modified_beam_search":
+ best_hyp = self.hyps.get_most_probable(length_norm=True)
+ return best_hyp.ys[self.params.context_size :] # noqa
+ else:
+ assert self.params.decoding_method == "fast_beam_search"
+ return self.hyp
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/decoder.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/decoder.py
new file mode 120000
index 000000000..33944d0d2
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/decoder.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/decoder.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/encoder_interface.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/encoder_interface.py
new file mode 120000
index 000000000..b9aa0ae08
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/encoder_interface.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/encoder_interface.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/export.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/export.py
new file mode 100755
index 000000000..5c06cc052
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/export.py
@@ -0,0 +1,320 @@
+#!/usr/bin/env python3
+#
+# Copyright 2021 Xiaomi Corporation (Author: Fangjun Kuang)
+#
+# 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.
+
+# This script converts several saved checkpoints
+# to a single one using model averaging.
+"""
+
+Usage:
+
+(1) Export to torchscript model using torch.jit.script()
+
+./pruned_transducer_stateless7_streaming/export.py \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 9 \
+ --jit 1
+
+It will generate a file `cpu_jit.pt` in the given `exp_dir`. You can later
+load it by `torch.jit.load("cpu_jit.pt")`.
+
+Note `cpu` in the name `cpu_jit.pt` means the parameters when loaded into Python
+are on CPU. You can use `to("cuda")` to move them to a CUDA device.
+
+Check
+https://github.com/k2-fsa/sherpa
+for how to use the exported models outside of icefall.
+
+(2) Export `model.state_dict()`
+
+./pruned_transducer_stateless7_streaming/export.py \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 20 \
+ --avg 10
+
+It will generate a file `pretrained.pt` in the given `exp_dir`. You can later
+load it by `icefall.checkpoint.load_checkpoint()`.
+
+To use the generated file with `pruned_transducer_stateless7_streaming/decode.py`,
+you can do:
+
+ cd /path/to/exp_dir
+ ln -s pretrained.pt epoch-9999.pt
+
+ cd /path/to/egs/librispeech/ASR
+ ./pruned_transducer_stateless7_streaming/decode.py \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --epoch 9999 \
+ --avg 1 \
+ --max-duration 600 \
+ --decoding-method greedy_search \
+ --bpe-model data/lang_bpe_500/bpe.model
+
+Check ./pretrained.py for its usage.
+
+Note: If you don't want to train a model from scratch, we have
+provided one for you. You can get it at
+
+https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless7-2022-11-11
+
+with the following commands:
+
+ sudo apt-get install git-lfs
+ git lfs install
+ git clone https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless7-2022-11-11
+ # You will find the pre-trained model in icefall-asr-librispeech-pruned-transducer-stateless7-2022-11-11/exp
+"""
+
+import argparse
+import logging
+from pathlib import Path
+
+import sentencepiece as spm
+import torch
+import torch.nn as nn
+from scaling_converter import convert_scaled_to_non_scaled
+from train import add_model_arguments, get_params, get_transducer_model
+
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from icefall.utils import str2bool
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=30,
+ help="""It specifies the checkpoint to use for decoding.
+ Note: Epoch counts from 1.
+ You can specify --avg to use more checkpoints for model averaging.""",
+ )
+
+ parser.add_argument(
+ "--iter",
+ type=int,
+ default=0,
+ help="""If positive, --epoch is ignored and it
+ will use the checkpoint exp_dir/checkpoint-iter.pt.
+ You can specify --avg to use more checkpoints for model averaging.
+ """,
+ )
+
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=9,
+ help="Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch' and '--iter'",
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=True,
+ help="Whether to load averaged model. Currently it only supports "
+ "using --epoch. If True, it would decode with the averaged model "
+ "over the epoch range from `epoch-avg` (excluded) to `epoch`."
+ "Actually only the models with epoch number of `epoch-avg` and "
+ "`epoch` are loaded for averaging. ",
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="pruned_transducer_stateless7_streaming/exp",
+ help="""It specifies the directory where all training related
+ files, e.g., checkpoints, log, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ default="data/lang_bpe_500/bpe.model",
+ help="Path to the BPE model",
+ )
+
+ parser.add_argument(
+ "--jit",
+ type=str2bool,
+ default=False,
+ help="""True to save a model after applying torch.jit.script.
+ It will generate a file named cpu_jit.pt
+
+ Check ./jit_pretrained.py for how to use it.
+ """,
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram; 2 means tri-gram",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+@torch.no_grad()
+def main():
+ args = get_parser().parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ params = get_params()
+ params.update(vars(args))
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"device: {device}")
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(params.bpe_model)
+
+ # is defined in local/train_bpe_model.py
+ params.blank_id = sp.piece_to_id("")
+ params.vocab_size = sp.get_piece_size()
+
+ logging.info(params)
+
+ logging.info("About to create model")
+ model = get_transducer_model(params)
+
+ model.to(device)
+
+ if not params.use_averaged_model:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ elif 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 i >= 1:
+ filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg + 1
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg + 1:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ filename_start = filenames[-1]
+ filename_end = filenames[0]
+ logging.info(
+ "Calculating the averaged model over iteration checkpoints"
+ f" from {filename_start} (excluded) to {filename_end}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+ else:
+ assert params.avg > 0, params.avg
+ start = params.epoch - params.avg
+ assert start >= 1, start
+ filename_start = f"{params.exp_dir}/epoch-{start}.pt"
+ filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
+ logging.info(
+ f"Calculating the averaged model over epoch range from "
+ f"{start} (excluded) to {params.epoch}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+
+ model.to("cpu")
+ model.eval()
+
+ if params.jit is True:
+ convert_scaled_to_non_scaled(model, inplace=True)
+ # We won't use the forward() method of the model in C++, so just ignore
+ # it here.
+ # Otherwise, one of its arguments is a ragged tensor and is not
+ # torch scriptabe.
+ model.__class__.forward = torch.jit.ignore(model.__class__.forward)
+ logging.info("Using torch.jit.script")
+ model = torch.jit.script(model)
+ filename = params.exp_dir / "cpu_jit.pt"
+ model.save(str(filename))
+ logging.info(f"Saved to {filename}")
+ else:
+ logging.info("Not using torchscript. Export model.state_dict()")
+ # Save it using a format so that it can be loaded
+ # by :func:`load_checkpoint`
+ filename = params.exp_dir / "pretrained.pt"
+ torch.save({"model": model.state_dict()}, str(filename))
+ logging.info(f"Saved to {filename}")
+
+
+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/pruned_transducer_stateless7_streaming/jit_pretrained.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/jit_pretrained.py
new file mode 100755
index 000000000..4fd5e1820
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/jit_pretrained.py
@@ -0,0 +1,278 @@
+#!/usr/bin/env python3
+# Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang)
+#
+# 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.
+"""
+This script loads torchscript models, exported by `torch.jit.script()`
+and uses them to decode waves.
+You can use the following command to get the exported models:
+
+./pruned_transducer_stateless7_streaming/export.py \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 20 \
+ --avg 10 \
+ --jit 1
+
+Usage of this script:
+
+./pruned_transducer_stateless7_streaming/jit_pretrained.py \
+ --nn-model-filename ./pruned_transducer_stateless7_streaming/exp/cpu_jit.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+"""
+
+import argparse
+import logging
+import math
+from typing import List
+
+import kaldifeat
+import sentencepiece as spm
+import torch
+import torchaudio
+from torch.nn.utils.rnn import pad_sequence
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--nn-model-filename",
+ type=str,
+ required=True,
+ help="Path to the torchscript model cpu_jit.pt",
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ help="""Path to bpe.model.""",
+ )
+
+ parser.add_argument(
+ "sound_files",
+ type=str,
+ nargs="+",
+ help="The input sound file(s) to transcribe. "
+ "Supported formats are those supported by torchaudio.load(). "
+ "For example, wav and flac are supported. "
+ "The sample rate has to be 16kHz.",
+ )
+
+ parser.add_argument(
+ "--decode-chunk-len",
+ type=int,
+ default=32,
+ help="The chunk size for decoding (in frames before subsampling)",
+ )
+
+ return parser
+
+
+def read_sound_files(
+ filenames: List[str], expected_sample_rate: float = 16000
+) -> List[torch.Tensor]:
+ """Read a list of sound files into a list 1-D float32 torch tensors.
+ Args:
+ filenames:
+ A list of sound filenames.
+ expected_sample_rate:
+ The expected sample rate of the sound files.
+ Returns:
+ Return a list of 1-D float32 torch tensors.
+ """
+ ans = []
+ for f in filenames:
+ wave, sample_rate = torchaudio.load(f)
+ assert (
+ sample_rate == expected_sample_rate
+ ), f"expected sample rate: {expected_sample_rate}. Given: {sample_rate}"
+ # We use only the first channel
+ ans.append(wave[0])
+ return ans
+
+
+def greedy_search(
+ model: torch.jit.ScriptModule,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+) -> List[List[int]]:
+ """Greedy search in batch mode. It hardcodes --max-sym-per-frame=1.
+ Args:
+ model:
+ The transducer model.
+ encoder_out:
+ A 3-D tensor of shape (N, T, C)
+ encoder_out_lens:
+ A 1-D tensor of shape (N,).
+ Returns:
+ Return the decoded results for each utterance.
+ """
+ assert encoder_out.ndim == 3
+ assert encoder_out.size(0) >= 1, encoder_out.size(0)
+
+ packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
+ input=encoder_out,
+ lengths=encoder_out_lens.cpu(),
+ batch_first=True,
+ enforce_sorted=False,
+ )
+
+ device = encoder_out.device
+ blank_id = 0 # hard-code to 0
+
+ batch_size_list = packed_encoder_out.batch_sizes.tolist()
+ N = encoder_out.size(0)
+
+ assert torch.all(encoder_out_lens > 0), encoder_out_lens
+ assert N == batch_size_list[0], (N, batch_size_list)
+
+ context_size = model.decoder.context_size
+ hyps = [[blank_id] * context_size for _ in range(N)]
+
+ decoder_input = torch.tensor(
+ hyps,
+ device=device,
+ dtype=torch.int64,
+ ) # (N, context_size)
+
+ decoder_out = model.decoder(
+ decoder_input,
+ need_pad=torch.tensor([False]),
+ ).squeeze(1)
+
+ offset = 0
+ for batch_size in batch_size_list:
+ start = offset
+ end = offset + batch_size
+ current_encoder_out = packed_encoder_out.data[start:end]
+ current_encoder_out = current_encoder_out
+ # current_encoder_out's shape: (batch_size, encoder_out_dim)
+ offset = end
+
+ decoder_out = decoder_out[:batch_size]
+
+ logits = model.joiner(
+ current_encoder_out,
+ decoder_out,
+ )
+ # logits'shape (batch_size, vocab_size)
+
+ assert logits.ndim == 2, logits.shape
+ y = logits.argmax(dim=1).tolist()
+ emitted = False
+ for i, v in enumerate(y):
+ if v != blank_id:
+ hyps[i].append(v)
+ emitted = True
+ if emitted:
+ # update decoder output
+ decoder_input = [h[-context_size:] for h in hyps[:batch_size]]
+ decoder_input = torch.tensor(
+ decoder_input,
+ device=device,
+ dtype=torch.int64,
+ )
+ decoder_out = model.decoder(
+ decoder_input,
+ need_pad=torch.tensor([False]),
+ )
+ decoder_out = decoder_out.squeeze(1)
+
+ sorted_ans = [h[context_size:] for h in hyps]
+ ans = []
+ unsorted_indices = packed_encoder_out.unsorted_indices.tolist()
+ for i in range(N):
+ ans.append(sorted_ans[unsorted_indices[i]])
+
+ return ans
+
+
+@torch.no_grad()
+def main():
+ parser = get_parser()
+ args = parser.parse_args()
+ logging.info(vars(args))
+
+ device = torch.device("cpu")
+
+ logging.info(f"device: {device}")
+
+ model = torch.jit.load(args.nn_model_filename)
+ model.encoder.decode_chunk_size = args.decode_chunk_len // 2
+
+ model.eval()
+
+ model.to(device)
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(args.bpe_model)
+
+ logging.info("Constructing Fbank computer")
+ opts = kaldifeat.FbankOptions()
+ opts.device = device
+ opts.frame_opts.dither = 0
+ opts.frame_opts.snip_edges = False
+ opts.frame_opts.samp_freq = 16000
+ opts.mel_opts.num_bins = 80
+
+ fbank = kaldifeat.Fbank(opts)
+
+ logging.info(f"Reading sound files: {args.sound_files}")
+ waves = read_sound_files(
+ filenames=args.sound_files,
+ )
+ waves = [w.to(device) for w in waves]
+
+ logging.info("Decoding started")
+ features = fbank(waves)
+ feature_lengths = [f.size(0) for f in features]
+
+ features = pad_sequence(
+ features,
+ batch_first=True,
+ padding_value=math.log(1e-10),
+ )
+
+ feature_lengths = torch.tensor(feature_lengths, device=device)
+
+ encoder_out, encoder_out_lens = model.encoder(
+ x=features,
+ x_lens=feature_lengths,
+ )
+
+ hyps = greedy_search(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ )
+ s = "\n"
+ for filename, hyp in zip(args.sound_files, hyps):
+ words = sp.decode(hyp)
+ s += f"{filename}:\n{words}\n\n"
+ logging.info(s)
+
+ logging.info("Decoding Done")
+
+
+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/pruned_transducer_stateless7_streaming/jit_trace_export.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/jit_trace_export.py
new file mode 100755
index 000000000..a164f3f69
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/jit_trace_export.py
@@ -0,0 +1,313 @@
+#!/usr/bin/env python3
+
+"""
+Usage:
+./pruned_transducer_stateless7_streaming/jit_trace_export.py \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 10 \
+ --use-averaged-model=True \
+ --decode-chunk-len 32
+"""
+
+import argparse
+import logging
+from pathlib import Path
+
+import sentencepiece as spm
+import torch
+from scaling_converter import convert_scaled_to_non_scaled
+from train import add_model_arguments, get_params, get_transducer_model
+
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from icefall.utils import AttributeDict, str2bool
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=28,
+ help="""It specifies the checkpoint to use for averaging.
+ Note: Epoch counts from 0.
+ You can specify --avg to use more checkpoints for model averaging.""",
+ )
+
+ parser.add_argument(
+ "--iter",
+ type=int,
+ default=0,
+ help="""If positive, --epoch is ignored and it
+ will use the checkpoint exp_dir/checkpoint-iter.pt.
+ You can specify --avg to use more checkpoints for model averaging.
+ """,
+ )
+
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=15,
+ help="Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch' and '--iter'",
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="pruned_transducer_stateless2/exp",
+ help="""It specifies the directory where all training related
+ files, e.g., checkpoints, log, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ default="data/lang_bpe_500/bpe.model",
+ help="Path to the BPE model",
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram; 2 means tri-gram",
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=True,
+ help="Whether to load averaged model. Currently it only supports "
+ "using --epoch. If True, it would decode with the averaged model "
+ "over the epoch range from `epoch-avg` (excluded) to `epoch`."
+ "Actually only the models with epoch number of `epoch-avg` and "
+ "`epoch` are loaded for averaging. ",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def export_encoder_model_jit_trace(
+ encoder_model: torch.nn.Module,
+ encoder_filename: str,
+ params: AttributeDict,
+) -> None:
+ """Export the given encoder model with torch.jit.trace()
+
+ Note: The warmup argument is fixed to 1.
+
+ Args:
+ encoder_model:
+ The input encoder model
+ encoder_filename:
+ The filename to save the exported model.
+ """
+ decode_chunk_len = params.decode_chunk_len # before subsampling
+ pad_length = 7
+ s = f"decode_chunk_len: {decode_chunk_len}"
+ logging.info(s)
+ assert encoder_model.decode_chunk_size == decode_chunk_len // 2, (
+ encoder_model.decode_chunk_size,
+ decode_chunk_len,
+ )
+
+ T = decode_chunk_len + pad_length
+
+ x = torch.zeros(1, T, 80, dtype=torch.float32)
+ x_lens = torch.full((1,), T, dtype=torch.int32)
+ states = encoder_model.get_init_state(device=x.device)
+
+ encoder_model.__class__.forward = encoder_model.__class__.streaming_forward
+ traced_model = torch.jit.trace(encoder_model, (x, x_lens, states))
+ traced_model.save(encoder_filename)
+ logging.info(f"Saved to {encoder_filename}")
+
+
+def export_decoder_model_jit_trace(
+ decoder_model: torch.nn.Module,
+ decoder_filename: str,
+) -> None:
+ """Export the given decoder model with torch.jit.trace()
+
+ Note: The argument need_pad is fixed to False.
+
+ Args:
+ decoder_model:
+ The input decoder model
+ decoder_filename:
+ The filename to save the exported model.
+ """
+ y = torch.zeros(10, decoder_model.context_size, dtype=torch.int64)
+ need_pad = torch.tensor([False])
+
+ traced_model = torch.jit.trace(decoder_model, (y, need_pad))
+ traced_model.save(decoder_filename)
+ logging.info(f"Saved to {decoder_filename}")
+
+
+def export_joiner_model_jit_trace(
+ joiner_model: torch.nn.Module,
+ joiner_filename: str,
+) -> None:
+ """Export the given joiner model with torch.jit.trace()
+
+ Note: The argument project_input is fixed to True. A user should not
+ project the encoder_out/decoder_out by himself/herself. The exported joiner
+ will do that for the user.
+
+ Args:
+ joiner_model:
+ The input joiner model
+ joiner_filename:
+ The filename to save the exported model.
+
+ """
+ encoder_out_dim = joiner_model.encoder_proj.weight.shape[1]
+ decoder_out_dim = joiner_model.decoder_proj.weight.shape[1]
+ encoder_out = torch.rand(1, encoder_out_dim, dtype=torch.float32)
+ decoder_out = torch.rand(1, decoder_out_dim, dtype=torch.float32)
+
+ traced_model = torch.jit.trace(joiner_model, (encoder_out, decoder_out))
+ traced_model.save(joiner_filename)
+ logging.info(f"Saved to {joiner_filename}")
+
+
+@torch.no_grad()
+def main():
+ args = get_parser().parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ params = get_params()
+ params.update(vars(args))
+
+ device = torch.device("cpu")
+
+ logging.info(f"device: {device}")
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(params.bpe_model)
+
+ # is defined in local/train_bpe_model.py
+ params.blank_id = sp.piece_to_id("")
+ params.vocab_size = sp.get_piece_size()
+
+ logging.info(params)
+
+ logging.info("About to create model")
+ model = get_transducer_model(params)
+
+ if not params.use_averaged_model:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ elif 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 i >= 1:
+ filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg + 1
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg + 1:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ filename_start = filenames[-1]
+ filename_end = filenames[0]
+ logging.info(
+ "Calculating the averaged model over iteration checkpoints"
+ f" from {filename_start} (excluded) to {filename_end}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+ else:
+ assert params.avg > 0, params.avg
+ start = params.epoch - params.avg
+ assert start >= 1, start
+ filename_start = f"{params.exp_dir}/epoch-{start}.pt"
+ filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
+ logging.info(
+ f"Calculating the averaged model over epoch range from "
+ f"{start} (excluded) to {params.epoch}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+
+ model.to("cpu")
+ model.eval()
+
+ convert_scaled_to_non_scaled(model, inplace=True)
+ logging.info("Using torch.jit.trace()")
+
+ logging.info("Exporting encoder")
+ encoder_filename = params.exp_dir / "encoder_jit_trace.pt"
+ export_encoder_model_jit_trace(model.encoder, encoder_filename, params)
+
+ logging.info("Exporting decoder")
+ decoder_filename = params.exp_dir / "decoder_jit_trace.pt"
+ export_decoder_model_jit_trace(model.decoder, decoder_filename)
+
+ logging.info("Exporting joiner")
+ joiner_filename = params.exp_dir / "joiner_jit_trace.pt"
+ export_joiner_model_jit_trace(model.joiner, joiner_filename)
+
+
+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/pruned_transducer_stateless7_streaming/jit_trace_pretrained.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/jit_trace_pretrained.py
new file mode 100755
index 000000000..f2ac1914d
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/jit_trace_pretrained.py
@@ -0,0 +1,295 @@
+#!/usr/bin/env python3
+# flake8: noqa
+# Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang, Zengwei Yao)
+#
+# 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.
+"""
+This script loads torchscript models exported by `torch.jit.trace()`
+and uses them to decode waves.
+You can use the following command to get the exported models:
+
+./pruned_transducer_stateless7_streaming/jit_trace_export.py \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 10 \
+ --use-averaged-model=True \
+ --decode-chunk-len 32
+
+Usage of this script:
+
+./pruned_transducer_stateless7_streaming/jit_trace_pretrained.py \
+ --encoder-model-filename ./pruned_transducer_stateless7_streaming/exp/encoder_jit_trace.pt \
+ --decoder-model-filename ./pruned_transducer_stateless7_streaming/exp/decoder_jit_trace.pt \
+ --joiner-model-filename ./pruned_transducer_stateless7_streaming/exp/joiner_jit_trace.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --decode-chunk-len 32 \
+ /path/to/foo.wav \
+"""
+
+import argparse
+import logging
+import math
+from typing import List, Optional
+
+import kaldifeat
+import sentencepiece as spm
+import torch
+import torchaudio
+from kaldifeat import FbankOptions, OnlineFbank, OnlineFeature
+from torch.nn.utils.rnn import pad_sequence
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--encoder-model-filename",
+ type=str,
+ required=True,
+ help="Path to the encoder torchscript model. ",
+ )
+
+ parser.add_argument(
+ "--decoder-model-filename",
+ type=str,
+ required=True,
+ help="Path to the decoder torchscript model. ",
+ )
+
+ parser.add_argument(
+ "--joiner-model-filename",
+ type=str,
+ required=True,
+ help="Path to the joiner torchscript model. ",
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ help="""Path to bpe.model.""",
+ )
+
+ parser.add_argument(
+ "--sample-rate",
+ type=int,
+ default=16000,
+ help="The sample rate of the input sound file",
+ )
+
+ parser.add_argument(
+ "--decode-chunk-len",
+ type=int,
+ default=32,
+ help="The chunk size for decoding (in frames before subsampling)",
+ )
+
+ parser.add_argument(
+ "sound_file",
+ type=str,
+ help="The input sound file(s) to transcribe. "
+ "Supported formats are those supported by torchaudio.load(). "
+ "For example, wav and flac are supported. "
+ "The sample rate has to be 16kHz.",
+ )
+
+ return parser
+
+
+def read_sound_files(
+ filenames: List[str], expected_sample_rate: float
+) -> List[torch.Tensor]:
+ """Read a list of sound files into a list 1-D float32 torch tensors.
+ Args:
+ filenames:
+ A list of sound filenames.
+ expected_sample_rate:
+ The expected sample rate of the sound files.
+ Returns:
+ Return a list of 1-D float32 torch tensors.
+ """
+ ans = []
+ for f in filenames:
+ wave, sample_rate = torchaudio.load(f)
+ assert (
+ sample_rate == expected_sample_rate
+ ), f"expected sample rate: {expected_sample_rate}. Given: {sample_rate}"
+ # We use only the first channel
+ ans.append(wave[0])
+ return ans
+
+
+def greedy_search(
+ decoder: torch.jit.ScriptModule,
+ joiner: torch.jit.ScriptModule,
+ encoder_out: torch.Tensor,
+ decoder_out: Optional[torch.Tensor] = None,
+ hyp: Optional[List[int]] = None,
+):
+ assert encoder_out.ndim == 2
+ context_size = 2
+ blank_id = 0
+
+ if decoder_out is None:
+ assert hyp is None, hyp
+ hyp = [blank_id] * context_size
+ decoder_input = torch.tensor(hyp, dtype=torch.int32).unsqueeze(0)
+ # decoder_input.shape (1,, 1 context_size)
+ decoder_out = decoder(decoder_input, torch.tensor([False])).squeeze(1)
+ else:
+ assert decoder_out.ndim == 2
+ assert hyp is not None, hyp
+
+ T = encoder_out.size(0)
+ for i in range(T):
+ cur_encoder_out = encoder_out[i : i + 1]
+ joiner_out = joiner(cur_encoder_out, decoder_out).squeeze(0)
+ y = joiner_out.argmax(dim=0).item()
+
+ if y != blank_id:
+ hyp.append(y)
+ decoder_input = hyp[-context_size:]
+
+ decoder_input = torch.tensor(decoder_input, dtype=torch.int32).unsqueeze(0)
+ decoder_out = decoder(decoder_input, torch.tensor([False])).squeeze(1)
+
+ return hyp, decoder_out
+
+
+def create_streaming_feature_extractor(sample_rate) -> OnlineFeature:
+ """Create a CPU streaming feature extractor.
+
+ At present, we assume it returns a fbank feature extractor with
+ fixed options. In the future, we will support passing in the options
+ from outside.
+
+ Returns:
+ Return a CPU streaming feature extractor.
+ """
+ opts = FbankOptions()
+ opts.device = "cpu"
+ opts.frame_opts.dither = 0
+ opts.frame_opts.snip_edges = False
+ opts.frame_opts.samp_freq = sample_rate
+ opts.mel_opts.num_bins = 80
+ return OnlineFbank(opts)
+
+
+@torch.no_grad()
+def main():
+ parser = get_parser()
+ args = parser.parse_args()
+ logging.info(vars(args))
+
+ device = torch.device("cpu")
+
+ logging.info(f"device: {device}")
+
+ encoder = torch.jit.load(args.encoder_model_filename)
+ decoder = torch.jit.load(args.decoder_model_filename)
+ joiner = torch.jit.load(args.joiner_model_filename)
+
+ encoder.eval()
+ decoder.eval()
+ joiner.eval()
+
+ encoder.to(device)
+ decoder.to(device)
+ joiner.to(device)
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(args.bpe_model)
+
+ logging.info("Constructing Fbank computer")
+ online_fbank = create_streaming_feature_extractor(args.sample_rate)
+
+ logging.info(f"Reading sound files: {args.sound_file}")
+ wave_samples = read_sound_files(
+ filenames=[args.sound_file],
+ expected_sample_rate=args.sample_rate,
+ )[0]
+ logging.info(wave_samples.shape)
+
+ logging.info("Decoding started")
+ chunk_length = args.decode_chunk_len
+ assert encoder.decode_chunk_size == chunk_length // 2, (
+ encoder.decode_chunk_size,
+ chunk_length,
+ )
+
+ # we subsample features with ((x_len - 7) // 2 + 1) // 2
+ pad_length = 7
+ T = chunk_length + pad_length
+
+ logging.info(f"chunk_length: {chunk_length}")
+
+ states = encoder.get_init_state(device)
+
+ tail_padding = torch.zeros(int(0.3 * args.sample_rate), dtype=torch.float32)
+
+ wave_samples = torch.cat([wave_samples, tail_padding])
+
+ chunk = int(0.25 * args.sample_rate) # 0.2 second
+ num_processed_frames = 0
+
+ hyp = None
+ decoder_out = None
+
+ start = 0
+ while start < wave_samples.numel():
+ logging.info(f"{start}/{wave_samples.numel()}")
+ end = min(start + chunk, wave_samples.numel())
+ samples = wave_samples[start:end]
+ start += chunk
+ online_fbank.accept_waveform(
+ sampling_rate=args.sample_rate,
+ waveform=samples,
+ )
+ while online_fbank.num_frames_ready - num_processed_frames >= T:
+ frames = []
+ for i in range(T):
+ frames.append(online_fbank.get_frame(num_processed_frames + i))
+ frames = torch.cat(frames, dim=0).unsqueeze(0)
+ x_lens = torch.tensor([T], dtype=torch.int32)
+ encoder_out, out_lens, states = encoder(
+ x=frames,
+ x_lens=x_lens,
+ states=states,
+ )
+ num_processed_frames += chunk_length
+
+ hyp, decoder_out = greedy_search(
+ decoder, joiner, encoder_out.squeeze(0), decoder_out, hyp
+ )
+
+ context_size = 2
+ logging.info(args.sound_file)
+ logging.info(sp.decode(hyp[context_size:]))
+
+ logging.info("Decoding Done")
+
+
+torch.set_num_threads(4)
+torch.set_num_interop_threads(1)
+torch._C._jit_set_profiling_executor(False)
+torch._C._jit_set_profiling_mode(False)
+torch._C._set_graph_executor_optimize(False)
+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/pruned_transducer_stateless7_streaming/joiner.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/joiner.py
new file mode 120000
index 000000000..ecfb6dd8a
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/joiner.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/joiner.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/model.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/model.py
new file mode 120000
index 000000000..e17d4f734
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/model.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/model.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/optim.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/optim.py
new file mode 120000
index 000000000..81ac4a89a
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/optim.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/optim.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/pretrained.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/pretrained.py
new file mode 100755
index 000000000..fb77fdd42
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/pretrained.py
@@ -0,0 +1,355 @@
+#!/usr/bin/env python3
+# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang)
+#
+# 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.
+"""
+This script loads a checkpoint and uses it to decode waves.
+You can generate the checkpoint with the following command:
+
+./pruned_transducer_stateless7_streaming/export.py \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 20 \
+ --avg 10
+
+Usage of this script:
+
+(1) greedy search
+./pruned_transducer_stateless7_streaming/pretrained.py \
+ --checkpoint ./pruned_transducer_stateless7_streaming/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method greedy_search \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+(2) beam search
+./pruned_transducer_stateless7_streaming/pretrained.py \
+ --checkpoint ./pruned_transducer_stateless7_streaming/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method beam_search \
+ --beam-size 4 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+(3) modified beam search
+./pruned_transducer_stateless7_streaming/pretrained.py \
+ --checkpoint ./pruned_transducer_stateless7_streaming/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method modified_beam_search \
+ --beam-size 4 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+(4) fast beam search
+./pruned_transducer_stateless7_streaming/pretrained.py \
+ --checkpoint ./pruned_transducer_stateless7_streaming/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method fast_beam_search \
+ --beam-size 4 \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+You can also use `./pruned_transducer_stateless7_streaming/exp/epoch-xx.pt`.
+
+Note: ./pruned_transducer_stateless7_streaming/exp/pretrained.pt is generated by
+./pruned_transducer_stateless7_streaming/export.py
+"""
+
+
+import argparse
+import logging
+import math
+from typing import List
+
+import k2
+import kaldifeat
+import sentencepiece as spm
+import torch
+import torchaudio
+from beam_search import (
+ beam_search,
+ fast_beam_search_one_best,
+ greedy_search,
+ greedy_search_batch,
+ modified_beam_search,
+)
+from torch.nn.utils.rnn import pad_sequence
+from train import add_model_arguments, get_params, get_transducer_model
+
+from icefall.utils import str2bool
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--checkpoint",
+ type=str,
+ required=True,
+ help="Path to the checkpoint. "
+ "The checkpoint is assumed to be saved by "
+ "icefall.checkpoint.save_checkpoint().",
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ help="""Path to bpe.model.""",
+ )
+
+ parser.add_argument(
+ "--method",
+ type=str,
+ default="greedy_search",
+ help="""Possible values are:
+ - greedy_search
+ - beam_search
+ - modified_beam_search
+ - fast_beam_search
+ """,
+ )
+
+ parser.add_argument(
+ "sound_files",
+ type=str,
+ nargs="+",
+ help="The input sound file(s) to transcribe. "
+ "Supported formats are those supported by torchaudio.load(). "
+ "For example, wav and flac are supported. "
+ "The sample rate has to be 16kHz.",
+ )
+
+ parser.add_argument(
+ "--sample-rate",
+ type=int,
+ default=16000,
+ help="The sample rate of the input sound file",
+ )
+
+ parser.add_argument(
+ "--beam-size",
+ type=int,
+ default=4,
+ help="""An integer indicating how many candidates we will keep for each
+ frame. Used only when --method is beam_search or
+ modified_beam_search.""",
+ )
+
+ parser.add_argument(
+ "--beam",
+ type=float,
+ default=4,
+ help="""A floating point value to calculate the cutoff score during beam
+ search (i.e., `cutoff = max-score - beam`), which is the same as the
+ `beam` in Kaldi.
+ Used only when --method is fast_beam_search""",
+ )
+
+ parser.add_argument(
+ "--max-contexts",
+ type=int,
+ default=4,
+ help="""Used only when --method is fast_beam_search""",
+ )
+
+ parser.add_argument(
+ "--max-states",
+ type=int,
+ default=8,
+ help="""Used only when --method is fast_beam_search""",
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram; 2 means tri-gram",
+ )
+ parser.add_argument(
+ "--max-sym-per-frame",
+ type=int,
+ default=1,
+ help="""Maximum number of symbols per frame. Used only when
+ --method is greedy_search.
+ """,
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def read_sound_files(
+ filenames: List[str], expected_sample_rate: float
+) -> List[torch.Tensor]:
+ """Read a list of sound files into a list 1-D float32 torch tensors.
+ Args:
+ filenames:
+ A list of sound filenames.
+ expected_sample_rate:
+ The expected sample rate of the sound files.
+ Returns:
+ Return a list of 1-D float32 torch tensors.
+ """
+ ans = []
+ for f in filenames:
+ wave, sample_rate = torchaudio.load(f)
+ assert (
+ sample_rate == expected_sample_rate
+ ), f"expected sample rate: {expected_sample_rate}. Given: {sample_rate}"
+ # We use only the first channel
+ ans.append(wave[0])
+ return ans
+
+
+@torch.no_grad()
+def main():
+ parser = get_parser()
+ args = parser.parse_args()
+
+ params = get_params()
+
+ params.update(vars(args))
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(params.bpe_model)
+
+ # is defined in local/train_bpe_model.py
+ params.blank_id = sp.piece_to_id("")
+ params.unk_id = sp.piece_to_id("")
+ params.vocab_size = sp.get_piece_size()
+
+ logging.info(f"{params}")
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"device: {device}")
+
+ logging.info("Creating model")
+ model = get_transducer_model(params)
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ checkpoint = torch.load(args.checkpoint, map_location="cpu")
+ model.load_state_dict(checkpoint["model"], strict=False)
+ model.to(device)
+ model.eval()
+ model.device = device
+
+ logging.info("Constructing Fbank computer")
+ opts = kaldifeat.FbankOptions()
+ opts.device = device
+ opts.frame_opts.dither = 0
+ opts.frame_opts.snip_edges = False
+ opts.frame_opts.samp_freq = params.sample_rate
+ opts.mel_opts.num_bins = params.feature_dim
+
+ fbank = kaldifeat.Fbank(opts)
+
+ logging.info(f"Reading sound files: {params.sound_files}")
+ waves = read_sound_files(
+ filenames=params.sound_files, expected_sample_rate=params.sample_rate
+ )
+ waves = [w.to(device) for w in waves]
+
+ logging.info("Decoding started")
+ features = fbank(waves)
+ feature_lengths = [f.size(0) for f in features]
+
+ features = pad_sequence(features, batch_first=True, padding_value=math.log(1e-10))
+
+ feature_lengths = torch.tensor(feature_lengths, device=device)
+
+ encoder_out, encoder_out_lens = model.encoder(x=features, x_lens=feature_lengths)
+
+ num_waves = encoder_out.size(0)
+ hyps = []
+ msg = f"Using {params.method}"
+ if params.method == "beam_search":
+ msg += f" with beam size {params.beam_size}"
+ logging.info(msg)
+
+ if params.method == "fast_beam_search":
+ decoding_graph = k2.trivial_graph(params.vocab_size - 1, device=device)
+ hyp_tokens = fast_beam_search_one_best(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.method == "modified_beam_search":
+ hyp_tokens = modified_beam_search(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=params.beam_size,
+ )
+
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ elif params.method == "greedy_search" and params.max_sym_per_frame == 1:
+ hyp_tokens = greedy_search_batch(
+ model=model,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ )
+ for hyp in sp.decode(hyp_tokens):
+ hyps.append(hyp.split())
+ else:
+ for i in range(num_waves):
+ # fmt: off
+ encoder_out_i = encoder_out[i:i+1, :encoder_out_lens[i]]
+ # fmt: on
+ if params.method == "greedy_search":
+ hyp = greedy_search(
+ model=model,
+ encoder_out=encoder_out_i,
+ max_sym_per_frame=params.max_sym_per_frame,
+ )
+ elif params.method == "beam_search":
+ hyp = beam_search(
+ model=model,
+ encoder_out=encoder_out_i,
+ beam=params.beam_size,
+ )
+ else:
+ raise ValueError(f"Unsupported method: {params.method}")
+
+ hyps.append(sp.decode(hyp).split())
+
+ s = "\n"
+ for filename, hyp in zip(params.sound_files, hyps):
+ words = " ".join(hyp)
+ s += f"{filename}:\n{words}\n\n"
+ logging.info(s)
+
+ logging.info("Decoding Done")
+
+
+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/pruned_transducer_stateless7_streaming/scaling.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/scaling.py
new file mode 120000
index 000000000..2428b74b9
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/scaling.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/scaling.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/scaling_converter.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/scaling_converter.py
new file mode 120000
index 000000000..b8b8ba432
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/scaling_converter.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/scaling_converter.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/streaming_beam_search.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/streaming_beam_search.py
new file mode 120000
index 000000000..3a5f89833
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/streaming_beam_search.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/streaming_beam_search.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/streaming_decode.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/streaming_decode.py
new file mode 100755
index 000000000..7a349ecb2
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/streaming_decode.py
@@ -0,0 +1,615 @@
+#!/usr/bin/env python3
+# Copyright 2022 Xiaomi Corporation (Authors: Wei Kang, Fangjun Kuang)
+#
+# 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.
+
+"""
+Usage:
+./pruned_transducer_stateless7_streaming/streaming_decode.py \
+ --epoch 28 \
+ --avg 15 \
+ --decode-chunk-len 32 \
+ --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+ --decoding_method greedy_search \
+ --num-decode-streams 2000
+"""
+
+import argparse
+import logging
+import math
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import k2
+import numpy as np
+import sentencepiece as spm
+import torch
+import torch.nn as nn
+from asr_datamodule import LibriSpeechAsrDataModule
+from decode_stream import DecodeStream
+from kaldifeat import Fbank, FbankOptions
+from lhotse import CutSet
+from streaming_beam_search import (
+ fast_beam_search_one_best,
+ greedy_search,
+ modified_beam_search,
+)
+from torch.nn.utils.rnn import pad_sequence
+from train import add_model_arguments, get_params, get_transducer_model
+from zipformer import stack_states, unstack_states
+
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from icefall.utils import (
+ AttributeDict,
+ setup_logger,
+ store_transcripts,
+ str2bool,
+ write_error_stats,
+)
+
+LOG_EPS = math.log(1e-10)
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=28,
+ help="""It specifies the checkpoint to use for decoding.
+ Note: Epoch counts from 0.
+ You can specify --avg to use more checkpoints for model averaging.""",
+ )
+
+ parser.add_argument(
+ "--iter",
+ type=int,
+ default=0,
+ help="""If positive, --epoch is ignored and it
+ will use the checkpoint exp_dir/checkpoint-iter.pt.
+ You can specify --avg to use more checkpoints for model averaging.
+ """,
+ )
+
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=15,
+ help="Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch' and '--iter'",
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=True,
+ help="Whether to load averaged model. Currently it only supports "
+ "using --epoch. If True, it would decode with the averaged model "
+ "over the epoch range from `epoch-avg` (excluded) to `epoch`."
+ "Actually only the models with epoch number of `epoch-avg` and "
+ "`epoch` are loaded for averaging. ",
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="pruned_transducer_stateless2/exp",
+ help="The experiment dir",
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ default="data/lang_bpe_500/bpe.model",
+ help="Path to the BPE model",
+ )
+
+ parser.add_argument(
+ "--decoding-method",
+ type=str,
+ default="greedy_search",
+ help="""Supported decoding methods are:
+ greedy_search
+ modified_beam_search
+ fast_beam_search
+ """,
+ )
+
+ parser.add_argument(
+ "--num_active_paths",
+ type=int,
+ default=4,
+ help="""An interger indicating how many candidates we will keep for each
+ frame. Used only when --decoding-method is modified_beam_search.""",
+ )
+
+ parser.add_argument(
+ "--beam",
+ type=float,
+ default=4,
+ help="""A floating point value to calculate the cutoff score during beam
+ search (i.e., `cutoff = max-score - beam`), which is the same as the
+ `beam` in Kaldi.
+ Used only when --decoding-method is fast_beam_search""",
+ )
+
+ parser.add_argument(
+ "--max-contexts",
+ type=int,
+ default=4,
+ help="""Used only when --decoding-method is
+ fast_beam_search""",
+ )
+
+ parser.add_argument(
+ "--max-states",
+ type=int,
+ default=32,
+ help="""Used only when --decoding-method is
+ fast_beam_search""",
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram; 2 means tri-gram",
+ )
+
+ parser.add_argument(
+ "--num-decode-streams",
+ type=int,
+ default=2000,
+ help="The number of streams that can be decoded parallel.",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def decode_one_chunk(
+ params: AttributeDict,
+ model: nn.Module,
+ decode_streams: List[DecodeStream],
+) -> List[int]:
+ """Decode one chunk frames of features for each decode_streams and
+ return the indexes of finished streams in a List.
+
+ Args:
+ params:
+ It's the return value of :func:`get_params`.
+ model:
+ The neural model.
+ decode_streams:
+ A List of DecodeStream, each belonging to a utterance.
+ Returns:
+ Return a List containing which DecodeStreams are finished.
+ """
+ device = model.device
+
+ features = []
+ feature_lens = []
+ states = []
+ processed_lens = []
+
+ for stream in decode_streams:
+ feat, feat_len = stream.get_feature_frames(params.decode_chunk_len)
+ features.append(feat)
+ feature_lens.append(feat_len)
+ states.append(stream.states)
+ processed_lens.append(stream.done_frames)
+
+ feature_lens = torch.tensor(feature_lens, device=device)
+ features = pad_sequence(features, batch_first=True, padding_value=LOG_EPS)
+
+ # We subsample features with ((x_len - 7) // 2 + 1) // 2 and the max downsampling
+ # factor in encoders is 8.
+ # After feature embedding (x_len - 7) // 2, we have (23 - 7) // 2 = 8.
+ tail_length = 23
+ if features.size(1) < tail_length:
+ pad_length = tail_length - features.size(1)
+ feature_lens += pad_length
+ features = torch.nn.functional.pad(
+ features,
+ (0, 0, 0, pad_length),
+ mode="constant",
+ value=LOG_EPS,
+ )
+
+ states = stack_states(states)
+ processed_lens = torch.tensor(processed_lens, device=device)
+
+ encoder_out, encoder_out_lens, new_states = model.encoder.streaming_forward(
+ x=features,
+ x_lens=feature_lens,
+ states=states,
+ )
+
+ encoder_out = model.joiner.encoder_proj(encoder_out)
+
+ if params.decoding_method == "greedy_search":
+ greedy_search(model=model, encoder_out=encoder_out, streams=decode_streams)
+ elif params.decoding_method == "fast_beam_search":
+ processed_lens = processed_lens + encoder_out_lens
+ fast_beam_search_one_best(
+ model=model,
+ encoder_out=encoder_out,
+ processed_lens=processed_lens,
+ streams=decode_streams,
+ beam=params.beam,
+ max_states=params.max_states,
+ max_contexts=params.max_contexts,
+ )
+ elif params.decoding_method == "modified_beam_search":
+ modified_beam_search(
+ model=model,
+ streams=decode_streams,
+ encoder_out=encoder_out,
+ num_active_paths=params.num_active_paths,
+ )
+ else:
+ raise ValueError(f"Unsupported decoding method: {params.decoding_method}")
+
+ states = unstack_states(new_states)
+
+ finished_streams = []
+ for i in range(len(decode_streams)):
+ decode_streams[i].states = states[i]
+ decode_streams[i].done_frames += encoder_out_lens[i]
+ if decode_streams[i].done:
+ finished_streams.append(i)
+
+ return finished_streams
+
+
+def decode_dataset(
+ cuts: CutSet,
+ params: AttributeDict,
+ model: nn.Module,
+ sp: spm.SentencePieceProcessor,
+ decoding_graph: Optional[k2.Fsa] = None,
+) -> Dict[str, List[Tuple[List[str], List[str]]]]:
+ """Decode dataset.
+
+ Args:
+ cuts:
+ Lhotse Cutset containing the dataset to decode.
+ params:
+ It is returned by :func:`get_params`.
+ model:
+ The neural model.
+ sp:
+ The BPE model.
+ decoding_graph:
+ The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+ only when --decoding_method is fast_beam_search.
+ Returns:
+ Return a dict, whose key may be "greedy_search" if greedy search
+ is used, or it may be "beam_7" if beam size of 7 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.
+ """
+ device = model.device
+
+ opts = FbankOptions()
+ opts.device = device
+ opts.frame_opts.dither = 0
+ opts.frame_opts.snip_edges = False
+ opts.frame_opts.samp_freq = 16000
+ opts.mel_opts.num_bins = 80
+
+ log_interval = 50
+
+ decode_results = []
+ # Contain decode streams currently running.
+ decode_streams = []
+ for num, cut in enumerate(cuts):
+ # each utterance has a DecodeStream.
+ initial_states = model.encoder.get_init_state(device=device)
+ decode_stream = DecodeStream(
+ params=params,
+ cut_id=cut.id,
+ initial_states=initial_states,
+ decoding_graph=decoding_graph,
+ device=device,
+ )
+
+ audio: np.ndarray = cut.load_audio()
+ # audio.shape: (1, num_samples)
+ assert len(audio.shape) == 2
+ assert audio.shape[0] == 1, "Should be single channel"
+ assert audio.dtype == np.float32, audio.dtype
+
+ # The trained model is using normalized samples
+ assert audio.max() <= 1, "Should be normalized to [-1, 1])"
+
+ samples = torch.from_numpy(audio).squeeze(0)
+
+ fbank = Fbank(opts)
+ feature = fbank(samples.to(device))
+ decode_stream.set_features(feature, tail_pad_len=params.decode_chunk_len)
+ decode_stream.ground_truth = cut.supervisions[0].text
+
+ decode_streams.append(decode_stream)
+
+ while len(decode_streams) >= params.num_decode_streams:
+ finished_streams = decode_one_chunk(
+ params=params, model=model, decode_streams=decode_streams
+ )
+ for i in sorted(finished_streams, reverse=True):
+ decode_results.append(
+ (
+ decode_streams[i].id,
+ decode_streams[i].ground_truth.split(),
+ sp.decode(decode_streams[i].decoding_result()).split(),
+ )
+ )
+ del decode_streams[i]
+
+ if num % log_interval == 0:
+ logging.info(f"Cuts processed until now is {num}.")
+
+ # decode final chunks of last sequences
+ while len(decode_streams):
+ finished_streams = decode_one_chunk(
+ params=params, model=model, decode_streams=decode_streams
+ )
+ for i in sorted(finished_streams, reverse=True):
+ decode_results.append(
+ (
+ decode_streams[i].id,
+ decode_streams[i].ground_truth.split(),
+ sp.decode(decode_streams[i].decoding_result()).split(),
+ )
+ )
+ del decode_streams[i]
+
+ if params.decoding_method == "greedy_search":
+ key = "greedy_search"
+ elif params.decoding_method == "fast_beam_search":
+ key = (
+ f"beam_{params.beam}_"
+ f"max_contexts_{params.max_contexts}_"
+ f"max_states_{params.max_states}"
+ )
+ elif params.decoding_method == "modified_beam_search":
+ key = f"num_active_paths_{params.num_active_paths}"
+ else:
+ raise ValueError(f"Unsupported decoding method: {params.decoding_method}")
+ return {key: decode_results}
+
+
+def save_results(
+ params: AttributeDict,
+ test_set_name: str,
+ results_dict: Dict[str, List[Tuple[List[str], List[str]]]],
+):
+ test_set_wers = dict()
+ for key, results in results_dict.items():
+ recog_path = (
+ params.res_dir / f"recogs-{test_set_name}-{key}-{params.suffix}.txt"
+ )
+ results = sorted(results)
+ 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.res_dir / f"errs-{test_set_name}-{key}-{params.suffix}.txt"
+ )
+ with open(errs_filename, "w") as f:
+ wer = write_error_stats(
+ f, f"{test_set_name}-{key}", results, enable_log=True
+ )
+ 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.res_dir / f"wer-summary-{test_set_name}-{key}-{params.suffix}.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()
+ args.exp_dir = Path(args.exp_dir)
+
+ params = get_params()
+ params.update(vars(args))
+
+ params.res_dir = params.exp_dir / "streaming" / params.decoding_method
+
+ if params.iter > 0:
+ params.suffix = f"iter-{params.iter}-avg-{params.avg}"
+ else:
+ params.suffix = f"epoch-{params.epoch}-avg-{params.avg}"
+
+ # for streaming
+ params.suffix += f"-streaming-chunk-size-{params.decode_chunk_len}"
+
+ # for fast_beam_search
+ if params.decoding_method == "fast_beam_search":
+ params.suffix += f"-beam-{params.beam}"
+ params.suffix += f"-max-contexts-{params.max_contexts}"
+ params.suffix += f"-max-states-{params.max_states}"
+
+ if params.use_averaged_model:
+ params.suffix += "-use-averaged-model"
+
+ setup_logger(f"{params.res_dir}/log-decode-{params.suffix}")
+ logging.info("Decoding started")
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"Device: {device}")
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(params.bpe_model)
+
+ # and is defined in local/train_bpe_model.py
+ params.blank_id = sp.piece_to_id("")
+ params.unk_id = sp.piece_to_id("")
+ params.vocab_size = sp.get_piece_size()
+
+ logging.info(params)
+
+ logging.info("About to create model")
+ model = get_transducer_model(params)
+
+ if not params.use_averaged_model:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ elif 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.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg + 1
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg + 1:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ filename_start = filenames[-1]
+ filename_end = filenames[0]
+ logging.info(
+ "Calculating the averaged model over iteration checkpoints"
+ f" from {filename_start} (excluded) to {filename_end}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+ else:
+ assert params.avg > 0, params.avg
+ start = params.epoch - params.avg
+ assert start >= 1, start
+ filename_start = f"{params.exp_dir}/epoch-{start}.pt"
+ filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
+ logging.info(
+ f"Calculating the averaged model over epoch range from "
+ f"{start} (excluded) to {params.epoch}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+
+ model.to(device)
+ model.eval()
+ model.device = device
+
+ decoding_graph = None
+ if params.decoding_method == "fast_beam_search":
+ decoding_graph = k2.trivial_graph(params.vocab_size - 1, device=device)
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ librispeech = LibriSpeechAsrDataModule(args)
+
+ test_clean_cuts = librispeech.test_clean_cuts()
+ test_other_cuts = librispeech.test_other_cuts()
+
+ test_sets = ["test-clean", "test-other"]
+ test_cuts = [test_clean_cuts, test_other_cuts]
+
+ for test_set, test_cut in zip(test_sets, test_cuts):
+ results_dict = decode_dataset(
+ cuts=test_cut,
+ params=params,
+ model=model,
+ sp=sp,
+ decoding_graph=decoding_graph,
+ )
+
+ 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/pruned_transducer_stateless7_streaming/test_model.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/test_model.py
new file mode 100755
index 000000000..5400df804
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/test_model.py
@@ -0,0 +1,150 @@
+#!/usr/bin/env python3
+# Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang)
+#
+# 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.
+
+
+"""
+To run this file, do:
+
+ cd icefall/egs/librispeech/ASR
+ python ./pruned_transducer_stateless7_streaming/test_model.py
+"""
+
+import torch
+from scaling_converter import convert_scaled_to_non_scaled
+from train import get_params, get_transducer_model
+
+
+def test_model():
+ params = get_params()
+ params.vocab_size = 500
+ params.blank_id = 0
+ params.context_size = 2
+ params.num_encoder_layers = "2,4,3,2,4"
+ params.feedforward_dims = "1024,1024,2048,2048,1024"
+ params.nhead = "8,8,8,8,8"
+ params.encoder_dims = "384,384,384,384,384"
+ params.attention_dims = "192,192,192,192,192"
+ params.encoder_unmasked_dims = "256,256,256,256,256"
+ params.zipformer_downsampling_factors = "1,2,4,8,2"
+ params.cnn_module_kernels = "31,31,31,31,31"
+ params.decoder_dim = 512
+ params.joiner_dim = 512
+ params.num_left_chunks = 4
+ params.short_chunk_size = 50
+ params.decode_chunk_len = 32
+ model = get_transducer_model(params)
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ print(f"Number of model parameters: {num_param}")
+
+ # Test jit script
+ convert_scaled_to_non_scaled(model, inplace=True)
+ # We won't use the forward() method of the model in C++, so just ignore
+ # it here.
+ # Otherwise, one of its arguments is a ragged tensor and is not
+ # torch scriptabe.
+ model.__class__.forward = torch.jit.ignore(model.__class__.forward)
+ print("Using torch.jit.script")
+ model = torch.jit.script(model)
+
+
+def test_model_jit_trace():
+ params = get_params()
+ params.vocab_size = 500
+ params.blank_id = 0
+ params.context_size = 2
+ params.num_encoder_layers = "2,4,3,2,4"
+ params.feedforward_dims = "1024,1024,2048,2048,1024"
+ params.nhead = "8,8,8,8,8"
+ params.encoder_dims = "384,384,384,384,384"
+ params.attention_dims = "192,192,192,192,192"
+ params.encoder_unmasked_dims = "256,256,256,256,256"
+ params.zipformer_downsampling_factors = "1,2,4,8,2"
+ params.cnn_module_kernels = "31,31,31,31,31"
+ params.decoder_dim = 512
+ params.joiner_dim = 512
+ params.num_left_chunks = 4
+ params.short_chunk_size = 50
+ params.decode_chunk_len = 32
+ model = get_transducer_model(params)
+ model.eval()
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ print(f"Number of model parameters: {num_param}")
+
+ convert_scaled_to_non_scaled(model, inplace=True)
+
+ # Test encoder
+ def _test_encoder():
+ encoder = model.encoder
+ assert encoder.decode_chunk_size == params.decode_chunk_len // 2, (
+ encoder.decode_chunk_size,
+ params.decode_chunk_len,
+ )
+ T = params.decode_chunk_len + 7
+
+ x = torch.zeros(1, T, 80, dtype=torch.float32)
+ x_lens = torch.full((1,), T, dtype=torch.int32)
+ states = encoder.get_init_state(device=x.device)
+ encoder.__class__.forward = encoder.__class__.streaming_forward
+ traced_encoder = torch.jit.trace(encoder, (x, x_lens, states))
+
+ states1 = encoder.get_init_state(device=x.device)
+ states2 = traced_encoder.get_init_state(device=x.device)
+ for i in range(5):
+ x = torch.randn(1, T, 80, dtype=torch.float32)
+ x_lens = torch.full((1,), T, dtype=torch.int32)
+ y1, _, states1 = encoder.streaming_forward(x, x_lens, states1)
+ y2, _, states2 = traced_encoder(x, x_lens, states2)
+ assert torch.allclose(y1, y2, atol=1e-6), (i, (y1 - y2).abs().mean())
+
+ # Test decoder
+ def _test_decoder():
+ decoder = model.decoder
+ y = torch.zeros(10, decoder.context_size, dtype=torch.int64)
+ need_pad = torch.tensor([False])
+
+ traced_decoder = torch.jit.trace(decoder, (y, need_pad))
+ d1 = decoder(y, need_pad)
+ d2 = traced_decoder(y, need_pad)
+ assert torch.equal(d1, d2), (d1 - d2).abs().mean()
+
+ # Test joiner
+ def _test_joiner():
+ joiner = model.joiner
+ encoder_out_dim = joiner.encoder_proj.weight.shape[1]
+ decoder_out_dim = joiner.decoder_proj.weight.shape[1]
+ encoder_out = torch.rand(1, encoder_out_dim, dtype=torch.float32)
+ decoder_out = torch.rand(1, decoder_out_dim, dtype=torch.float32)
+
+ traced_joiner = torch.jit.trace(joiner, (encoder_out, decoder_out))
+ j1 = joiner(encoder_out, decoder_out)
+ j2 = traced_joiner(encoder_out, decoder_out)
+ assert torch.equal(j1, j2), (j1 - j2).abs().mean()
+
+ _test_encoder()
+ _test_decoder()
+ _test_joiner()
+
+
+def main():
+ test_model()
+ test_model_jit_trace()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/train.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/train.py
new file mode 100755
index 000000000..2bdc882a5
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/train.py
@@ -0,0 +1,1264 @@
+#!/usr/bin/env python3
+# Copyright 2021-2022 Xiaomi Corp. (authors: Fangjun Kuang,
+# Wei Kang,
+# Mingshuang Luo,)
+# Zengwei Yao)
+#
+# 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.
+"""
+Usage:
+
+export CUDA_VISIBLE_DEVICES="0,1,2,3"
+
+./pruned_transducer_stateless7_streaming/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --exp-dir pruned_transducer_stateless7_streaming/exp \
+ --full-libri 1 \
+ --max-duration 300
+
+# For mix precision training:
+
+./pruned_transducer_stateless7_streaming/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --use-fp16 1 \
+ --exp-dir pruned_transducer_stateless7_streaming/exp \
+ --full-libri 1 \
+ --max-duration 550
+"""
+
+
+import argparse
+import copy
+import logging
+import warnings
+from pathlib import Path
+from shutil import copyfile
+from typing import Any, Dict, Optional, Tuple, Union
+
+import k2
+import optim
+import sentencepiece as spm
+import torch
+import torch.multiprocessing as mp
+import torch.nn as nn
+from asr_datamodule import LibriSpeechAsrDataModule
+from decoder import Decoder
+from joiner import Joiner
+from lhotse.cut import Cut
+from lhotse.dataset.sampling.base import CutSampler
+from lhotse.utils import fix_random_seed
+from model import Transducer
+from optim import Eden, ScaledAdam
+from torch import Tensor
+from torch.cuda.amp import GradScaler
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.tensorboard import SummaryWriter
+from zipformer import Zipformer
+
+from icefall import diagnostics
+from icefall.checkpoint import load_checkpoint, remove_checkpoints
+from icefall.checkpoint import save_checkpoint as save_checkpoint_impl
+from icefall.checkpoint import (
+ save_checkpoint_with_global_batch_idx,
+ update_averaged_model,
+)
+from icefall.dist import cleanup_dist, setup_dist
+from icefall.env import get_env_info
+from icefall.hooks import register_inf_check_hooks
+from icefall.utils import AttributeDict, MetricsTracker, setup_logger, str2bool
+
+LRSchedulerType = Union[torch.optim.lr_scheduler._LRScheduler, optim.LRScheduler]
+
+
+def set_batch_count(model: Union[nn.Module, DDP], batch_count: float) -> None:
+ if isinstance(model, DDP):
+ # get underlying nn.Module
+ model = model.module
+ for module in model.modules():
+ if hasattr(module, "batch_count"):
+ module.batch_count = batch_count
+
+
+def add_model_arguments(parser: argparse.ArgumentParser):
+ parser.add_argument(
+ "--num-encoder-layers",
+ type=str,
+ default="2,4,3,2,4",
+ help="Number of zipformer encoder layers, comma separated.",
+ )
+
+ parser.add_argument(
+ "--feedforward-dims",
+ type=str,
+ default="1024,1024,2048,2048,1024",
+ help="Feedforward dimension of the zipformer encoder layers, comma separated.",
+ )
+
+ parser.add_argument(
+ "--nhead",
+ type=str,
+ default="8,8,8,8,8",
+ help="Number of attention heads in the zipformer encoder layers.",
+ )
+
+ parser.add_argument(
+ "--encoder-dims",
+ type=str,
+ default="384,384,384,384,384",
+ help="Embedding dimension in the 2 blocks of zipformer encoder layers, comma separated",
+ )
+
+ parser.add_argument(
+ "--attention-dims",
+ type=str,
+ default="192,192,192,192,192",
+ help="""Attention dimension in the 2 blocks of zipformer encoder layers, comma separated;
+ not the same as embedding dimension.""",
+ )
+
+ parser.add_argument(
+ "--encoder-unmasked-dims",
+ type=str,
+ default="256,256,256,256,256",
+ help="Unmasked dimensions in the encoders, relates to augmentation during training. "
+ "Must be <= each of encoder_dims. Empirically, less than 256 seems to make performance "
+ " worse.",
+ )
+
+ parser.add_argument(
+ "--zipformer-downsampling-factors",
+ type=str,
+ default="1,2,4,8,2",
+ help="Downsampling factor for each stack of encoder layers.",
+ )
+
+ parser.add_argument(
+ "--cnn-module-kernels",
+ type=str,
+ default="31,31,31,31,31",
+ help="Sizes of kernels in convolution modules",
+ )
+
+ parser.add_argument(
+ "--decoder-dim",
+ type=int,
+ default=512,
+ help="Embedding dimension in the decoder model.",
+ )
+
+ parser.add_argument(
+ "--joiner-dim",
+ type=int,
+ default=512,
+ help="""Dimension used in the joiner model.
+ Outputs from the encoder and decoder model are projected
+ to this dimension before adding.
+ """,
+ )
+
+ parser.add_argument(
+ "--short-chunk-size",
+ type=int,
+ default=50,
+ help="""Chunk length of dynamic training, the chunk size would be either
+ max sequence length of current batch or uniformly sampled from (1, short_chunk_size).
+ """,
+ )
+
+ parser.add_argument(
+ "--num-left-chunks",
+ type=int,
+ default=4,
+ help="How many left context can be seen in chunks when calculating attention.",
+ )
+
+ parser.add_argument(
+ "--decode-chunk-len",
+ type=int,
+ default=32,
+ help="The chunk size for decoding (in frames before subsampling)",
+ )
+
+
+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.",
+ )
+
+ parser.add_argument(
+ "--num-epochs",
+ type=int,
+ default=30,
+ help="Number of epochs to train.",
+ )
+
+ parser.add_argument(
+ "--start-epoch",
+ type=int,
+ default=1,
+ help="""Resume training from this epoch. It should be positive.
+ If larger than 1, it will load checkpoint from
+ exp-dir/epoch-{start_epoch-1}.pt
+ """,
+ )
+
+ parser.add_argument(
+ "--start-batch",
+ type=int,
+ default=0,
+ help="""If positive, --start-epoch is ignored and
+ it loads the checkpoint from exp-dir/checkpoint-{start_batch}.pt
+ """,
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="pruned_transducer_stateless7_streaming/exp",
+ help="""The experiment dir.
+ It specifies the directory where all training related
+ files, e.g., checkpoints, log, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ default="data/lang_bpe_500/bpe.model",
+ help="Path to the BPE model",
+ )
+
+ parser.add_argument(
+ "--base-lr", type=float, default=0.05, help="The base learning rate."
+ )
+
+ parser.add_argument(
+ "--lr-batches",
+ type=float,
+ default=5000,
+ help="""Number of steps that affects how rapidly the learning rate
+ decreases. We suggest not to change this.""",
+ )
+
+ parser.add_argument(
+ "--lr-epochs",
+ type=float,
+ default=3.5,
+ help="""Number of epochs that affects how rapidly the learning rate decreases.
+ """,
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram; 2 means tri-gram",
+ )
+
+ parser.add_argument(
+ "--prune-range",
+ type=int,
+ default=5,
+ help="The prune range for rnnt loss, it means how many symbols(context)"
+ "we are using to compute the loss",
+ )
+
+ parser.add_argument(
+ "--lm-scale",
+ type=float,
+ default=0.25,
+ help="The scale to smooth the loss with lm "
+ "(output of prediction network) part.",
+ )
+
+ parser.add_argument(
+ "--am-scale",
+ type=float,
+ default=0.0,
+ help="The scale to smooth the loss with am (output of encoder network) part.",
+ )
+
+ parser.add_argument(
+ "--simple-loss-scale",
+ type=float,
+ default=0.5,
+ help="To get pruning ranges, we will calculate a simple version"
+ "loss(joiner is just addition), this simple loss also uses for"
+ "training (as a regularization item). We will scale the simple loss"
+ "with this parameter before adding to the final loss.",
+ )
+
+ parser.add_argument(
+ "--seed",
+ type=int,
+ default=42,
+ help="The seed for random generators intended for reproducibility",
+ )
+
+ parser.add_argument(
+ "--print-diagnostics",
+ type=str2bool,
+ default=False,
+ help="Accumulate stats on activations, print them and exit.",
+ )
+
+ parser.add_argument(
+ "--inf-check",
+ type=str2bool,
+ default=False,
+ help="Add hooks to check for infinite module outputs and gradients.",
+ )
+
+ parser.add_argument(
+ "--save-every-n",
+ type=int,
+ default=2000,
+ help="""Save checkpoint after processing this number of batches"
+ periodically. We save checkpoint to exp-dir/ whenever
+ params.batch_idx_train % save_every_n == 0. The checkpoint filename
+ has the form: f'exp-dir/checkpoint-{params.batch_idx_train}.pt'
+ Note: It also saves checkpoint to `exp-dir/epoch-xxx.pt` at the
+ end of each epoch where `xxx` is the epoch number counting from 0.
+ """,
+ )
+
+ parser.add_argument(
+ "--keep-last-k",
+ type=int,
+ default=30,
+ help="""Only keep this number of checkpoints on disk.
+ For instance, if it is 3, there are only 3 checkpoints
+ in the exp-dir with filenames `checkpoint-xxx.pt`.
+ It does not affect checkpoints with name `epoch-xxx.pt`.
+ """,
+ )
+
+ parser.add_argument(
+ "--average-period",
+ type=int,
+ default=200,
+ help="""Update the averaged model, namely `model_avg`, after processing
+ this number of batches. `model_avg` is a separate version of model,
+ in which each floating-point parameter is the average of all the
+ parameters from the start of training. Each time we take the average,
+ we do: `model_avg = model * (average_period / batch_idx_train) +
+ model_avg * ((batch_idx_train - average_period) / batch_idx_train)`.
+ """,
+ )
+
+ parser.add_argument(
+ "--use-fp16",
+ type=str2bool,
+ default=False,
+ help="Whether to use half precision training.",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def get_params() -> AttributeDict:
+ """Return a dict containing training parameters.
+
+ All training related parameters that are not passed from the commandline
+ are 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`:
+
+ - 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
+
+ - reset_interval: Reset statistics if batch_idx % reset_interval is 0
+
+ - valid_interval: Run validation if batch_idx % valid_interval is 0
+
+ - feature_dim: The model input dim. It has to match the one used
+ in computing features.
+
+ - subsampling_factor: The subsampling factor for the model.
+
+ - encoder_dim: Hidden dim for multi-head attention model.
+
+ - num_decoder_layers: Number of decoder layer of transformer decoder.
+
+ - warm_step: The warmup period that dictates the decay of the
+ scale on "simple" (un-pruned) loss.
+ """
+ params = AttributeDict(
+ {
+ "best_train_loss": float("inf"),
+ "best_valid_loss": float("inf"),
+ "best_train_epoch": -1,
+ "best_valid_epoch": -1,
+ "batch_idx_train": 0,
+ "log_interval": 50,
+ "reset_interval": 200,
+ "valid_interval": 3000, # For the 100h subset, use 800
+ # parameters for zipformer
+ "feature_dim": 80,
+ "subsampling_factor": 4, # not passed in, this is fixed.
+ "warm_step": 2000,
+ "env_info": get_env_info(),
+ }
+ )
+
+ return params
+
+
+def get_encoder_model(params: AttributeDict) -> nn.Module:
+ # TODO: We can add an option to switch between Zipformer and Transformer
+ def to_int_tuple(s: str):
+ return tuple(map(int, s.split(",")))
+
+ encoder = Zipformer(
+ num_features=params.feature_dim,
+ output_downsampling_factor=2,
+ zipformer_downsampling_factors=to_int_tuple(
+ params.zipformer_downsampling_factors
+ ),
+ encoder_dims=to_int_tuple(params.encoder_dims),
+ attention_dim=to_int_tuple(params.attention_dims),
+ encoder_unmasked_dims=to_int_tuple(params.encoder_unmasked_dims),
+ nhead=to_int_tuple(params.nhead),
+ feedforward_dim=to_int_tuple(params.feedforward_dims),
+ cnn_module_kernels=to_int_tuple(params.cnn_module_kernels),
+ num_encoder_layers=to_int_tuple(params.num_encoder_layers),
+ num_left_chunks=params.num_left_chunks,
+ short_chunk_size=params.short_chunk_size,
+ decode_chunk_size=params.decode_chunk_len // 2,
+ )
+ return encoder
+
+
+def get_decoder_model(params: AttributeDict) -> nn.Module:
+ decoder = Decoder(
+ vocab_size=params.vocab_size,
+ decoder_dim=params.decoder_dim,
+ blank_id=params.blank_id,
+ context_size=params.context_size,
+ )
+ return decoder
+
+
+def get_joiner_model(params: AttributeDict) -> nn.Module:
+ joiner = Joiner(
+ encoder_dim=int(params.encoder_dims.split(",")[-1]),
+ decoder_dim=params.decoder_dim,
+ joiner_dim=params.joiner_dim,
+ vocab_size=params.vocab_size,
+ )
+ return joiner
+
+
+def get_transducer_model(params: AttributeDict) -> nn.Module:
+ encoder = get_encoder_model(params)
+ decoder = get_decoder_model(params)
+ joiner = get_joiner_model(params)
+
+ model = Transducer(
+ encoder=encoder,
+ decoder=decoder,
+ joiner=joiner,
+ encoder_dim=int(params.encoder_dims.split(",")[-1]),
+ decoder_dim=params.decoder_dim,
+ joiner_dim=params.joiner_dim,
+ vocab_size=params.vocab_size,
+ )
+ return model
+
+
+def load_checkpoint_if_available(
+ params: AttributeDict,
+ model: nn.Module,
+ model_avg: nn.Module = None,
+ optimizer: Optional[torch.optim.Optimizer] = None,
+ scheduler: Optional[LRSchedulerType] = None,
+) -> Optional[Dict[str, Any]]:
+ """Load checkpoint from file.
+
+ If params.start_batch is positive, it will load the checkpoint from
+ `params.exp_dir/checkpoint-{params.start_batch}.pt`. Otherwise, if
+ params.start_epoch is larger than 1, it will load the checkpoint from
+ `params.start_epoch - 1`.
+
+ Apart from loading state dict for `model` and `optimizer` 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.
+ model_avg:
+ The stored model averaged from the start of training.
+ optimizer:
+ The optimizer that we are using.
+ scheduler:
+ The scheduler that we are using.
+ Returns:
+ Return a dict containing previously saved training info.
+ """
+ if params.start_batch > 0:
+ filename = params.exp_dir / f"checkpoint-{params.start_batch}.pt"
+ elif params.start_epoch > 1:
+ filename = params.exp_dir / f"epoch-{params.start_epoch-1}.pt"
+ else:
+ return None
+
+ assert filename.is_file(), f"{filename} does not exist!"
+
+ saved_params = load_checkpoint(
+ filename,
+ model=model,
+ model_avg=model_avg,
+ 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]
+
+ if params.start_batch > 0:
+ if "cur_epoch" in saved_params:
+ params["start_epoch"] = saved_params["cur_epoch"]
+
+ if "cur_batch_idx" in saved_params:
+ params["cur_batch_idx"] = saved_params["cur_batch_idx"]
+
+ return saved_params
+
+
+def save_checkpoint(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ model_avg: Optional[nn.Module] = None,
+ optimizer: Optional[torch.optim.Optimizer] = None,
+ scheduler: Optional[LRSchedulerType] = None,
+ sampler: Optional[CutSampler] = None,
+ scaler: Optional[GradScaler] = 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.
+ model_avg:
+ The stored model averaged from the start of training.
+ optimizer:
+ The optimizer used in the training.
+ sampler:
+ The sampler for the training dataset.
+ scaler:
+ The scaler used for mix precision training.
+ """
+ if rank != 0:
+ return
+ filename = params.exp_dir / f"epoch-{params.cur_epoch}.pt"
+ save_checkpoint_impl(
+ filename=filename,
+ model=model,
+ model_avg=model_avg,
+ params=params,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=sampler,
+ scaler=scaler,
+ 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: Union[nn.Module, DDP],
+ sp: spm.SentencePieceProcessor,
+ batch: dict,
+ is_training: bool,
+) -> Tuple[Tensor, MetricsTracker]:
+ """
+ Compute transducer 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 Zipformer in our case.
+ batch:
+ A batch of data. See `lhotse.dataset.K2SpeechRecognitionDataset()`
+ for the content in it.
+ 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.
+ warmup: a floating point value which increases throughout training;
+ values >= 1.0 are fully warmed up and have all modules present.
+ """
+ device = model.device if isinstance(model, DDP) else next(model.parameters()).device
+ feature = batch["inputs"]
+ # at entry, feature is (N, T, C)
+ assert feature.ndim == 3
+ feature = feature.to(device)
+
+ supervisions = batch["supervisions"]
+ feature_lens = supervisions["num_frames"].to(device)
+
+ batch_idx_train = params.batch_idx_train
+ warm_step = params.warm_step
+
+ texts = batch["supervisions"]["text"]
+ y = sp.encode(texts, out_type=int)
+ y = k2.RaggedTensor(y).to(device)
+
+ with torch.set_grad_enabled(is_training):
+ simple_loss, pruned_loss = model(
+ x=feature,
+ x_lens=feature_lens,
+ y=y,
+ prune_range=params.prune_range,
+ am_scale=params.am_scale,
+ lm_scale=params.lm_scale,
+ )
+
+ s = params.simple_loss_scale
+ # take down the scale on the simple loss from 1.0 at the start
+ # to params.simple_loss scale by warm_step.
+ simple_loss_scale = (
+ s
+ if batch_idx_train >= warm_step
+ else 1.0 - (batch_idx_train / warm_step) * (1.0 - s)
+ )
+ pruned_loss_scale = (
+ 1.0
+ if batch_idx_train >= warm_step
+ else 0.1 + 0.9 * (batch_idx_train / warm_step)
+ )
+
+ loss = simple_loss_scale * simple_loss + pruned_loss_scale * pruned_loss
+
+ assert loss.requires_grad == is_training
+
+ info = MetricsTracker()
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ info["frames"] = (feature_lens // params.subsampling_factor).sum().item()
+
+ # Note: We use reduction=sum while computing the loss.
+ info["loss"] = loss.detach().cpu().item()
+ info["simple_loss"] = simple_loss.detach().cpu().item()
+ info["pruned_loss"] = pruned_loss.detach().cpu().item()
+
+ return loss, info
+
+
+def compute_validation_loss(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ sp: spm.SentencePieceProcessor,
+ valid_dl: torch.utils.data.DataLoader,
+ world_size: int = 1,
+) -> MetricsTracker:
+ """Run the validation process."""
+ model.eval()
+
+ tot_loss = MetricsTracker()
+
+ for batch_idx, batch in enumerate(valid_dl):
+ loss, loss_info = compute_loss(
+ params=params,
+ model=model,
+ sp=sp,
+ batch=batch,
+ is_training=False,
+ )
+ assert loss.requires_grad is False
+ tot_loss = tot_loss + loss_info
+
+ if world_size > 1:
+ tot_loss.reduce(loss.device)
+
+ loss_value = tot_loss["loss"] / tot_loss["frames"]
+ if loss_value < params.best_valid_loss:
+ params.best_valid_epoch = params.cur_epoch
+ params.best_valid_loss = loss_value
+
+ return tot_loss
+
+
+def train_one_epoch(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ optimizer: torch.optim.Optimizer,
+ scheduler: LRSchedulerType,
+ sp: spm.SentencePieceProcessor,
+ train_dl: torch.utils.data.DataLoader,
+ valid_dl: torch.utils.data.DataLoader,
+ scaler: GradScaler,
+ model_avg: Optional[nn.Module] = None,
+ tb_writer: Optional[SummaryWriter] = None,
+ world_size: int = 1,
+ rank: int = 0,
+) -> 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.
+ scheduler:
+ The learning rate scheduler, we call step() every step.
+ train_dl:
+ Dataloader for the training dataset.
+ valid_dl:
+ Dataloader for the validation dataset.
+ scaler:
+ The scaler used for mix precision training.
+ model_avg:
+ The stored model averaged from the start of training.
+ tb_writer:
+ Writer to write log messages to tensorboard.
+ world_size:
+ Number of nodes in DDP training. If it is 1, DDP is disabled.
+ rank:
+ The rank of the node in DDP training. If no DDP is used, it should
+ be set to 0.
+ """
+ model.train()
+
+ tot_loss = MetricsTracker()
+
+ cur_batch_idx = params.get("cur_batch_idx", 0)
+
+ for batch_idx, batch in enumerate(train_dl):
+ if batch_idx < cur_batch_idx:
+ continue
+ cur_batch_idx = batch_idx
+
+ params.batch_idx_train += 1
+ batch_size = len(batch["supervisions"]["text"])
+
+ try:
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, loss_info = compute_loss(
+ params=params,
+ model=model,
+ sp=sp,
+ batch=batch,
+ is_training=True,
+ )
+ # summary stats
+ tot_loss = (tot_loss * (1 - 1 / params.reset_interval)) + loss_info
+
+ # NOTE: We use reduction==sum and loss is computed over utterances
+ # in the batch and there is no normalization to it so far.
+ scaler.scale(loss).backward()
+ set_batch_count(model, params.batch_idx_train)
+ scheduler.step_batch(params.batch_idx_train)
+
+ scaler.step(optimizer)
+ scaler.update()
+ optimizer.zero_grad()
+ except: # noqa
+ display_and_save_batch(batch, params=params, sp=sp)
+ raise
+
+ if params.print_diagnostics and batch_idx == 5:
+ return
+
+ if (
+ rank == 0
+ and params.batch_idx_train > 0
+ and params.batch_idx_train % params.average_period == 0
+ ):
+ update_averaged_model(
+ params=params,
+ model_cur=model,
+ model_avg=model_avg,
+ )
+
+ if (
+ params.batch_idx_train > 0
+ and params.batch_idx_train % params.save_every_n == 0
+ ):
+ params.cur_batch_idx = batch_idx
+ save_checkpoint_with_global_batch_idx(
+ out_dir=params.exp_dir,
+ global_batch_idx=params.batch_idx_train,
+ model=model,
+ model_avg=model_avg,
+ params=params,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=train_dl.sampler,
+ scaler=scaler,
+ rank=rank,
+ )
+ del params.cur_batch_idx
+ remove_checkpoints(
+ out_dir=params.exp_dir,
+ topk=params.keep_last_k,
+ rank=rank,
+ )
+
+ if batch_idx % 100 == 0 and params.use_fp16:
+ # If the grad scale was less than 1, try increasing it. The _growth_interval
+ # of the grad scaler is configurable, but we can't configure it to have different
+ # behavior depending on the current grad scale.
+ cur_grad_scale = scaler._scale.item()
+ if cur_grad_scale < 1.0 or (cur_grad_scale < 8.0 and batch_idx % 400 == 0):
+ scaler.update(cur_grad_scale * 2.0)
+ if cur_grad_scale < 0.01:
+ logging.warning(f"Grad scale is small: {cur_grad_scale}")
+ if cur_grad_scale < 1.0e-05:
+ raise RuntimeError(
+ f"grad_scale is too small, exiting: {cur_grad_scale}"
+ )
+
+ if batch_idx % params.log_interval == 0:
+ cur_lr = scheduler.get_last_lr()[0]
+ cur_grad_scale = scaler._scale.item() if params.use_fp16 else 1.0
+
+ logging.info(
+ f"Epoch {params.cur_epoch}, "
+ f"batch {batch_idx}, loss[{loss_info}], "
+ f"tot_loss[{tot_loss}], batch size: {batch_size}, "
+ f"lr: {cur_lr:.2e}, "
+ + (f"grad_scale: {scaler._scale.item()}" if params.use_fp16 else "")
+ )
+
+ if tb_writer is not None:
+ tb_writer.add_scalar(
+ "train/learning_rate", cur_lr, params.batch_idx_train
+ )
+
+ loss_info.write_summary(
+ tb_writer, "train/current_", params.batch_idx_train
+ )
+ tot_loss.write_summary(tb_writer, "train/tot_", params.batch_idx_train)
+ if params.use_fp16:
+ tb_writer.add_scalar(
+ "train/grad_scale",
+ cur_grad_scale,
+ params.batch_idx_train,
+ )
+
+ if batch_idx % params.valid_interval == 0 and not params.print_diagnostics:
+ logging.info("Computing validation loss")
+ valid_info = compute_validation_loss(
+ params=params,
+ model=model,
+ sp=sp,
+ valid_dl=valid_dl,
+ world_size=world_size,
+ )
+ model.train()
+ logging.info(f"Epoch {params.cur_epoch}, validation: {valid_info}")
+ logging.info(
+ f"Maximum memory allocated so far is {torch.cuda.max_memory_allocated()//1000000}MB"
+ )
+ if tb_writer is not None:
+ valid_info.write_summary(
+ tb_writer, "train/valid_", params.batch_idx_train
+ )
+
+ loss_value = tot_loss["loss"] / tot_loss["frames"]
+ params.train_loss = loss_value
+ 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))
+ if params.full_libri is False:
+ params.valid_interval = 1600
+
+ fix_random_seed(params.seed)
+ 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")
+
+ if args.tensorboard and rank == 0:
+ tb_writer = SummaryWriter(log_dir=f"{params.exp_dir}/tensorboard")
+ else:
+ tb_writer = None
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", rank)
+ logging.info(f"Device: {device}")
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(params.bpe_model)
+
+ # is defined in local/train_bpe_model.py
+ params.blank_id = sp.piece_to_id("")
+ params.vocab_size = sp.get_piece_size()
+
+ logging.info(params)
+
+ logging.info("About to create model")
+ model = get_transducer_model(params)
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ assert params.save_every_n >= params.average_period
+ model_avg: Optional[nn.Module] = None
+ if rank == 0:
+ # model_avg is only used with rank 0
+ model_avg = copy.deepcopy(model).to(torch.float64)
+
+ assert params.start_epoch > 0, params.start_epoch
+ checkpoints = load_checkpoint_if_available(
+ params=params, model=model, model_avg=model_avg
+ )
+
+ model.to(device)
+ if world_size > 1:
+ logging.info("Using DDP")
+ model = DDP(model, device_ids=[rank], find_unused_parameters=True)
+
+ parameters_names = []
+ parameters_names.append(
+ [name_param_pair[0] for name_param_pair in model.named_parameters()]
+ )
+ optimizer = ScaledAdam(
+ model.parameters(),
+ lr=params.base_lr,
+ clipping_scale=2.0,
+ parameters_names=parameters_names,
+ )
+
+ scheduler = Eden(optimizer, params.lr_batches, params.lr_epochs)
+
+ if checkpoints and "optimizer" in checkpoints:
+ logging.info("Loading optimizer state dict")
+ optimizer.load_state_dict(checkpoints["optimizer"])
+
+ if (
+ checkpoints
+ and "scheduler" in checkpoints
+ and checkpoints["scheduler"] is not None
+ ):
+ logging.info("Loading scheduler state dict")
+ scheduler.load_state_dict(checkpoints["scheduler"])
+
+ if params.print_diagnostics:
+ opts = diagnostics.TensorDiagnosticOptions(
+ 2**22
+ ) # allow 4 megabytes per sub-module
+ diagnostic = diagnostics.attach_diagnostics(model, opts)
+
+ if params.inf_check:
+ register_inf_check_hooks(model)
+
+ librispeech = LibriSpeechAsrDataModule(args)
+
+ train_cuts = librispeech.train_clean_100_cuts()
+ if params.full_libri:
+ train_cuts += librispeech.train_clean_360_cuts()
+ train_cuts += librispeech.train_other_500_cuts()
+
+ def remove_short_and_long_utt(c: Cut):
+ # Keep only utterances with duration between 1 second and 20 seconds
+ #
+ # Caution: There is a reason to select 20.0 here. Please see
+ # ../local/display_manifest_statistics.py
+ #
+ # You should use ../local/display_manifest_statistics.py to get
+ # an utterance duration distribution for your dataset to select
+ # the threshold
+ if c.duration < 1.0 or c.duration > 20.0:
+ logging.warning(
+ f"Exclude cut with ID {c.id} from training. Duration: {c.duration}"
+ )
+ return False
+
+ # In pruned RNN-T, we require that T >= S
+ # where T is the number of feature frames after subsampling
+ # and S is the number of tokens in the utterance
+
+ # In ./zipformer.py, the conv module uses the following expression
+ # for subsampling
+ T = ((c.num_frames - 7) // 2 + 1) // 2
+ tokens = sp.encode(c.supervisions[0].text, out_type=str)
+
+ if T < len(tokens):
+ logging.warning(
+ f"Exclude cut with ID {c.id} from training. "
+ f"Number of frames (before subsampling): {c.num_frames}. "
+ f"Number of frames (after subsampling): {T}. "
+ f"Text: {c.supervisions[0].text}. "
+ f"Tokens: {tokens}. "
+ f"Number of tokens: {len(tokens)}"
+ )
+ return False
+
+ return True
+
+ train_cuts = train_cuts.filter(remove_short_and_long_utt)
+
+ if params.start_batch > 0 and checkpoints and "sampler" in checkpoints:
+ # We only load the sampler's state dict when it loads a checkpoint
+ # saved in the middle of an epoch
+ sampler_state_dict = checkpoints["sampler"]
+ else:
+ sampler_state_dict = None
+
+ train_dl = librispeech.train_dataloaders(
+ train_cuts, sampler_state_dict=sampler_state_dict
+ )
+
+ valid_cuts = librispeech.dev_clean_cuts()
+ valid_cuts += librispeech.dev_other_cuts()
+ valid_dl = librispeech.valid_dataloaders(valid_cuts)
+
+ if not params.print_diagnostics:
+ scan_pessimistic_batches_for_oom(
+ model=model,
+ train_dl=train_dl,
+ optimizer=optimizer,
+ sp=sp,
+ params=params,
+ )
+
+ scaler = GradScaler(enabled=params.use_fp16, init_scale=1.0)
+ if checkpoints and "grad_scaler" in checkpoints:
+ logging.info("Loading grad scaler state dict")
+ scaler.load_state_dict(checkpoints["grad_scaler"])
+
+ for epoch in range(params.start_epoch, params.num_epochs + 1):
+ scheduler.step_epoch(epoch - 1)
+ fix_random_seed(params.seed + epoch - 1)
+ train_dl.sampler.set_epoch(epoch - 1)
+
+ if tb_writer is not None:
+ tb_writer.add_scalar("train/epoch", epoch, params.batch_idx_train)
+
+ params.cur_epoch = epoch
+
+ train_one_epoch(
+ params=params,
+ model=model,
+ model_avg=model_avg,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sp=sp,
+ train_dl=train_dl,
+ valid_dl=valid_dl,
+ scaler=scaler,
+ tb_writer=tb_writer,
+ world_size=world_size,
+ rank=rank,
+ )
+
+ if params.print_diagnostics:
+ diagnostic.print_diagnostics()
+ break
+
+ save_checkpoint(
+ params=params,
+ model=model,
+ model_avg=model_avg,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=train_dl.sampler,
+ scaler=scaler,
+ rank=rank,
+ )
+
+ logging.info("Done!")
+
+ if world_size > 1:
+ torch.distributed.barrier()
+ cleanup_dist()
+
+
+def display_and_save_batch(
+ batch: dict,
+ params: AttributeDict,
+ sp: spm.SentencePieceProcessor,
+) -> None:
+ """Display the batch statistics and save the batch into disk.
+
+ Args:
+ batch:
+ A batch of data. See `lhotse.dataset.K2SpeechRecognitionDataset()`
+ for the content in it.
+ params:
+ Parameters for training. See :func:`get_params`.
+ sp:
+ The BPE model.
+ """
+ from lhotse.utils import uuid4
+
+ filename = f"{params.exp_dir}/batch-{uuid4()}.pt"
+ logging.info(f"Saving batch to {filename}")
+ torch.save(batch, filename)
+
+ supervisions = batch["supervisions"]
+ features = batch["inputs"]
+
+ logging.info(f"features shape: {features.shape}")
+
+ y = sp.encode(supervisions["text"], out_type=int)
+ num_tokens = sum(len(i) for i in y)
+ logging.info(f"num tokens: {num_tokens}")
+
+
+def scan_pessimistic_batches_for_oom(
+ model: Union[nn.Module, DDP],
+ train_dl: torch.utils.data.DataLoader,
+ optimizer: torch.optim.Optimizer,
+ sp: spm.SentencePieceProcessor,
+ params: AttributeDict,
+):
+ from lhotse.dataset import find_pessimistic_batches
+
+ logging.info(
+ "Sanity check -- see if any of the batches in epoch 1 would cause OOM."
+ )
+ batches, crit_values = find_pessimistic_batches(train_dl.sampler)
+ for criterion, cuts in batches.items():
+ batch = train_dl.dataset[cuts]
+ try:
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, _ = compute_loss(
+ params=params,
+ model=model,
+ sp=sp,
+ batch=batch,
+ is_training=True,
+ )
+ loss.backward()
+ optimizer.zero_grad()
+ except Exception as e:
+ if "CUDA out of memory" in str(e):
+ logging.error(
+ "Your GPU ran out of memory with the current "
+ "max_duration setting. We recommend decreasing "
+ "max_duration and trying again.\n"
+ f"Failing criterion: {criterion} "
+ f"(={crit_values[criterion]}) ..."
+ )
+ display_and_save_batch(batch, params=params, sp=sp)
+ raise
+ logging.info(
+ f"Maximum memory allocated so far is {torch.cuda.max_memory_allocated()//1000000}MB"
+ )
+
+
+def main():
+ parser = get_parser()
+ LibriSpeechAsrDataModule.add_arguments(parser)
+ args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ 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)
+
+
+torch.set_num_threads(1)
+torch.set_num_interop_threads(1)
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/zipformer.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/zipformer.py
new file mode 100644
index 000000000..88beb38c1
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/zipformer.py
@@ -0,0 +1,2881 @@
+#!/usr/bin/env python3
+# Copyright 2022 Xiaomi Corp. (authors: Daniel Povey,)
+# Zengwei Yao)
+#
+# 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.
+
+import copy
+import itertools
+import logging
+import math
+import random
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+from encoder_interface import EncoderInterface
+from scaling import (
+ ScaledLinear, # not as in other dirs.. just scales down initial parameter values.
+)
+from scaling import (
+ ActivationBalancer,
+ BasicNorm,
+ DoubleSwish,
+ Identity,
+ MaxEig,
+ ScaledConv1d,
+ Whiten,
+ _diag,
+ penalize_abs_values_gt,
+ random_clamp,
+ softmax,
+)
+from torch import Tensor, nn
+
+from icefall.dist import get_rank
+from icefall.utils import make_pad_mask, subsequent_chunk_mask
+
+
+def stack_states(state_list: List[List[Tensor]]) -> List[Tensor]:
+ """Stack list of zipformer states that correspond to separate utterances
+ into a single emformer state, so that it can be used as an input for
+ zipformer when those utterances are formed into a batch.
+
+ Note:
+ It is the inverse of :func:`unstack_states`.
+
+ Args:
+ state_list:
+ Each element in state_list corresponding to the internal state
+ of the zipformer model for a single utterance.
+ ``states[i]`` is a list of 7 * num_encoders elements of i-th utterance.
+ ``states[i][0:num_encoders]`` is the cached numbers of past frames.
+ ``states[i][num_encoders:2*num_encoders]`` is the cached average tensors.
+ ``states[i][2*num_encoders:3*num_encoders]`` is the cached key tensors of the first attention modules.
+ ``states[i][3*num_encoders:4*num_encoders]`` is the cached value tensors of the first attention modules.
+ ``states[i][4*num_encoders:5*num_encoders]`` is the cached value tensors of the second attention modules.
+ ``states[i][5*num_encoders:6*num_encoders]`` is the cached left contexts of the first convolution modules.
+ ``states[i][6*num_encoders:7*num_encoders]`` is the cached left contexts of the second convolution modules.
+
+ Returns:
+ A new state corresponding to a batch of utterances.
+ See the input argument of :func:`unstack_states` for the meaning
+ of the returned tensor.
+ """
+ batch_size = len(state_list)
+ assert len(state_list[0]) % 7 == 0, len(state_list[0])
+ num_encoders = len(state_list[0]) // 7
+
+ cached_len = []
+ cached_avg = []
+ cached_key = []
+ cached_val = []
+ cached_val2 = []
+ cached_conv1 = []
+ cached_conv2 = []
+
+ # For cached_len
+ len_list = [state_list[n][0:num_encoders] for n in range(batch_size)]
+ for i in range(num_encoders):
+ # len_avg: (num_layers, batch_size)
+ len_avg = torch.cat([len_list[n][i] for n in range(batch_size)], dim=1)
+ cached_len.append(len_avg)
+
+ # For cached_avg
+ avg_list = [
+ state_list[n][num_encoders : 2 * num_encoders] for n in range(batch_size)
+ ]
+ for i in range(num_encoders):
+ # avg: (num_layers, batch_size, D)
+ avg = torch.cat([avg_list[n][i] for n in range(batch_size)], dim=1)
+ cached_avg.append(avg)
+
+ # For cached_key
+ key_list = [
+ state_list[n][2 * num_encoders : 3 * num_encoders] for n in range(batch_size)
+ ]
+ for i in range(num_encoders):
+ # key: (num_layers, left_context_size, batch_size, D)
+ key = torch.cat([key_list[n][i] for n in range(batch_size)], dim=2)
+ cached_key.append(key)
+
+ # For cached_val
+ val_list = [
+ state_list[n][3 * num_encoders : 4 * num_encoders] for n in range(batch_size)
+ ]
+ for i in range(num_encoders):
+ # val: (num_layers, left_context_size, batch_size, D)
+ val = torch.cat([val_list[n][i] for n in range(batch_size)], dim=2)
+ cached_val.append(val)
+
+ # For cached_val2
+ val2_list = [
+ state_list[n][4 * num_encoders : 5 * num_encoders] for n in range(batch_size)
+ ]
+ for i in range(num_encoders):
+ # val2: (num_layers, left_context_size, batch_size, D)
+ val2 = torch.cat([val2_list[n][i] for n in range(batch_size)], dim=2)
+ cached_val2.append(val2)
+
+ # For cached_conv1
+ conv1_list = [
+ state_list[n][5 * num_encoders : 6 * num_encoders] for n in range(batch_size)
+ ]
+ for i in range(num_encoders):
+ # conv1: (num_layers, batch_size, D, kernel-1)
+ conv1 = torch.cat([conv1_list[n][i] for n in range(batch_size)], dim=1)
+ cached_conv1.append(conv1)
+
+ # For cached_conv2
+ conv2_list = [
+ state_list[n][6 * num_encoders : 7 * num_encoders] for n in range(batch_size)
+ ]
+ for i in range(num_encoders):
+ # conv2: (num_layers, batch_size, D, kernel-1)
+ conv2 = torch.cat([conv2_list[n][i] for n in range(batch_size)], dim=1)
+ cached_conv2.append(conv2)
+
+ states = (
+ cached_len
+ + cached_avg
+ + cached_key
+ + cached_val
+ + cached_val2
+ + cached_conv1
+ + cached_conv2
+ )
+ return states
+
+
+def unstack_states(states: List[Tensor]) -> List[List[Tensor]]:
+ """Unstack the zipformer state corresponding to a batch of utterances
+ into a list of states, where the i-th entry is the state from the i-th
+ utterance in the batch.
+
+ Note:
+ It is the inverse of :func:`stack_states`.
+
+ Args:
+ states:
+ A list of 7 * num_encoders elements:
+ ``states[0:num_encoders]`` is the cached numbers of past frames.
+ ``states[num_encoders:2*num_encoders]`` is the cached average tensors.
+ ``states[2*num_encoders:3*num_encoders]`` is the cached key tensors of the first attention modules.
+ ``states[3*num_encoders:4*num_encoders]`` is the cached value tensors of the first attention modules.
+ ``states[4*num_encoders:5*num_encoders]`` is the cached value tensors of the second attention modules.
+ ``states[5*num_encoders:6*num_encoders]`` is the cached left contexts of the first convolution modules.
+ ``states[6*num_encoders:7*num_encoders]`` is the cached left contexts of the second convolution modules.
+
+ Returns:
+ A list of states.
+ ``states[i]`` is a list of 7 * num_encoders elements of i-th utterance.
+ """
+ assert len(states) % 7 == 0, len(states)
+ num_encoders = len(states) // 7
+ (
+ cached_len,
+ cached_avg,
+ cached_key,
+ cached_val,
+ cached_val2,
+ cached_conv1,
+ cached_conv2,
+ ) = (states[i * num_encoders : (i + 1) * num_encoders] for i in range(7))
+
+ batch_size = cached_len[0].shape[1]
+
+ len_list = [[] for _ in range(batch_size)]
+ for i in range(num_encoders):
+ # cached_len[i]: (num_layers, batch_size)
+ len_avg = cached_len[i].chunk(chunks=batch_size, dim=1)
+ for n in range(batch_size):
+ len_list[n].append(len_avg[n])
+
+ avg_list = [[] for _ in range(batch_size)]
+ for i in range(num_encoders):
+ # cached_avg[i]: (num_layers, batch_size, D)
+ avg = cached_avg[i].chunk(chunks=batch_size, dim=1)
+ for n in range(batch_size):
+ avg_list[n].append(avg[n])
+
+ key_list = [[] for _ in range(batch_size)]
+ for i in range(num_encoders):
+ # cached_key[i]: (num_layers, left_context, batch_size, D)
+ key = cached_key[i].chunk(chunks=batch_size, dim=2)
+ for n in range(batch_size):
+ key_list[n].append(key[n])
+
+ val_list = [[] for _ in range(batch_size)]
+ for i in range(num_encoders):
+ # cached_val[i]: (num_layers, left_context, batch_size, D)
+ val = cached_val[i].chunk(chunks=batch_size, dim=2)
+ for n in range(batch_size):
+ val_list[n].append(val[n])
+
+ val2_list = [[] for _ in range(batch_size)]
+ for i in range(num_encoders):
+ # cached_val2[i]: (num_layers, left_context, batch_size, D)
+ val2 = cached_val2[i].chunk(chunks=batch_size, dim=2)
+ for n in range(batch_size):
+ val2_list[n].append(val2[n])
+
+ conv1_list = [[] for _ in range(batch_size)]
+ for i in range(num_encoders):
+ # cached_conv1[i]: (num_layers, batch_size, D, kernel-1)
+ conv1 = cached_conv1[i].chunk(chunks=batch_size, dim=1)
+ for n in range(batch_size):
+ conv1_list[n].append(conv1[n])
+
+ conv2_list = [[] for _ in range(batch_size)]
+ for i in range(num_encoders):
+ # cached_conv2[i]: (num_layers, batch_size, D, kernel-1)
+ conv2 = cached_conv2[i].chunk(chunks=batch_size, dim=1)
+ for n in range(batch_size):
+ conv2_list[n].append(conv2[n])
+
+ state_list = [
+ (
+ len_list[i]
+ + avg_list[i]
+ + key_list[i]
+ + val_list[i]
+ + val2_list[i]
+ + conv1_list[i]
+ + conv2_list[i]
+ )
+ for i in range(batch_size)
+ ]
+ return state_list
+
+
+class Zipformer(EncoderInterface):
+ """
+ Args:
+ num_features (int): Number of input features
+ d_model: (int,int): embedding dimension of 2 encoder stacks
+ attention_dim: (int,int): attention dimension of 2 encoder stacks
+ nhead (int, int): number of heads
+ dim_feedforward (int, int): feedforward dimension in 2 encoder stacks
+ num_encoder_layers (int): number of encoder layers
+ dropout (float): dropout rate
+ cnn_module_kernel (int): Kernel size of convolution module
+ vgg_frontend (bool): whether to use vgg frontend.
+ warmup_batches (float): number of batches to warm up over
+ """
+
+ def __init__(
+ self,
+ num_features: int,
+ output_downsampling_factor: int = 2,
+ encoder_dims: Tuple[int] = (384, 384),
+ attention_dim: Tuple[int] = (256, 256),
+ encoder_unmasked_dims: Tuple[int] = (256, 256),
+ zipformer_downsampling_factors: Tuple[int] = (2, 4),
+ nhead: Tuple[int] = (8, 8),
+ feedforward_dim: Tuple[int] = (1536, 2048),
+ num_encoder_layers: Tuple[int] = (12, 12),
+ dropout: float = 0.1,
+ cnn_module_kernels: Tuple[int] = (31, 31),
+ pos_dim: int = 4,
+ num_left_chunks: int = 4,
+ short_chunk_threshold: float = 0.75,
+ short_chunk_size: int = 50,
+ decode_chunk_size: int = 16,
+ warmup_batches: float = 4000.0,
+ ) -> None:
+ super(Zipformer, self).__init__()
+
+ self.num_features = num_features
+ assert 0 < encoder_dims[0] <= encoder_dims[1]
+ self.encoder_dims = encoder_dims
+ self.encoder_unmasked_dims = encoder_unmasked_dims
+ self.zipformer_downsampling_factors = zipformer_downsampling_factors
+ self.output_downsampling_factor = output_downsampling_factor
+
+ self.num_left_chunks = num_left_chunks
+ self.short_chunk_threshold = short_chunk_threshold
+ self.short_chunk_size = short_chunk_size
+
+ # Used in decoding
+ self.decode_chunk_size = decode_chunk_size
+
+ # will be written to, see set_batch_count()
+ self.batch_count = 0
+ self.warmup_end = warmup_batches
+
+ for u, d in zip(encoder_unmasked_dims, encoder_dims):
+ assert u <= d, (u, d)
+
+ # self.encoder_embed converts the input of shape (N, T, num_features)
+ # to the shape (N, (T - 7)//2, encoder_dims).
+ # That is, it does two things simultaneously:
+ # (1) subsampling: T -> (T - 7)//2
+ # (2) embedding: num_features -> encoder_dims
+ self.encoder_embed = Conv2dSubsampling(
+ num_features, encoder_dims[0], dropout=dropout
+ )
+
+ # each one will be ZipformerEncoder or DownsampledZipformerEncoder
+ encoders = []
+
+ self.num_encoders = len(encoder_dims)
+ for i in range(self.num_encoders):
+ encoder_layer = ZipformerEncoderLayer(
+ encoder_dims[i],
+ attention_dim[i],
+ nhead[i],
+ feedforward_dim[i],
+ dropout,
+ cnn_module_kernels[i],
+ pos_dim,
+ )
+
+ # For the segment of the warmup period, we let the Conv2dSubsampling
+ # layer learn something. Then we start to warm up the other encoders.
+ encoder = ZipformerEncoder(
+ encoder_layer,
+ num_encoder_layers[i],
+ dropout,
+ warmup_begin=warmup_batches * (i + 1) / (self.num_encoders + 1),
+ warmup_end=warmup_batches * (i + 2) / (self.num_encoders + 1),
+ )
+
+ if zipformer_downsampling_factors[i] != 1:
+ encoder = DownsampledZipformerEncoder(
+ encoder,
+ input_dim=encoder_dims[i - 1] if i > 0 else encoder_dims[0],
+ output_dim=encoder_dims[i],
+ downsample=zipformer_downsampling_factors[i],
+ )
+ encoders.append(encoder)
+ self.encoders = nn.ModuleList(encoders)
+
+ # initializes self.skip_layers and self.skip_modules
+ self._init_skip_modules()
+
+ self.downsample_output = AttentionDownsample(
+ encoder_dims[-1], encoder_dims[-1], downsample=output_downsampling_factor
+ )
+
+ def _get_layer_skip_dropout_prob(self):
+ if not self.training:
+ return 0.0
+ batch_count = self.batch_count
+ min_dropout_prob = 0.025
+
+ if batch_count > self.warmup_end:
+ return min_dropout_prob
+ else:
+ return 0.5 - (batch_count / self.warmup_end) * (0.5 - min_dropout_prob)
+
+ def _init_skip_modules(self):
+ """
+ If self.zipformer_downampling_factors = (1, 2, 4, 8, 4, 2), then at the input of layer
+ indexed 4 (in zero indexing), with has subsapling_factor=4, we combine the output of
+ layers 2 and 3; and at the input of layer indexed 5, which which has subsampling_factor=2,
+ we combine the outputs of layers 1 and 5.
+ """
+ skip_layers = []
+ skip_modules = []
+ z = self.zipformer_downsampling_factors
+ for i in range(len(z)):
+ if i <= 1 or z[i - 1] <= z[i]:
+ skip_layers.append(None)
+ skip_modules.append(SimpleCombinerIdentity())
+ else:
+ # TEMP
+ for j in range(i - 2, -1, -1):
+ if z[j] <= z[i] or j == 0:
+ # TEMP logging statement.
+ logging.info(
+ f"At encoder stack {i}, which has downsampling_factor={z[i]}, we will "
+ f"combine the outputs of layers {j} and {i-1}, with downsampling_factors={z[j]} and {z[i-1]}."
+ )
+ skip_layers.append(j)
+ skip_modules.append(
+ SimpleCombiner(
+ self.encoder_dims[j],
+ self.encoder_dims[i - 1],
+ min_weight=(0.0, 0.25),
+ )
+ )
+ break
+ self.skip_layers = skip_layers
+ self.skip_modules = nn.ModuleList(skip_modules)
+
+ def get_feature_masks(self, x: torch.Tensor) -> List[float]:
+ # Note: The actual return type is Union[List[float], List[Tensor]],
+ # but to make torch.jit.script() work, we use List[float]
+ """
+ In eval mode, returns [1.0] * num_encoders; in training mode, returns a number of
+ randomized feature masks, one per encoder.
+ On e.g. 15% of frames, these masks will zero out all enocder dims larger than
+ some supplied number, e.g. >256, so in effect on those frames we are using
+ a smaller encoer dim.
+
+ We generate the random masks at this level because we want the 2 masks to 'agree'
+ all the way up the encoder stack. This will mean that the 1st mask will have
+ mask values repeated self.zipformer_subsampling_factor times.
+
+ Args:
+ x: the embeddings (needed for the shape and dtype and device), of shape
+ (num_frames, batch_size, encoder_dims0)
+ """
+ num_encoders = len(self.encoder_dims)
+ if torch.jit.is_scripting() or not self.training:
+ return [1.0] * num_encoders
+
+ (num_frames0, batch_size, _encoder_dims0) = x.shape
+
+ assert self.encoder_dims[0] == _encoder_dims0, (
+ self.encoder_dims,
+ _encoder_dims0,
+ )
+
+ max_downsampling_factor = max(self.zipformer_downsampling_factors)
+
+ num_frames_max = num_frames0 + max_downsampling_factor - 1
+
+ feature_mask_dropout_prob = 0.15
+
+ # frame_mask_max shape: (num_frames_max, batch_size, 1)
+ frame_mask_max = (
+ torch.rand(num_frames_max, batch_size, 1, device=x.device)
+ > feature_mask_dropout_prob
+ ).to(x.dtype)
+
+ feature_masks = []
+ for i in range(num_encoders):
+ ds = self.zipformer_downsampling_factors[i]
+ upsample_factor = max_downsampling_factor // ds
+
+ frame_mask = (
+ frame_mask_max.unsqueeze(1)
+ .expand(num_frames_max, upsample_factor, batch_size, 1)
+ .reshape(num_frames_max * upsample_factor, batch_size, 1)
+ )
+ num_frames = (num_frames0 + ds - 1) // ds
+ frame_mask = frame_mask[:num_frames]
+ feature_mask = torch.ones(
+ num_frames,
+ batch_size,
+ self.encoder_dims[i],
+ dtype=x.dtype,
+ device=x.device,
+ )
+ u = self.encoder_unmasked_dims[i]
+ feature_mask[:, :, u:] *= frame_mask
+ feature_masks.append(feature_mask)
+
+ return feature_masks
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ x_lens: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Args:
+ x:
+ The input tensor. Its shape is (batch_size, seq_len, feature_dim).
+ x_lens:
+ A tensor of shape (batch_size,) containing the number of frames in
+ `x` before padding.
+ chunk_size:
+ The chunk size used in evaluation mode.
+ Returns:
+ Return a tuple containing 2 tensors:
+ - embeddings: its shape is (batch_size, output_seq_len, encoder_dims[-1])
+ - lengths, a tensor of shape (batch_size,) containing the number
+ of frames in `embeddings` before padding.
+ """
+ x = self.encoder_embed(x)
+
+ x = x.permute(1, 0, 2) # (N, T, C) -> (T, N, C)
+
+ lengths = (x_lens - 7) >> 1
+ assert x.size(0) == lengths.max().item(), (x.shape, lengths, lengths.max())
+ mask = make_pad_mask(lengths)
+
+ outputs = []
+ feature_masks = self.get_feature_masks(x)
+
+ if self.training:
+ # Training mode
+ max_ds = max(self.zipformer_downsampling_factors)
+ # Generate dynamic chunk-wise attention mask during training
+ max_len = x.size(0) // max_ds
+ short_chunk_size = self.short_chunk_size // max_ds
+ chunk_size = torch.randint(1, max_len, (1,)).item()
+ if chunk_size > (max_len * self.short_chunk_threshold):
+ # Full attention
+ chunk_size = x.size(0)
+ else:
+ # Chunk-wise attention
+ chunk_size = chunk_size % short_chunk_size + 1
+ chunk_size *= max_ds
+ else:
+ chunk_size = self.decode_chunk_size
+ # Evaluation mode
+ for ds in self.zipformer_downsampling_factors:
+ assert chunk_size % ds == 0, (chunk_size, ds)
+
+ attn_mask = ~subsequent_chunk_mask(
+ size=x.size(0),
+ chunk_size=chunk_size,
+ num_left_chunks=self.num_left_chunks,
+ device=x.device,
+ )
+
+ for i, (module, skip_module) in enumerate(
+ zip(self.encoders, self.skip_modules)
+ ):
+ ds = self.zipformer_downsampling_factors[i]
+ k = self.skip_layers[i]
+ if isinstance(k, int):
+ layer_skip_dropout_prob = self._get_layer_skip_dropout_prob()
+ if torch.jit.is_scripting():
+ x = skip_module(outputs[k], x)
+ elif (not self.training) or random.random() > layer_skip_dropout_prob:
+ x = skip_module(outputs[k], x)
+ x = module(
+ x,
+ feature_mask=feature_masks[i],
+ src_key_padding_mask=None if mask is None else mask[..., ::ds],
+ attn_mask=attn_mask[::ds, ::ds],
+ )
+ outputs.append(x)
+
+ x = self.downsample_output(x)
+ # class Downsample has this rounding behavior..
+ assert self.output_downsampling_factor == 2, self.output_downsampling_factor
+ lengths = (lengths + 1) >> 1
+
+ x = x.permute(1, 0, 2) # (T, N, C) ->(N, T, C)
+
+ return x, lengths
+
+ def streaming_forward(
+ self,
+ x: torch.Tensor,
+ x_lens: torch.Tensor,
+ states: List[Tensor],
+ ) -> Tuple[Tensor, Tensor, List[Tensor]]:
+ """
+ Args:
+ x:
+ The input tensor. Its shape is (batch_size, seq_len, feature_dim).
+ seq_len is the input chunk length.
+ x_lens:
+ A tensor of shape (batch_size,) containing the number of frames in
+ `x` before padding.
+ states:
+ A list of 7 * num_encoders elements:
+ ``states[0:num_encoders]`` is the cached numbers of past frames.
+ ``states[num_encoders:2*num_encoders]`` is the cached average tensors.
+ ``states[2*num_encoders:3*num_encoders]`` is the cached key tensors of the first attention modules.
+ ``states[3*num_encoders:4*num_encoders]`` is the cached value tensors of the first attention modules.
+ ``states[4*num_encoders:5*num_encoders]`` is the cached value tensors of the second attention modules.
+ ``states[5*num_encoders:6*num_encoders]`` is the cached left contexts of the first convolution modules.
+ ``states[6*num_encoders:7*num_encoders]`` is the cached left contexts of the second convolution modules.
+
+ Returns:
+ Return a tuple containing 3 tensors:
+ - embeddings: its shape is (batch_size, output_seq_len, encoder_dims[-1])
+ - lengths, a tensor of shape (batch_size,) containing the number
+ of frames in `embeddings` before padding.
+ - updated states.
+ """
+ assert len(states) == 7 * self.num_encoders, (len(states), self.num_encoders)
+
+ cached_len = states[: self.num_encoders]
+ cached_avg = states[self.num_encoders : 2 * self.num_encoders]
+ cached_key = states[2 * self.num_encoders : 3 * self.num_encoders]
+ cached_val = states[3 * self.num_encoders : 4 * self.num_encoders]
+ cached_val2 = states[4 * self.num_encoders : 5 * self.num_encoders]
+ cached_conv1 = states[5 * self.num_encoders : 6 * self.num_encoders]
+ cached_conv2 = states[6 * self.num_encoders : 7 * self.num_encoders]
+
+ x = self.encoder_embed(x)
+ x = x.permute(1, 0, 2) # (N, T, C) -> (T, N, C)
+ lengths = (x_lens - 7) >> 1
+ assert x.size(0) == lengths.max().item(), (x.shape, lengths, lengths.max())
+
+ outputs = []
+ new_cached_len = []
+ new_cached_avg = []
+ new_cached_key = []
+ new_cached_val = []
+ new_cached_val2 = []
+ new_cached_conv1 = []
+ new_cached_conv2 = []
+
+ for i, (module, skip_module) in enumerate(
+ zip(self.encoders, self.skip_modules)
+ ):
+ k = self.skip_layers[i]
+ if isinstance(k, int):
+ x = skip_module(outputs[k], x)
+ x, len_avg, avg, key, val, val2, conv1, conv2 = module.streaming_forward(
+ x,
+ cached_len=cached_len[i],
+ cached_avg=cached_avg[i],
+ cached_key=cached_key[i],
+ cached_val=cached_val[i],
+ cached_val2=cached_val2[i],
+ cached_conv1=cached_conv1[i],
+ cached_conv2=cached_conv2[i],
+ )
+ outputs.append(x)
+ # Update caches
+ new_cached_len.append(len_avg)
+ new_cached_avg.append(avg)
+ new_cached_key.append(key)
+ new_cached_val.append(val)
+ new_cached_val2.append(val2)
+ new_cached_conv1.append(conv1)
+ new_cached_conv2.append(conv2)
+
+ x = self.downsample_output(x)
+ # class Downsample has this rounding behavior..
+ assert self.output_downsampling_factor == 2, self.output_downsampling_factor
+ lengths = (lengths + 1) >> 1
+
+ x = x.permute(1, 0, 2) # (T, N, C) ->(N, T, C)
+
+ new_states = (
+ new_cached_len
+ + new_cached_avg
+ + new_cached_key
+ + new_cached_val
+ + new_cached_val2
+ + new_cached_conv1
+ + new_cached_conv2
+ )
+ return x, lengths, new_states
+
+ @torch.jit.export
+ def get_init_state(
+ self,
+ device: torch.device = torch.device("cpu"),
+ ) -> List[Tensor]:
+ """Get initial states.
+ A list of 7 * num_encoders elements:
+ ``states[0:num_encoders]`` is the cached numbers of past frames.
+ ``states[num_encoders:2*num_encoders]`` is the cached average tensors.
+ ``states[2*num_encoders:3*num_encoders]`` is the cached key tensors of the first attention modules.
+ ``states[3*num_encoders:4*num_encoders]`` is the cached value tensors of the first attention modules.
+ ``states[4*num_encoders:5*num_encoders]`` is the cached value tensors of the second attention modules.
+ ``states[5*num_encoders:6*num_encoders]`` is the cached left contexts of the first convolution modules.
+ ``states[6*num_encoders:7*num_encoders]`` is the cached left contexts of the second convolution modules.
+ """
+ cached_len = []
+ cached_avg = []
+ cached_key = []
+ cached_val = []
+ cached_val2 = []
+ cached_conv1 = []
+ cached_conv2 = []
+
+ left_context_len = self.decode_chunk_size * self.num_left_chunks
+
+ for i, encoder in enumerate(self.encoders):
+ num_layers = encoder.num_layers
+ ds = self.zipformer_downsampling_factors[i]
+
+ len_avg = torch.zeros(num_layers, 1, dtype=torch.int32, device=device)
+ cached_len.append(len_avg)
+
+ avg = torch.zeros(num_layers, 1, encoder.d_model, device=device)
+ cached_avg.append(avg)
+
+ key = torch.zeros(
+ num_layers,
+ left_context_len // ds,
+ 1,
+ encoder.attention_dim,
+ device=device,
+ )
+ cached_key.append(key)
+
+ val = torch.zeros(
+ num_layers,
+ left_context_len // ds,
+ 1,
+ encoder.attention_dim // 2,
+ device=device,
+ )
+ cached_val.append(val)
+
+ val2 = torch.zeros(
+ num_layers,
+ left_context_len // ds,
+ 1,
+ encoder.attention_dim // 2,
+ device=device,
+ )
+ cached_val2.append(val2)
+
+ conv1 = torch.zeros(
+ num_layers,
+ 1,
+ encoder.d_model,
+ encoder.cnn_module_kernel - 1,
+ device=device,
+ )
+ cached_conv1.append(conv1)
+
+ conv2 = torch.zeros(
+ num_layers,
+ 1,
+ encoder.d_model,
+ encoder.cnn_module_kernel - 1,
+ device=device,
+ )
+ cached_conv2.append(conv2)
+
+ states = (
+ cached_len
+ + cached_avg
+ + cached_key
+ + cached_val
+ + cached_val2
+ + cached_conv1
+ + cached_conv2
+ )
+ return states
+
+
+class ZipformerEncoderLayer(nn.Module):
+ """
+ ZipformerEncoderLayer is made up of self-attn, feedforward and convolution networks.
+
+ Args:
+ d_model: the number of expected features in the input (required).
+ nhead: the number of heads in the multiheadattention models (required).
+ feedforward_dim: 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.
+
+ Examples::
+ >>> encoder_layer = ZipformerEncoderLayer(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,
+ attention_dim: int,
+ nhead: int,
+ feedforward_dim: int = 2048,
+ dropout: float = 0.1,
+ cnn_module_kernel: int = 31,
+ pos_dim: int = 4,
+ ) -> None:
+ super(ZipformerEncoderLayer, self).__init__()
+
+ self.d_model = d_model
+ self.attention_dim = attention_dim
+ self.cnn_module_kernel = cnn_module_kernel
+
+ # will be written to, see set_batch_count()
+ self.batch_count = 0
+
+ self.self_attn = RelPositionMultiheadAttention(
+ d_model,
+ attention_dim,
+ nhead,
+ pos_dim,
+ dropout=0.0,
+ )
+
+ self.pooling = PoolingModule(d_model)
+
+ self.feed_forward1 = FeedforwardModule(d_model, feedforward_dim, dropout)
+
+ self.feed_forward2 = FeedforwardModule(d_model, feedforward_dim, dropout)
+
+ self.feed_forward3 = FeedforwardModule(d_model, feedforward_dim, dropout)
+
+ self.conv_module1 = ConvolutionModule(d_model, cnn_module_kernel)
+
+ self.conv_module2 = ConvolutionModule(d_model, cnn_module_kernel)
+
+ self.norm_final = BasicNorm(d_model)
+
+ self.bypass_scale = nn.Parameter(torch.tensor(0.5))
+
+ # try to ensure the output is close to zero-mean (or at least, zero-median).
+ self.balancer = ActivationBalancer(
+ d_model,
+ channel_dim=-1,
+ min_positive=0.45,
+ max_positive=0.55,
+ max_abs=6.0,
+ )
+ self.whiten = Whiten(
+ num_groups=1, whitening_limit=5.0, prob=(0.025, 0.25), grad_scale=0.01
+ )
+
+ def get_bypass_scale(self):
+ if torch.jit.is_scripting() or not self.training:
+ return self.bypass_scale
+ if random.random() < 0.1:
+ # ensure we get grads if self.bypass_scale becomes out of range
+ return self.bypass_scale
+ # hardcode warmup period for bypass scale
+ warmup_period = 20000.0
+ initial_clamp_min = 0.75
+ final_clamp_min = 0.25
+ if self.batch_count > warmup_period:
+ clamp_min = final_clamp_min
+ else:
+ clamp_min = initial_clamp_min - (self.batch_count / warmup_period) * (
+ initial_clamp_min - final_clamp_min
+ )
+ return self.bypass_scale.clamp(min=clamp_min, max=1.0)
+
+ def get_dynamic_dropout_rate(self):
+ # return dropout rate for the dynamic modules (self_attn, pooling, convolution); this
+ # starts at 0.2 and rapidly decreases to 0. Its purpose is to keep the training stable
+ # at the beginning, by making the network focus on the feedforward modules.
+ if torch.jit.is_scripting() or not self.training:
+ return 0.0
+ warmup_period = 2000.0
+ initial_dropout_rate = 0.2
+ final_dropout_rate = 0.0
+ if self.batch_count > warmup_period:
+ return final_dropout_rate
+ else:
+ return initial_dropout_rate - (
+ initial_dropout_rate * final_dropout_rate
+ ) * (self.batch_count / warmup_period)
+
+ def forward(
+ self,
+ src: Tensor,
+ pos_emb: Tensor,
+ attn_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).
+ batch_split: if not None, this layer will only be applied to
+
+ 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
+ """
+ src_orig = src
+
+ # macaron style feed forward module
+ src = src + self.feed_forward1(src)
+
+ # dropout rate for submodules that interact with time.
+ dynamic_dropout = self.get_dynamic_dropout_rate()
+
+ # pooling module
+ if torch.jit.is_scripting():
+ src = src + self.pooling(src, src_key_padding_mask=src_key_padding_mask)
+ elif random.random() >= dynamic_dropout:
+ src = src + self.pooling(src, src_key_padding_mask=src_key_padding_mask)
+
+ if torch.jit.is_scripting():
+ src_att, attn_weights = self.self_attn(
+ src,
+ pos_emb=pos_emb,
+ attn_mask=attn_mask,
+ key_padding_mask=src_key_padding_mask,
+ )
+ src = src + src_att
+
+ src = src + self.conv_module1(
+ src, src_key_padding_mask=src_key_padding_mask
+ )
+
+ src = src + self.feed_forward2(src)
+
+ src = src + self.self_attn.forward2(src, attn_weights)
+
+ src = src + self.conv_module2(
+ src, src_key_padding_mask=src_key_padding_mask
+ )
+ else:
+ use_self_attn = random.random() >= dynamic_dropout
+ if use_self_attn:
+ src_att, attn_weights = self.self_attn(
+ src,
+ pos_emb=pos_emb,
+ attn_mask=attn_mask,
+ key_padding_mask=src_key_padding_mask,
+ )
+ src = src + src_att
+
+ if random.random() >= dynamic_dropout:
+ src = src + self.conv_module1(
+ src, src_key_padding_mask=src_key_padding_mask
+ )
+
+ src = src + self.feed_forward2(src)
+
+ if use_self_attn:
+ src = src + self.self_attn.forward2(src, attn_weights)
+
+ if random.random() >= dynamic_dropout:
+ src = src + self.conv_module2(
+ src, src_key_padding_mask=src_key_padding_mask
+ )
+
+ src = src + self.feed_forward3(src)
+
+ src = self.norm_final(self.balancer(src))
+
+ delta = src - src_orig
+
+ src = src_orig + delta * self.get_bypass_scale()
+
+ return self.whiten(src)
+
+ def streaming_forward(
+ self,
+ src: Tensor,
+ pos_emb: Tensor,
+ cached_len: Tensor,
+ cached_avg: Tensor,
+ cached_key: Tensor,
+ cached_val: Tensor,
+ cached_val2: Tensor,
+ cached_conv1: Tensor,
+ cached_conv2: Tensor,
+ ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+ """
+ Pass the input through the encoder layer.
+
+ Args:
+ src: the sequence to the encoder layer (required).
+ pos_emb: Positional embedding tensor (required).
+ cached_len: processed number of past frames.
+ cached_avg: cached average of past frames.
+ cached_key: cached key tensor of left context for the first attention module.
+ cached_val: cached value tensor of left context for the first attention module.
+ cached_val2: cached value tensor of left context for the second attention module.
+ cached_conv1: cached left context for the first convolution module.
+ cached_conv2: cached left context for the second convolution module.
+
+ Shape:
+ src: (S, N, E).
+ pos_emb: (N, left_context_len+2*S-1, E)
+ cached_len: (N,)
+ N is the batch size.
+ cached_avg: (N, C).
+ N is the batch size, C is the feature dimension.
+ cached_key: (left_context_len, N, K).
+ N is the batch size, K is the key dimension.
+ cached_val: (left_context_len, N, V).
+ N is the batch size, V is the key dimension.
+ cached_val2: (left_context_len, N, V).
+ N is the batch size, V is the key dimension.
+ cached_conv1: (N, C, kernel_size-1).
+ N is the batch size, C is the convolution channels.
+ cached_conv2: (N, C, kernel_size-1).
+ N is the batch size, C is the convolution channels.
+ """
+ src_orig = src
+
+ # macaron style feed forward module
+ src = src + self.feed_forward1(src)
+
+ src_pool, cached_len, cached_avg = self.pooling.streaming_forward(
+ src,
+ cached_len=cached_len,
+ cached_avg=cached_avg,
+ )
+ src = src + src_pool
+
+ (
+ src_attn,
+ attn_weights,
+ cached_key,
+ cached_val,
+ ) = self.self_attn.streaming_forward(
+ src,
+ pos_emb=pos_emb,
+ cached_key=cached_key,
+ cached_val=cached_val,
+ )
+ src = src + src_attn
+
+ src_conv, cached_conv1 = self.conv_module1.streaming_forward(
+ src,
+ cache=cached_conv1,
+ )
+ src = src + src_conv
+
+ src = src + self.feed_forward2(src)
+
+ src_attn, cached_val2 = self.self_attn.streaming_forward2(
+ src,
+ attn_weights,
+ cached_val=cached_val2,
+ )
+ src = src + src_attn
+
+ src_conv, cached_conv2 = self.conv_module2.streaming_forward(
+ src,
+ cache=cached_conv2,
+ )
+ src = src + src_conv
+
+ src = src + self.feed_forward3(src)
+
+ src = self.norm_final(self.balancer(src))
+
+ delta = src - src_orig
+
+ src = src_orig + delta * self.bypass_scale
+
+ return (
+ src,
+ cached_len,
+ cached_avg,
+ cached_key,
+ cached_val,
+ cached_val2,
+ cached_conv1,
+ cached_conv2,
+ )
+
+
+class ZipformerEncoder(nn.Module):
+ r"""ZipformerEncoder is a stack of N encoder layers
+
+ Args:
+ encoder_layer: an instance of the ZipformerEncoderLayer() class (required).
+ num_layers: the number of sub-encoder-layers in the encoder (required).
+
+ Examples::
+ >>> encoder_layer = ZipformerEncoderLayer(d_model=512, nhead=8)
+ >>> zipformer_encoder = ZipformerEncoder(encoder_layer, num_layers=6)
+ >>> src = torch.rand(10, 32, 512)
+ >>> out = zipformer_encoder(src)
+ """
+
+ def __init__(
+ self,
+ encoder_layer: nn.Module,
+ num_layers: int,
+ dropout: float,
+ warmup_begin: float,
+ warmup_end: float,
+ ) -> None:
+ super().__init__()
+ # will be written to, see set_batch_count() Note: in inference time this
+ # may be zero but should be treated as large, we can check if
+ # self.training is true.
+ self.batch_count = 0
+ self.warmup_begin = warmup_begin
+ self.warmup_end = warmup_end
+ # module_seed is for when we need a random number that is unique to the module but
+ # shared across jobs. It's used to randomly select how many layers to drop,
+ # so that we can keep this consistent across worker tasks (for efficiency).
+ self.module_seed = torch.randint(0, 1000, ()).item()
+
+ self.encoder_pos = RelPositionalEncoding(encoder_layer.d_model, dropout)
+
+ self.layers = nn.ModuleList(
+ [copy.deepcopy(encoder_layer) for i in range(num_layers)]
+ )
+ self.num_layers = num_layers
+
+ self.d_model = encoder_layer.d_model
+ self.attention_dim = encoder_layer.attention_dim
+ self.cnn_module_kernel = encoder_layer.cnn_module_kernel
+
+ assert 0 <= warmup_begin <= warmup_end, (warmup_begin, warmup_end)
+
+ delta = (1.0 / num_layers) * (warmup_end - warmup_begin)
+ cur_begin = warmup_begin
+ for i in range(num_layers):
+ self.layers[i].warmup_begin = cur_begin
+ cur_begin += delta
+ self.layers[i].warmup_end = cur_begin
+
+ def get_layers_to_drop(self, rnd_seed: int):
+ ans = set()
+ if not self.training:
+ return ans
+
+ batch_count = self.batch_count
+ num_layers = len(self.layers)
+
+ def get_layerdrop_prob(layer: int) -> float:
+ layer_warmup_begin = self.layers[layer].warmup_begin
+ layer_warmup_end = self.layers[layer].warmup_end
+
+ initial_layerdrop_prob = 0.5
+ final_layerdrop_prob = 0.05
+
+ if batch_count == 0:
+ # As a special case, if batch_count == 0, return 0 (drop no
+ # layers). This is rather ugly, I'm afraid; it is intended to
+ # enable our scan_pessimistic_batches_for_oom() code to work correctly
+ # so if we are going to get OOM it will happen early.
+ # also search for 'batch_count' with quotes in this file to see
+ # how we initialize the warmup count to a random number between
+ # 0 and 10.
+ return 0.0
+ elif batch_count < layer_warmup_begin:
+ return initial_layerdrop_prob
+ elif batch_count > layer_warmup_end:
+ return final_layerdrop_prob
+ else:
+ # linearly interpolate
+ t = (batch_count - layer_warmup_begin) / layer_warmup_end
+ assert 0.0 <= t < 1.001, t
+ return initial_layerdrop_prob + t * (
+ final_layerdrop_prob - initial_layerdrop_prob
+ )
+
+ shared_rng = random.Random(batch_count + self.module_seed)
+ independent_rng = random.Random(rnd_seed)
+
+ layerdrop_probs = [get_layerdrop_prob(i) for i in range(num_layers)]
+ tot = sum(layerdrop_probs)
+ # Instead of drawing the samples independently, we first randomly decide
+ # how many layers to drop out, using the same random number generator between
+ # jobs so that all jobs drop out the same number (this is for speed).
+ # Then we use an approximate approach to drop out the individual layers
+ # with their specified probs while reaching this exact target.
+ num_to_drop = int(tot) + int(shared_rng.random() < (tot - int(tot)))
+
+ layers = list(range(num_layers))
+ independent_rng.shuffle(layers)
+
+ # go through the shuffled layers until we get the required number of samples.
+ if num_to_drop > 0:
+ for layer in itertools.cycle(layers):
+ if independent_rng.random() < layerdrop_probs[layer]:
+ ans.add(layer)
+ if len(ans) == num_to_drop:
+ break
+ if shared_rng.random() < 0.005 or __name__ == "__main__":
+ logging.info(
+ f"warmup_begin={self.warmup_begin:.1f}, warmup_end={self.warmup_end:.1f}, "
+ f"batch_count={batch_count:.1f}, num_to_drop={num_to_drop}, layers_to_drop={ans}"
+ )
+ return ans
+
+ def forward(
+ self,
+ src: Tensor,
+ # Note: The type of feature_mask should be Union[float, Tensor],
+ # but to make torch.jit.script() work, we use `float` here
+ feature_mask: float = 1.0,
+ attn_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).
+ feature_mask: something that broadcasts with src, that we'll multiply `src`
+ by at every layer.
+ 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
+
+ Returns: (x, x_no_combine), both of shape (S, N, E)
+ """
+ pos_emb = self.encoder_pos(src)
+ output = src
+
+ if torch.jit.is_scripting():
+ layers_to_drop = []
+ else:
+ rnd_seed = src.numel() + random.randint(0, 1000)
+ layers_to_drop = self.get_layers_to_drop(rnd_seed)
+
+ output = output * feature_mask
+
+ for i, mod in enumerate(self.layers):
+ if not torch.jit.is_scripting():
+ if i in layers_to_drop:
+ continue
+ output = mod(
+ output,
+ pos_emb,
+ attn_mask=attn_mask,
+ src_key_padding_mask=src_key_padding_mask,
+ )
+
+ output = output * feature_mask
+
+ return output
+
+ @torch.jit.export
+ def streaming_forward(
+ self,
+ src: Tensor,
+ cached_len: Tensor,
+ cached_avg: Tensor,
+ cached_key: Tensor,
+ cached_val: Tensor,
+ cached_val2: Tensor,
+ cached_conv1: Tensor,
+ cached_conv2: Tensor,
+ ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+ r"""Pass the input through the encoder layers in turn.
+
+ Args:
+ src: the sequence to the encoder (required).
+ cached_len: number of past frames.
+ cached_avg: cached average of past frames.
+ cached_key: cached key tensor for first attention module.
+ cached_val: cached value tensor for first attention module.
+ cached_val2: cached value tensor for second attention module.
+ cached_conv1: cached left contexts for the first convolution module.
+ cached_conv2: cached left contexts for the second convolution module.
+
+ Shape:
+ src: (S, N, E).
+ cached_len: (N,)
+ N is the batch size.
+ cached_avg: (num_layers, N, C).
+ N is the batch size, C is the feature dimension.
+ cached_key: (num_layers, left_context_len, N, K).
+ N is the batch size, K is the key dimension.
+ cached_val: (num_layers, left_context_len, N, V).
+ N is the batch size, V is the key dimension.
+ cached_val2: (num_layers, left_context_len, N, V).
+ N is the batch size, V is the key dimension.
+ cached_conv1: (num_layers, N, C, kernel_size-1).
+ N is the batch size, C is the convolution channels.
+ cached_conv2: (num_layers, N, C, kernel_size-1).
+ N is the batch size, C is the convolution channels.
+
+ Returns: A tuple of 8 tensors:
+ - output tensor
+ - updated cached number of past frmaes.
+ - updated cached average of past frmaes.
+ - updated cached key tensor of of the first attention module.
+ - updated cached value tensor of of the first attention module.
+ - updated cached value tensor of of the second attention module.
+ - updated cached left contexts of the first convolution module.
+ - updated cached left contexts of the second convolution module.
+ """
+ assert cached_len.size(0) == self.num_layers, (
+ cached_len.size(0),
+ self.num_layers,
+ )
+ assert cached_avg.size(0) == self.num_layers, (
+ cached_avg.size(0),
+ self.num_layers,
+ )
+ assert cached_key.size(0) == self.num_layers, (
+ cached_key.size(0),
+ self.num_layers,
+ )
+ assert cached_val.size(0) == self.num_layers, (
+ cached_val.size(0),
+ self.num_layers,
+ )
+ assert cached_val2.size(0) == self.num_layers, (
+ cached_val2.size(0),
+ self.num_layers,
+ )
+ assert cached_conv1.size(0) == self.num_layers, (
+ cached_conv1.size(0),
+ self.num_layers,
+ )
+ assert cached_conv2.size(0) == self.num_layers, (
+ cached_conv2.size(0),
+ self.num_layers,
+ )
+
+ left_context_len = cached_key.shape[1]
+ pos_emb = self.encoder_pos(src, left_context_len)
+ output = src
+
+ new_cached_len = []
+ new_cached_avg = []
+ new_cached_key = []
+ new_cached_val = []
+ new_cached_val2 = []
+ new_cached_conv1 = []
+ new_cached_conv2 = []
+ for i, mod in enumerate(self.layers):
+ output, len_avg, avg, key, val, val2, conv1, conv2 = mod.streaming_forward(
+ output,
+ pos_emb,
+ cached_len=cached_len[i],
+ cached_avg=cached_avg[i],
+ cached_key=cached_key[i],
+ cached_val=cached_val[i],
+ cached_val2=cached_val2[i],
+ cached_conv1=cached_conv1[i],
+ cached_conv2=cached_conv2[i],
+ )
+ # Update caches
+ new_cached_len.append(len_avg)
+ new_cached_avg.append(avg)
+ new_cached_key.append(key)
+ new_cached_val.append(val)
+ new_cached_val2.append(val2)
+ new_cached_conv1.append(conv1)
+ new_cached_conv2.append(conv2)
+
+ return (
+ output,
+ torch.stack(new_cached_len, dim=0),
+ torch.stack(new_cached_avg, dim=0),
+ torch.stack(new_cached_key, dim=0),
+ torch.stack(new_cached_val, dim=0),
+ torch.stack(new_cached_val2, dim=0),
+ torch.stack(new_cached_conv1, dim=0),
+ torch.stack(new_cached_conv2, dim=0),
+ )
+
+
+class DownsampledZipformerEncoder(nn.Module):
+ r"""
+ DownsampledZipformerEncoder is a zipformer encoder evaluated at a reduced frame rate,
+ after convolutional downsampling, and then upsampled again at the output, and combined
+ with the origin input, so that the output has the same shape as the input.
+ """
+
+ def __init__(
+ self, encoder: nn.Module, input_dim: int, output_dim: int, downsample: int
+ ):
+ super(DownsampledZipformerEncoder, self).__init__()
+ self.downsample_factor = downsample
+ self.downsample = AttentionDownsample(input_dim, output_dim, downsample)
+ self.encoder = encoder
+ self.num_layers = encoder.num_layers
+ self.d_model = encoder.d_model
+ self.attention_dim = encoder.attention_dim
+ self.cnn_module_kernel = encoder.cnn_module_kernel
+ self.upsample = SimpleUpsample(output_dim, downsample)
+ self.out_combiner = SimpleCombiner(
+ input_dim, output_dim, min_weight=(0.0, 0.25)
+ )
+
+ def forward(
+ self,
+ src: Tensor,
+ # Note: the type of feature_mask should be Unino[float, Tensor],
+ # but to make torch.jit.script() happ, we use float here
+ feature_mask: float = 1.0,
+ attn_mask: Optional[Tensor] = None,
+ src_key_padding_mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ r"""Downsample, go through encoder, upsample.
+
+ Args:
+ src: the sequence to the encoder (required).
+ feature_mask: something that broadcasts with src, that we'll multiply `src`
+ by at every layer. feature_mask is expected to be already downsampled by
+ self.downsample_factor.
+ attn_mask: attention mask (optional). Should be downsampled already.
+ src_key_padding_mask: the mask for the src keys per batch (optional). Should be downsampled already.
+
+ Shape:
+ src: (S, N, E).
+ attn_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
+
+ Returns: output of shape (S, N, F) where F is the number of output features
+ (output_dim to constructor)
+ """
+ src_orig = src
+ src = self.downsample(src)
+
+ src = self.encoder(
+ src,
+ feature_mask=feature_mask,
+ attn_mask=attn_mask,
+ src_key_padding_mask=src_key_padding_mask,
+ )
+ src = self.upsample(src)
+ # remove any extra frames that are not a multiple of downsample_factor
+ src = src[: src_orig.shape[0]]
+
+ return self.out_combiner(src_orig, src)
+
+ def streaming_forward(
+ self,
+ src: Tensor,
+ cached_len: Tensor,
+ cached_avg: Tensor,
+ cached_key: Tensor,
+ cached_val: Tensor,
+ cached_val2: Tensor,
+ cached_conv1: Tensor,
+ cached_conv2: Tensor,
+ ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+ r"""Downsample, go through encoder, upsample.
+
+ Args:
+ src: the sequence to the encoder (required).
+ cached_avg: cached average value of past frames.
+ cached_len: length of past frames.
+ cached_key: cached key tensor for the first attention module.
+ cached_val: cached value tensor for the first attention module.
+ cached_val2: cached value tensor for the second attention module.
+ cached_conv1: cached left context for the first convolution module.
+ cached_conv2: cached left context for the second convolution module.
+
+ Shape:
+ src: (S, N, E).
+ cached_len: (N,)
+ N is the batch size.
+ cached_avg: (num_layers, N, C).
+ N is the batch size, C is the feature dimension.
+ cached_key: (num_layers, left_context_len, N, K).
+ N is the batch size, K is the key dimension.
+ cached_val: (num_layers, left_context_len, N, V).
+ N is the batch size, V is the key dimension.
+ cached_val2: (num_layers, left_context_len, N, V).
+ N is the batch size, V is the key dimension.
+ cached_conv1: (num_layers, N, C, kernel_size-1).
+ N is the batch size, C is the convolution channels.
+ cached_conv2: (num_layers, N, C, kernel_size-1).
+ N is the batch size, C is the convolution channels.
+ Returns: output of shape (S, N, F) where F is the number of output features
+ (output_dim to constructor)
+ """
+ src_orig = src
+ src = self.downsample(src)
+
+ (
+ src,
+ cached_len,
+ cached_avg,
+ cached_key,
+ cached_val,
+ cached_val2,
+ cached_conv1,
+ cached_conv2,
+ ) = self.encoder.streaming_forward(
+ src,
+ cached_len=cached_len,
+ cached_avg=cached_avg,
+ cached_key=cached_key,
+ cached_val=cached_val,
+ cached_val2=cached_val2,
+ cached_conv1=cached_conv1,
+ cached_conv2=cached_conv2,
+ )
+ src = self.upsample(src)
+ # remove any extra frames that are not a multiple of downsample_factor
+ src = src[: src_orig.shape[0]]
+
+ return (
+ self.out_combiner(src_orig, src),
+ cached_len,
+ cached_avg,
+ cached_key,
+ cached_val,
+ cached_val2,
+ cached_conv1,
+ cached_conv2,
+ )
+
+
+class AttentionDownsample(torch.nn.Module):
+ """
+ Does downsampling with attention, by weighted sum, and a projection..
+ """
+
+ def __init__(self, in_channels: int, out_channels: int, downsample: int):
+ """
+ Require out_channels > in_channels.
+ """
+ super(AttentionDownsample, self).__init__()
+ self.query = nn.Parameter(torch.randn(in_channels) * (in_channels**-0.5))
+
+ # fill in the extra dimensions with a projection of the input
+ if out_channels > in_channels:
+ self.extra_proj = nn.Linear(
+ in_channels * downsample, out_channels - in_channels, bias=False
+ )
+ else:
+ self.extra_proj = None
+ self.downsample = downsample
+
+ def forward(self, src: Tensor) -> Tensor:
+ """
+ x: (seq_len, 1, in_channels)
+ Returns a tensor of shape
+ ( (seq_len+downsample-1)//downsample, batch_size, out_channels)
+ """
+ (seq_len, batch_size, in_channels) = src.shape
+ ds = self.downsample
+ d_seq_len = (seq_len + ds - 1) // ds
+
+ # Pad to an exact multiple of self.downsample
+ if seq_len != d_seq_len * ds:
+ # right-pad src, repeating the last element.
+ pad = d_seq_len * ds - seq_len
+ src_extra = src[src.shape[0] - 1 :].expand(pad, src.shape[1], src.shape[2])
+ src = torch.cat((src, src_extra), dim=0)
+ assert src.shape[0] == d_seq_len * ds, (src.shape[0], d_seq_len, ds)
+
+ src = src.reshape(d_seq_len, ds, batch_size, in_channels)
+ scores = (src * self.query).sum(dim=-1, keepdim=True)
+
+ if not torch.jit.is_scripting() and not torch.jit.is_tracing():
+ scores = penalize_abs_values_gt(scores, limit=10.0, penalty=1.0e-04)
+
+ weights = scores.softmax(dim=1)
+
+ # ans1 is the first `in_channels` channels of the output
+ ans = (src * weights).sum(dim=1)
+ src = src.permute(0, 2, 1, 3).reshape(d_seq_len, batch_size, ds * in_channels)
+
+ if self.extra_proj is not None:
+ ans2 = self.extra_proj(src)
+ ans = torch.cat((ans, ans2), dim=2)
+ return ans
+
+
+class SimpleUpsample(torch.nn.Module):
+ """
+ A very simple form of upsampling that mostly just repeats the input, but
+ also adds a position-specific bias.
+ """
+
+ def __init__(self, num_channels: int, upsample: int):
+ super(SimpleUpsample, self).__init__()
+ self.bias = nn.Parameter(torch.randn(upsample, num_channels) * 0.01)
+
+ def forward(self, src: Tensor) -> Tensor:
+ """
+ x: (seq_len, batch_size, num_channels)
+ Returns a tensor of shape
+ ( (seq_len*upsample), batch_size, num_channels)
+ """
+ upsample = self.bias.shape[0]
+ (seq_len, batch_size, num_channels) = src.shape
+ src = src.unsqueeze(1).expand(seq_len, upsample, batch_size, num_channels)
+ src = src + self.bias.unsqueeze(1)
+ src = src.reshape(seq_len * upsample, batch_size, num_channels)
+ return src
+
+
+class SimpleCombinerIdentity(nn.Module):
+ def __init__(self, *args, **kwargs):
+ super().__init__()
+
+ def forward(self, src1: Tensor, src2: Tensor) -> Tensor:
+ return src1
+
+
+class SimpleCombiner(torch.nn.Module):
+ """
+ A very simple way of combining 2 vectors of 2 different dims, via a
+ learned weighted combination in the shared part of the dim.
+ Args:
+ dim1: the dimension of the first input, e.g. 256
+ dim2: the dimension of the second input, e.g. 384.
+ The output will have the same dimension as dim2.
+ """
+
+ def __init__(self, dim1: int, dim2: int, min_weight: Tuple[float] = (0.0, 0.0)):
+ super(SimpleCombiner, self).__init__()
+ assert dim2 >= dim1, (dim2, dim1)
+ self.weight1 = nn.Parameter(torch.zeros(()))
+ self.min_weight = min_weight
+
+ def forward(self, src1: Tensor, src2: Tensor) -> Tensor:
+ """
+ src1: (*, dim1)
+ src2: (*, dim2)
+
+ Returns: a tensor of shape (*, dim2)
+ """
+ assert src1.shape[:-1] == src2.shape[:-1], (src1.shape, src2.shape)
+
+ weight1 = self.weight1
+ if not torch.jit.is_scripting():
+ if (
+ self.training
+ and random.random() < 0.25
+ and self.min_weight != (0.0, 0.0)
+ ):
+ weight1 = weight1.clamp(
+ min=self.min_weight[0], max=1.0 - self.min_weight[1]
+ )
+
+ src1 = src1 * weight1
+ src2 = src2 * (1.0 - weight1)
+
+ src1_dim = src1.shape[-1]
+ src2_dim = src2.shape[-1]
+ if src1_dim != src2_dim:
+ if src1_dim < src2_dim:
+ src1 = torch.nn.functional.pad(src1, (0, src2_dim - src1_dim))
+ else:
+ src1 = src1[:src2_dim]
+
+ return src1 + src2
+
+
+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 a PositionalEncoding object."""
+ super(RelPositionalEncoding, self).__init__()
+ self.d_model = d_model
+ self.dropout = torch.nn.Dropout(dropout_rate)
+ self.pe = None
+ self.extend_pe(torch.tensor(0.0).expand(max_len))
+
+ def extend_pe(self, x: Tensor, left_context_len: int = 0) -> None:
+ """Reset the positional encodings."""
+ x_size_left = x.size(0) + left_context_len
+ 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_left * 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 Tensor:
+ """Add positional encoding.
+
+ Args:
+ x (torch.Tensor): Input tensor (time, batch, `*`).
+ left_context_len: (int): Length of cached left context.
+
+ Returns:
+ torch.Tensor: Encoded tensor (batch, left_context_len + 2*time-1, `*`).
+
+ """
+ self.extend_pe(x, left_context_len)
+ x_size_left = x.size(0) + left_context_len
+ pos_emb = self.pe[
+ :,
+ self.pe.size(1) // 2
+ - x_size_left
+ + 1 : self.pe.size(1) // 2 # noqa E203
+ + x.size(0),
+ ]
+ return self.dropout(pos_emb)
+
+
+class RelPositionMultiheadAttention(nn.Module):
+ r"""Multi-Head Attention layer with relative position encoding
+
+ This is a quite heavily modified from: "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context",
+ we have to write up the differences.
+
+
+ Args:
+ embed_dim: total dimension of the model.
+ attention_dim: dimension in the attention module, may be less or more than embed_dim
+ but must be a multiple of num_heads.
+ 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,
+ attention_dim: int,
+ num_heads: int,
+ pos_dim: int,
+ dropout: float = 0.0,
+ ) -> None:
+ super(RelPositionMultiheadAttention, self).__init__()
+ self.embed_dim = embed_dim
+ self.attention_dim = attention_dim
+ self.num_heads = num_heads
+ self.dropout = dropout
+ self.head_dim = attention_dim // num_heads
+ self.pos_dim = pos_dim
+ assert self.head_dim % 2 == 0, self.head_dim
+ assert self.head_dim * num_heads == attention_dim, (
+ self.head_dim,
+ num_heads,
+ attention_dim,
+ )
+
+ # the initial_scale is supposed to take over the "scaling" factor of
+ # head_dim ** -0.5, dividing it between the query and key.
+ in_proj_dim = (
+ 2 * attention_dim
+ + attention_dim // 2 # query, key
+ + pos_dim * num_heads # value
+ ) # positional encoding query
+
+ self.in_proj = ScaledLinear(
+ embed_dim, in_proj_dim, bias=True, initial_scale=self.head_dim**-0.25
+ )
+
+ # self.whiten_values is applied on the values in forward();
+ # it just copies the keys but prevents low-rank distribution by modifying grads.
+ self.whiten_values = Whiten(
+ num_groups=num_heads,
+ whitening_limit=2.0,
+ prob=(0.025, 0.25),
+ grad_scale=0.025,
+ )
+ self.whiten_keys = Whiten(
+ num_groups=num_heads,
+ whitening_limit=2.0,
+ prob=(0.025, 0.25),
+ grad_scale=0.025,
+ )
+
+ # linear transformation for positional encoding.
+ self.linear_pos = ScaledLinear(
+ embed_dim, num_heads * pos_dim, bias=False, initial_scale=0.05
+ )
+
+ # the following are for diagnosics only, see --print-diagnostics option.
+ # they only copy their inputs.
+ self.copy_pos_query = Identity()
+ self.copy_query = Identity()
+
+ self.out_proj = ScaledLinear(
+ attention_dim // 2, embed_dim, bias=True, initial_scale=0.05
+ )
+
+ self.in_proj2 = nn.Linear(embed_dim, attention_dim // 2, bias=False)
+ self.out_proj2 = ScaledLinear(
+ attention_dim // 2, embed_dim, bias=True, initial_scale=0.05
+ )
+ # self.whiten_values2 is applied on the values in forward2()
+ self.whiten_values2 = Whiten(
+ num_groups=num_heads,
+ whitening_limit=2.0,
+ prob=(0.025, 0.25),
+ grad_scale=0.025,
+ )
+
+ def forward(
+ self,
+ x: Tensor,
+ pos_emb: Tensor,
+ key_padding_mask: Optional[Tensor] = None,
+ attn_mask: Optional[Tensor] = None,
+ ) -> Tuple[Tensor, Tensor]:
+ r"""
+ Args:
+ x: input to be projected to query, key, value
+ 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
+ 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:
+ - x: :math:`(L, N, E)` where L is the target 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.
+
+ - Returns: (attn_output, attn_weights)
+
+ - attn_output: :math:`(S, N, E)` where S is the sequence length, N is the batch size,
+ E is the embedding dimension.
+ - attn_weights: :math:`(N * N, S, S)` where N is the batch size, H is the num-heads
+ and S is the sequence length.
+ """
+ x, weights = self.multi_head_attention_forward(
+ self.in_proj(x),
+ self.linear_pos(pos_emb),
+ self.attention_dim,
+ self.num_heads,
+ self.dropout,
+ self.out_proj.weight,
+ self.out_proj.bias,
+ training=self.training,
+ key_padding_mask=key_padding_mask,
+ attn_mask=attn_mask,
+ )
+ return x, weights
+
+ def streaming_forward(
+ self,
+ x: Tensor,
+ pos_emb: Tensor,
+ cached_key: Tensor,
+ cached_val: Tensor,
+ ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+ r"""
+ Args:
+ x: input to be projected to query, key, value
+ pos_emb: Positional embedding tensor
+ 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:
+ - x: :math:`(L, N, E)` where L is the target 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.
+ - 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.
+ - cached_key: :math:`(left_context_len, N, K)`, where N is the batch size, K is the key dimension.
+ - cached_val: :math:`(left_context_len, N, V)`, where N is the batch size, V is the value dimension.
+
+ - Returns: (attn_output, attn_weights, cached_key, cached_val)
+
+ - attn_output: :math:`(S, N, E)` where S is the sequence length, N is the batch size,
+ E is the embedding dimension.
+ - attn_weights: :math:`(N * N, S, S)` where N is the batch size, H is the num-heads
+ and S is the sequence length.
+ - cached_key: :math:`(left_context_len, N, K)`, updated cached attention key tensor of
+ left context
+ - cached_val: :math:`(left_context_len, N, K)`, updated cached attention value tensor of
+ """
+ (
+ x,
+ weights,
+ cached_key,
+ cached_val,
+ ) = self.streaming_multi_head_attention_forward(
+ self.in_proj(x),
+ self.linear_pos(pos_emb),
+ self.attention_dim,
+ self.num_heads,
+ self.out_proj.weight,
+ self.out_proj.bias,
+ cached_key=cached_key,
+ cached_val=cached_val,
+ )
+ return x, weights, cached_key, cached_val
+
+ def multi_head_attention_forward(
+ self,
+ x_proj: Tensor,
+ pos: Tensor,
+ attention_dim: int,
+ num_heads: int,
+ dropout_p: float,
+ out_proj_weight: Tensor,
+ out_proj_bias: Tensor,
+ training: bool = True,
+ key_padding_mask: Optional[Tensor] = None,
+ attn_mask: Optional[Tensor] = None,
+ ) -> Tuple[Tensor, Tensor]:
+ r"""
+ Args:
+ x_proj: the projected input, to be split into query, key, value.
+ pos: head-specific biases arising from the positional embeddings.
+ attention_dim: dimension inside attention mechanism
+ num_heads: parallel attention heads.
+ 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.
+ 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:
+ - x: :math:`(L, N, 7 * A // 2)` where L is the target sequence length, N is the batch size, A is
+ the attention dimension. Will be split into (query, key, value, pos).
+ - pos: :math:`(N, 2*L-1, A//2)` or :math:`(1, 2*L-1, A//2)` where L is the sequence
+ length, N is the batch size, and A is the attention dim.
+ - 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_weights: :math:`(N * H, S, S)` where N is the batch size,
+ H is the num-heads, S is the sequence length.
+ """
+
+ seq_len, bsz, _ = x_proj.size()
+
+ head_dim = attention_dim // num_heads
+ pos_dim = self.pos_dim # positional-encoding dim per head
+ assert (
+ head_dim * num_heads == attention_dim
+ ), f"attention_dim must be divisible by num_heads: {head_dim}, {num_heads}, {attention_dim}"
+
+ # self-attention
+ q = x_proj[..., 0:attention_dim]
+ k = x_proj[..., attention_dim : 2 * attention_dim]
+ value_dim = attention_dim // 2
+ v = x_proj[..., 2 * attention_dim : 2 * attention_dim + value_dim]
+ # p is the position-encoding query, its dimension is num_heads*pos_dim..
+ p = x_proj[..., 2 * attention_dim + value_dim :]
+
+ k = self.whiten_keys(k) # does nothing in the forward pass.
+ v = self.whiten_values(v) # does nothing in the forward pass.
+ q = self.copy_query(q) # for diagnostics only, does nothing.
+ p = self.copy_pos_query(p) # for diagnostics only, does nothing.
+
+ 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, seq_len, seq_len]:
+ 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,
+ seq_len,
+ seq_len,
+ ]:
+ 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.reshape(seq_len, bsz, num_heads, head_dim)
+ p = p.reshape(seq_len, bsz, num_heads, pos_dim)
+ k = k.reshape(seq_len, bsz, num_heads, head_dim)
+ v = v.reshape(seq_len, bsz * num_heads, head_dim // 2).transpose(0, 1)
+
+ 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) == seq_len, "{} == {}".format(
+ key_padding_mask.size(1), seq_len
+ )
+
+ q = q.permute(1, 2, 0, 3) # (batch, head, time1, head_dim)
+ p = p.permute(1, 2, 0, 3) # (batch, head, time1, pos_dim)
+ k = k.permute(1, 2, 3, 0) # (batch, head, d_k, time2)
+
+ seq_len2 = 2 * seq_len - 1
+ pos = pos.reshape(1, seq_len2, num_heads, pos_dim).permute(0, 2, 3, 1)
+ # pos shape now: (batch, head, pos_dim, seq_len2)
+
+ # (batch, head, time1, pos_dim) x (1, head, pos_dim, seq_len2) -> (batch, head, time1, seq_len2)
+ # [where seq_len2 represents relative position.]
+ pos_weights = torch.matmul(p, pos)
+ # the following .as_strided() expression converts the last axis of pos_weights from relative
+ # to absolute position. I don't know whether I might have got the time-offsets backwards or
+ # not, but let this code define which way round it is supposed to be.
+ pos_weights = pos_weights.as_strided(
+ (bsz, num_heads, seq_len, seq_len),
+ (
+ pos_weights.stride(0),
+ pos_weights.stride(1),
+ pos_weights.stride(2) - pos_weights.stride(3),
+ pos_weights.stride(3),
+ ),
+ storage_offset=pos_weights.stride(3) * (seq_len - 1),
+ )
+
+ # caution: they are really scores at this point.
+ attn_output_weights = torch.matmul(q, k) + pos_weights
+
+ if not torch.jit.is_scripting():
+ if training and random.random() < 0.1:
+ # This is a harder way of limiting the attention scores to not be too large.
+ # It incurs a penalty if any of them has an absolute value greater than 50.0.
+ # this should be outside the normal range of the attention scores. We use
+ # this mechanism instead of, say, a limit on entropy, because once the entropy
+ # gets very small gradients through the softmax can become very small, and
+ # some mechanisms like that become ineffective.
+ attn_output_weights = penalize_abs_values_gt(
+ attn_output_weights, limit=25.0, penalty=1.0e-04
+ )
+
+ # attn_output_weights: (batch, head, time1, time2)
+ attn_output_weights = attn_output_weights.view(
+ bsz * num_heads, seq_len, seq_len
+ )
+
+ if attn_mask is not None:
+ if attn_mask.dtype == torch.bool:
+ attn_output_weights = attn_output_weights.masked_fill(
+ attn_mask, float("-inf")
+ )
+ else:
+ attn_output_weights = attn_output_weights + attn_mask
+
+ if key_padding_mask is not None:
+ attn_output_weights = attn_output_weights.view(
+ bsz, num_heads, seq_len, seq_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, seq_len, seq_len
+ )
+
+ # Using this version of softmax, defined in scaling.py,
+ # should save a little of the memory used in backprop by, if
+ # we are in automatic mixed precision mode (amp) == autocast,
+ # only storing the half-precision output for backprop purposes.
+ attn_output_weights = softmax(attn_output_weights, dim=-1)
+
+ # If we are using chunk-wise attention mask and setting a limited
+ # num_left_chunks, the attention may only see the padding values which
+ # will also be masked out by `key_padding_mask`. At this circumstances,
+ # the whole column of `attn_output_weights` will be `-inf`
+ # (i.e. be `nan` after softmax). So we fill `0.0` at the masking
+ # positions to avoid invalid loss value below.
+ if (
+ attn_mask is not None
+ and attn_mask.dtype == torch.bool
+ and key_padding_mask is not None
+ ):
+ if attn_mask.size(0) != 1:
+ attn_mask = attn_mask.view(bsz, num_heads, seq_len, seq_len)
+ combined_mask = attn_mask | key_padding_mask.unsqueeze(1).unsqueeze(2)
+ else:
+ # attn_mask.shape == (1, tgt_len, src_len)
+ combined_mask = attn_mask.unsqueeze(0) | key_padding_mask.unsqueeze(
+ 1
+ ).unsqueeze(2)
+
+ attn_output_weights = attn_output_weights.view(
+ bsz, num_heads, seq_len, seq_len
+ )
+ attn_output_weights = attn_output_weights.masked_fill(combined_mask, 0.0)
+ attn_output_weights = attn_output_weights.view(
+ bsz * num_heads, seq_len, seq_len
+ )
+
+ 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, seq_len, head_dim // 2]
+ attn_output = (
+ attn_output.transpose(0, 1)
+ .contiguous()
+ .view(seq_len, bsz, attention_dim // 2)
+ )
+ attn_output = nn.functional.linear(attn_output, out_proj_weight, out_proj_bias)
+
+ return attn_output, attn_output_weights
+
+ def streaming_multi_head_attention_forward(
+ self,
+ x_proj: Tensor,
+ pos: Tensor,
+ attention_dim: int,
+ num_heads: int,
+ out_proj_weight: Tensor,
+ out_proj_bias: Tensor,
+ cached_key: Tensor,
+ cached_val: Tensor,
+ ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+ r"""
+ Args:
+ x_proj: the projected input, to be split into query, key, value.
+ pos: head-specific biases arising from the positional embeddings.
+ attention_dim: dimension inside attention mechanism
+ num_heads: parallel attention heads.
+ out_proj_weight, out_proj_bias: the output projection weight and bias.
+ cached_key: cached attention key tensor of left context.
+ cached_val: cached attention value tensor of left context.
+
+ Shape:
+ Inputs:
+ - x: :math:`(L, N, 7 * A // 2)` where L is the target sequence length, N is the batch size, A is
+ the attention dimension. Will be split into (query, key, value, pos).
+ - pos: :math:`(N, 2*L-1, A//2)` or :math:`(1, 2*L-1, A//2)` where L is the sequence
+ length, N is the batch size, and A is the attention dim.
+ 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.
+
+ 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_weights: :math:`(N * H, S, S)` where N is the batch size,
+ H is the num-heads, S is the sequence length.
+ - cached_key: :math:`(left_context_len, N, K)`, updated cached attention key tensor of left context.
+ - cached_val: :math:`(left_context_len, N, K)`, updated cached attention value tensor of left context.
+ """
+
+ seq_len, bsz, _ = x_proj.size()
+
+ head_dim = attention_dim // num_heads
+ pos_dim = self.pos_dim # positional-encoding dim per head
+ assert (
+ head_dim * num_heads == attention_dim
+ ), f"attention_dim must be divisible by num_heads: {head_dim}, {num_heads}, {attention_dim}"
+
+ # self-attention
+ q = x_proj[..., 0:attention_dim]
+ k = x_proj[..., attention_dim : 2 * attention_dim]
+ value_dim = attention_dim // 2
+ v = x_proj[..., 2 * attention_dim : 2 * attention_dim + value_dim]
+ # p is the position-encoding query, its dimension is num_heads*pos_dim..
+ p = x_proj[..., 2 * attention_dim + value_dim :]
+
+ left_context_len = cached_key.shape[0]
+ assert left_context_len > 0, left_context_len
+ assert cached_key.shape[0] == cached_val.shape[0], (
+ cached_key.shape,
+ cached_val.shape,
+ )
+ # Pad cached left contexts
+ k = torch.cat([cached_key, k], dim=0)
+ v = torch.cat([cached_val, v], dim=0)
+ # Update cached left contexts
+ cached_key = k[-left_context_len:, ...]
+ cached_val = v[-left_context_len:, ...]
+
+ # The length of key and value
+ kv_len = k.shape[0]
+
+ q = q.reshape(seq_len, bsz, num_heads, head_dim)
+ p = p.reshape(seq_len, bsz, num_heads, pos_dim)
+ k = k.reshape(kv_len, bsz, num_heads, head_dim)
+ v = v.reshape(kv_len, bsz * num_heads, head_dim // 2).transpose(0, 1)
+
+ q = q.permute(1, 2, 0, 3) # (batch, head, time1, head_dim)
+ p = p.permute(1, 2, 0, 3) # (batch, head, time1, pos_dim)
+ k = k.permute(1, 2, 3, 0) # (batch, head, d_k, time2)
+
+ seq_len2 = 2 * seq_len - 1 + left_context_len
+ pos = pos.reshape(1, seq_len2, num_heads, pos_dim).permute(0, 2, 3, 1)
+ # pos shape now: (batch, head, pos_dim, seq_len2)
+
+ # (batch, head, time1, pos_dim) x (1, head, pos_dim, seq_len2) -> (batch, head, time1, seq_len2)
+ # [where seq_len2 represents relative position.]
+ pos_weights = torch.matmul(p, pos)
+ # the following .as_strided() expression converts the last axis of pos_weights from relative
+ # to absolute position. I don't know whether I might have got the time-offsets backwards or
+ # not, but let this code define which way round it is supposed to be.
+ pos_weights = pos_weights.as_strided(
+ (bsz, num_heads, seq_len, kv_len),
+ (
+ pos_weights.stride(0),
+ pos_weights.stride(1),
+ pos_weights.stride(2) - pos_weights.stride(3),
+ pos_weights.stride(3),
+ ),
+ storage_offset=pos_weights.stride(3) * (seq_len - 1),
+ )
+
+ # caution: they are really scores at this point.
+ attn_output_weights = torch.matmul(q, k) + pos_weights
+
+ # attn_output_weights: (batch, head, time1, time2)
+ attn_output_weights = attn_output_weights.view(bsz * num_heads, seq_len, kv_len)
+
+ # Using this version of softmax, defined in scaling.py,
+ # should save a little of the memory used in backprop by, if
+ # we are in automatic mixed precision mode (amp) == autocast,
+ # only storing the half-precision output for backprop purposes.
+ attn_output_weights = softmax(attn_output_weights, dim=-1)
+
+ attn_output = torch.bmm(attn_output_weights, v)
+ assert list(attn_output.size()) == [bsz * num_heads, seq_len, head_dim // 2]
+ attn_output = (
+ attn_output.transpose(0, 1)
+ .contiguous()
+ .view(seq_len, bsz, attention_dim // 2)
+ )
+ attn_output = nn.functional.linear(attn_output, out_proj_weight, out_proj_bias)
+
+ return attn_output, attn_output_weights, cached_key, cached_val
+
+ def forward2(
+ self,
+ x: Tensor,
+ attn_weights: Tensor,
+ ) -> Tensor:
+ """
+ Second forward function, where we re-use the attn_weights returned by the first forward function
+ but with different input.
+ Args:
+ x: input, of shape (seq_len, batch_size, embed_dim)
+ attn_weights: attention weights returned by forward(), of shape (batch_size * num_heads, seq_len, seq_len)
+ Returns:
+ output of the same shape as x, i.e. (seq_len, batch_size, embed_dim)
+ """
+ num_heads = self.num_heads
+ (seq_len, bsz, embed_dim) = x.shape
+ head_dim = self.attention_dim // num_heads
+ # v: (tgt_len, bsz, embed_dim // 2)
+ v = self.in_proj2(x)
+ v = self.whiten_values2(v) # does nothing in the forward pass.
+ v = v.reshape(seq_len, bsz * num_heads, head_dim // 2).transpose(0, 1)
+
+ # now v: (bsz * num_heads, seq_len, head_dim // 2)
+ attn_output = torch.bmm(attn_weights, v)
+
+ if not torch.jit.is_scripting():
+ if random.random() < 0.001 or __name__ == "__main__":
+ self._print_attn_stats(attn_weights, attn_output)
+
+ # attn_output: (bsz * num_heads, seq_len, head_dim)
+ attn_output = (
+ attn_output.transpose(0, 1)
+ .contiguous()
+ .view(seq_len, bsz, self.attention_dim // 2)
+ )
+ # returned value is of shape (seq_len, bsz, embed_dim), like x.
+ return self.out_proj2(attn_output)
+
+ def streaming_forward2(
+ self,
+ x: Tensor,
+ attn_weights: Tensor,
+ cached_val: Tensor,
+ ) -> Tuple[Tensor, Tensor]:
+ """
+ Second forward function, where we re-use the attn_weights returned by the first forward function
+ but with different input.
+ Args:
+ x: input, of shape (seq_len, batch_size, embed_dim)
+ attn_weights: attention weights returned by forward(), of shape (batch_size * num_heads, seq_len, seq_len)
+ cached_val: cached attention value tensor of left context.
+ Returns:
+ - output of the same shape as x, i.e. (seq_len, batch_size, embed_dim)
+ - updated cached attention value tensor of left context.
+ """
+ num_heads = self.num_heads
+ (seq_len, bsz, embed_dim) = x.shape
+ head_dim = self.attention_dim // num_heads
+ # v: (tgt_len, bsz, embed_dim // 2)
+ v = self.in_proj2(x)
+
+ left_context_len = cached_val.shape[0]
+ assert left_context_len > 0, left_context_len
+ v = torch.cat([cached_val, v], dim=0)
+ cached_val = v[-left_context_len:]
+
+ seq_len2 = left_context_len + seq_len
+ v = v.reshape(seq_len2, bsz * num_heads, head_dim // 2).transpose(0, 1)
+
+ # now v: (bsz * num_heads, seq_len, head_dim // 2)
+ attn_output = torch.bmm(attn_weights, v)
+
+ # attn_output: (bsz * num_heads, seq_len, head_dim)
+ attn_output = (
+ attn_output.transpose(0, 1)
+ .contiguous()
+ .view(seq_len, bsz, self.attention_dim // 2)
+ )
+ # returned value is of shape (seq_len, bsz, embed_dim), like x.
+ return self.out_proj2(attn_output), cached_val
+
+ def _print_attn_stats(self, attn_weights: Tensor, attn_output: Tensor):
+ # attn_weights: (batch_size * num_heads, seq_len, seq_len)
+ # attn_output: (bsz * num_heads, seq_len, head_dim)
+ (n, seq_len, head_dim) = attn_output.shape
+ num_heads = self.num_heads
+ bsz = n // num_heads
+
+ with torch.no_grad():
+ with torch.cuda.amp.autocast(enabled=False):
+ attn_weights = attn_weights.to(torch.float32)
+ attn_output = attn_output.to(torch.float32)
+ attn_weights_entropy = (
+ -((attn_weights + 1.0e-20).log() * attn_weights)
+ .sum(dim=-1)
+ .reshape(bsz, num_heads, seq_len)
+ .mean(dim=(0, 2))
+ )
+ attn_output = attn_output.reshape(bsz, num_heads, seq_len, head_dim)
+ attn_output = attn_output.permute(1, 0, 2, 3).reshape(
+ num_heads, bsz * seq_len, head_dim
+ )
+ attn_output_mean = attn_output.mean(dim=1, keepdim=True)
+ attn_output = attn_output - attn_output_mean
+ attn_covar = torch.matmul(attn_output.transpose(1, 2), attn_output) / (
+ bsz * seq_len
+ )
+ # attn_covar: (num_heads, head_dim, head_dim)
+ # eigs, _ = torch.symeig(attn_covar)
+ # logging.info(f"attn_weights_entropy = {attn_weights_entropy}, output_eigs = {eigs}")
+
+ attn_covar = _diag(attn_covar).mean(dim=1) # (num_heads,)
+ embed_dim = self.in_proj2.weight.shape[1]
+ in_proj_covar = (
+ self.in_proj2.weight.reshape(num_heads, head_dim, embed_dim) ** 2
+ ).mean(dim=(1, 2))
+ out_proj_covar = (
+ self.out_proj2.weight.reshape(embed_dim, num_heads, head_dim) ** 2
+ ).mean(dim=(0, 2))
+ logging.info(
+ f"attn_weights_entropy = {attn_weights_entropy}, covar={attn_covar}, in_proj_covar={in_proj_covar}, out_proj_covar={out_proj_covar}"
+ )
+
+
+class PoolingModule(nn.Module):
+ """
+ Averages the input over the time dimension and project with a square matrix.
+ """
+
+ def __init__(self, d_model: int):
+ super().__init__()
+ self.proj = ScaledLinear(d_model, d_model, initial_scale=0.1, bias=False)
+
+ def forward(
+ self,
+ x: Tensor,
+ src_key_padding_mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ """
+ Args:
+ x: a Tensor of shape (T, N, C)
+ src_key_padding_mask: a Tensor of bool, of shape (N, T), with True in masked
+ positions.
+
+ Returns:
+ - output, a Tensor of shape (T, N, C).
+ """
+ if src_key_padding_mask is not None:
+ # False in padding positions
+ padding_mask = src_key_padding_mask.logical_not().to(x.dtype) # (N, T)
+ # Cumulated numbers of frames from start
+ cum_mask = padding_mask.cumsum(dim=1) # (N, T)
+ x = x.cumsum(dim=0) # (T, N, C)
+ pooling_mask = padding_mask / cum_mask
+ pooling_mask = pooling_mask.transpose(0, 1).contiguous().unsqueeze(-1)
+ # now pooling_mask: (T, N, 1)
+ x = x * pooling_mask # (T, N, C)
+ else:
+ num_frames = x.shape[0]
+ cum_mask = torch.arange(1, num_frames + 1).unsqueeze(1) # (T, 1)
+ x = x.cumsum(dim=0) # (T, N, C)
+ pooling_mask = (1.0 / cum_mask).unsqueeze(2)
+ # now pooling_mask: (T, N, 1)
+ x = x * pooling_mask
+
+ x = self.proj(x)
+ return x
+
+ def streaming_forward(
+ self,
+ x: Tensor,
+ cached_len: Tensor,
+ cached_avg: Tensor,
+ ) -> Tuple[Tensor, Tensor, Tensor]:
+ """
+ Args:
+ x: a Tensor of shape (T, N, C)
+ cached_len: a Tensor of int, of shape (N,), containing the number of
+ past frames in batch.
+ cached_avg: a Tensor of shape (N, C), the average over all past frames
+ in batch.
+
+ Returns:
+ A tuple of 2 tensors:
+ - output, a Tensor of shape (T, N, C).
+ - updated cached_avg, a Tensor of shape (N, C).
+ """
+ x = x.cumsum(dim=0) # (T, N, C)
+ x = x + (cached_avg * cached_len.unsqueeze(1)).unsqueeze(0)
+ # Cumulated numbers of frames from start
+ cum_mask = torch.arange(1, x.size(0) + 1, device=x.device)
+ cum_mask = cum_mask.unsqueeze(1) + cached_len.unsqueeze(0) # (T, N)
+ pooling_mask = (1.0 / cum_mask).unsqueeze(2)
+ # now pooling_mask: (T, N, 1)
+ x = x * pooling_mask # (T, N, C)
+
+ cached_len = cached_len + x.size(0)
+ cached_avg = x[-1]
+
+ x = self.proj(x)
+ return x, cached_len, cached_avg
+
+
+class FeedforwardModule(nn.Module):
+ """Feedforward module in Zipformer model."""
+
+ def __init__(self, d_model: int, feedforward_dim: int, dropout: float):
+ super(FeedforwardModule, self).__init__()
+ self.in_proj = nn.Linear(d_model, feedforward_dim)
+ self.balancer = ActivationBalancer(
+ feedforward_dim, channel_dim=-1, max_abs=10.0, min_prob=0.25
+ )
+ self.activation = DoubleSwish()
+ self.dropout = nn.Dropout(dropout)
+ self.out_proj = ScaledLinear(feedforward_dim, d_model, initial_scale=0.01)
+
+ def forward(self, x: Tensor):
+ x = self.in_proj(x)
+ x = self.balancer(x)
+ x = self.activation(x)
+ x = self.dropout(x)
+ x = self.out_proj(x)
+ return x
+
+
+class ConvolutionModule(nn.Module):
+ """ConvolutionModule in Zipformer model.
+ Modified from https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/zipformer/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, kernel_size
+
+ self.pointwise_conv1 = nn.Conv1d(
+ channels,
+ 2 * channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ bias=bias,
+ )
+
+ # after pointwise_conv1 we put x through a gated linear unit (nn.functional.glu).
+ # For most layers the normal rms value of channels of x seems to be in the range 1 to 4,
+ # but sometimes, for some reason, for layer 0 the rms ends up being very large,
+ # between 50 and 100 for different channels. This will cause very peaky and
+ # sparse derivatives for the sigmoid gating function, which will tend to make
+ # the loss function not learn effectively. (for most layers the average absolute values
+ # are in the range 0.5..9.0, and the average p(x>0), i.e. positive proportion,
+ # at the output of pointwise_conv1.output is around 0.35 to 0.45 for different
+ # layers, which likely breaks down as 0.5 for the "linear" half and
+ # 0.2 to 0.3 for the part that goes into the sigmoid. The idea is that if we
+ # constrain the rms values to a reasonable range via a constraint of max_abs=10.0,
+ # it will be in a better position to start learning something, i.e. to latch onto
+ # the correct range.
+ self.deriv_balancer1 = ActivationBalancer(
+ 2 * channels,
+ channel_dim=1,
+ max_abs=10.0,
+ min_positive=0.05,
+ max_positive=1.0,
+ )
+
+ # Will pad cached left context
+ self.lorder = kernel_size - 1
+ self.depthwise_conv = nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ stride=1,
+ padding=0,
+ groups=channels,
+ bias=bias,
+ )
+
+ self.deriv_balancer2 = ActivationBalancer(
+ channels,
+ channel_dim=1,
+ min_positive=0.05,
+ max_positive=1.0,
+ max_abs=20.0,
+ )
+
+ self.activation = DoubleSwish()
+
+ self.pointwise_conv2 = ScaledConv1d(
+ channels,
+ channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ bias=bias,
+ initial_scale=0.05,
+ )
+
+ def forward(
+ self,
+ x: Tensor,
+ src_key_padding_mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ """Compute convolution module.
+
+ Args:
+ x: Input tensor (#time, batch, channels).
+ src_key_padding_mask: the mask for the src keys per batch (optional):
+ (batch, #time), contains bool in masked positions.
+
+ Returns:
+ - 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 = self.deriv_balancer1(x)
+ x = nn.functional.glu(x, dim=1) # (batch, channels, time)
+
+ if src_key_padding_mask is not None:
+ x.masked_fill_(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
+
+ # 1D Depthwise Conv
+ # Make depthwise_conv causal by
+ # manualy padding self.lorder zeros to the left
+ x = nn.functional.pad(x, (self.lorder, 0), "constant", 0.0)
+ x = self.depthwise_conv(x)
+
+ x = self.deriv_balancer2(x)
+ x = self.activation(x)
+
+ x = self.pointwise_conv2(x) # (batch, channel, time)
+
+ return x.permute(2, 0, 1)
+
+ def streaming_forward(
+ self,
+ x: Tensor,
+ cache: Tensor,
+ ) -> Tuple[Tensor, Tensor]:
+ """Compute convolution module.
+
+ Args:
+ x: Input tensor (#time, batch, channels).
+ src_key_padding_mask: the mask for the src keys per batch:
+ (batch, #time), contains bool in masked positions.
+ cache: Cached left context for depthwise_conv, with shape of
+ (batch, channels, #kernel_size-1). Only used in real streaming decoding.
+
+ Returns:
+ A tuple of 2 tensors:
+ - Output tensor (#time, batch, channels).
+ - New cached left context, with shape of (batch, channels, #kernel_size-1).
+ """
+ # 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 = self.deriv_balancer1(x)
+ x = nn.functional.glu(x, dim=1) # (batch, channels, time)
+
+ # 1D Depthwise Conv
+ assert cache.shape == (x.size(0), x.size(1), self.lorder), (
+ cache.shape,
+ (x.size(0), x.size(1), self.lorder),
+ )
+ x = torch.cat([cache, x], dim=2)
+ # Update cache
+ cache = x[:, :, -self.lorder :]
+ x = self.depthwise_conv(x)
+
+ x = self.deriv_balancer2(x)
+ x = self.activation(x)
+
+ x = self.pointwise_conv2(x) # (batch, channel, time)
+
+ return x.permute(2, 0, 1), cache
+
+
+class Conv2dSubsampling(nn.Module):
+ """Convolutional 2D subsampling (to 1/4 length).
+
+ Convert an input of shape (N, T, idim) to an output
+ with shape (N, T', odim), where
+ T' = (T-3)//2 - 2 == (T-7)//2
+
+ It is based on
+ https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/transformer/subsampling.py # noqa
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ layer1_channels: int = 8,
+ layer2_channels: int = 32,
+ layer3_channels: int = 128,
+ dropout: float = 0.1,
+ ) -> None:
+ """
+ Args:
+ in_channels:
+ Number of channels in. The input shape is (N, T, in_channels).
+ Caution: It requires: T >=7, in_channels >=7
+ out_channels
+ Output dim. The output shape is (N, (T-7)//2, out_channels)
+ layer1_channels:
+ Number of channels in layer1
+ layer2_channels:
+ Number of channels in layer2
+ layer3_channels:
+ Number of channels in layer3
+ """
+ assert in_channels >= 7, in_channels
+ super().__init__()
+
+ self.conv = nn.Sequential(
+ nn.Conv2d(
+ in_channels=1,
+ out_channels=layer1_channels,
+ kernel_size=3,
+ padding=(0, 1), # (time, freq)
+ ),
+ ActivationBalancer(layer1_channels, channel_dim=1),
+ DoubleSwish(),
+ nn.Conv2d(
+ in_channels=layer1_channels,
+ out_channels=layer2_channels,
+ kernel_size=3,
+ stride=2,
+ padding=0,
+ ),
+ ActivationBalancer(layer2_channels, channel_dim=1),
+ DoubleSwish(),
+ nn.Conv2d(
+ in_channels=layer2_channels,
+ out_channels=layer3_channels,
+ kernel_size=3,
+ stride=(1, 2), # (time, freq)
+ ),
+ ActivationBalancer(layer3_channels, channel_dim=1),
+ DoubleSwish(),
+ )
+ out_height = (((in_channels - 1) // 2) - 1) // 2
+ self.out = ScaledLinear(out_height * layer3_channels, out_channels)
+ self.dropout = nn.Dropout(dropout)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """Subsample x.
+
+ Args:
+ x:
+ Its shape is (N, T, idim).
+
+ Returns:
+ Return a tensor of shape (N, (T-7)//2, odim)
+ """
+ # On entry, x is (N, T, idim)
+ x = x.unsqueeze(1) # (N, T, idim) -> (N, 1, T, idim) i.e., (N, C, H, W)
+ x = self.conv(x)
+ # Now x is of shape (N, odim, (T-7)//2, ((idim-1)//2 - 1)//2)
+ b, c, t, f = x.size()
+ x = self.out(x.transpose(1, 2).reshape(b, t, c * f))
+ # Now x is of shape (N, (T-7)//2, odim)
+ x = self.dropout(x)
+ return x
+
+
+def _test_zipformer_main():
+ feature_dim = 50
+ batch_size = 5
+ seq_len = 47
+ feature_dim = 50
+ # Just make sure the forward pass runs.
+
+ c = Zipformer(
+ num_features=feature_dim,
+ encoder_dims=(64, 96),
+ encoder_unmasked_dims=(48, 64),
+ nhead=(4, 4),
+ decode_chunk_size=4,
+ )
+ # Just make sure the forward pass runs.
+ f = c(
+ torch.randn(batch_size, seq_len, feature_dim),
+ torch.full((batch_size,), seq_len, dtype=torch.int64),
+ )
+ assert ((seq_len - 7) // 2 + 1) // 2 == f[0].shape[1], (seq_len, f.shape[1])
+ f[0].sum().backward()
+ c.eval()
+ f = c(
+ torch.randn(batch_size, seq_len, feature_dim),
+ torch.full((batch_size,), seq_len, dtype=torch.int64),
+ )
+ f # to remove flake8 warnings
+
+
+def _test_conv2d_subsampling():
+ num_features = 80
+ encoder_dims = 384
+ dropout = 0.1
+ encoder_embed = Conv2dSubsampling(num_features, encoder_dims, dropout=dropout)
+ for i in range(20, 40):
+ x = torch.rand(2, i, num_features)
+ y = encoder_embed(x)
+ assert (x.shape[1] - 7) // 2 == y.shape[1], (x.shape[1], y.shape[1])
+
+
+def _test_pooling_module():
+ N, S, C = 2, 12, 32
+ chunk_len = 4
+ m = PoolingModule(d_model=C)
+
+ # test chunk-wise forward with padding_mask
+ x = torch.randn(S, N, C)
+ y = m(x)
+ cached_len = torch.zeros(N, dtype=torch.int32)
+ cached_avg = torch.zeros(N, C)
+ for i in range(S // chunk_len):
+ start = i * chunk_len
+ end = start + chunk_len
+ x_chunk = x[start:end]
+ y_chunk, cached_len, cached_avg = m.streaming_forward(
+ x_chunk,
+ cached_len=cached_len,
+ cached_avg=cached_avg,
+ )
+ assert torch.allclose(y_chunk, y[start:end]), (y_chunk, y[start:end])
+
+
+def _test_state_stack_unstack():
+ m = Zipformer(
+ num_features=80,
+ encoder_dims=(64, 96),
+ encoder_unmasked_dims=(48, 64),
+ nhead=(4, 4),
+ zipformer_downsampling_factors=(4, 8),
+ num_left_chunks=2,
+ decode_chunk_size=8,
+ )
+ s1 = m.get_init_state()
+ s2 = m.get_init_state()
+ states = stack_states([s1, s2])
+ new_s1, new_s2 = unstack_states(states)
+ for i in range(m.num_encoders * 7):
+ for x, y in zip(s1[i], new_s1[i]):
+ assert torch.equal(x, y)
+ for x, y in zip(s2[i], new_s2[i]):
+ assert torch.equal(x, y)
+
+
+if __name__ == "__main__":
+ logging.getLogger().setLevel(logging.INFO)
+ torch.set_num_threads(1)
+ torch.set_num_interop_threads(1)
+ _test_zipformer_main()
+ _test_conv2d_subsampling()
+ _test_pooling_module()
+ _test_state_stack_unstack()
diff --git a/egs/librispeech/ASR/transducer_lstm/train.py b/egs/librispeech/ASR/transducer_lstm/train.py
index 792708bc0..a6f2bd08c 100755
--- a/egs/librispeech/ASR/transducer_lstm/train.py
+++ b/egs/librispeech/ASR/transducer_lstm/train.py
@@ -629,18 +629,8 @@ def run(rank, world_size, args):
# Keep only utterances with duration between 1 second and 20 seconds
return 1.0 <= c.duration <= 20.0
- num_in_total = len(train_cuts)
-
train_cuts = train_cuts.filter(remove_short_and_long_utt)
- num_left = len(train_cuts)
- num_removed = num_in_total - num_left
- removed_percent = num_removed / num_in_total * 100
-
- logging.info(f"Before removing short and long utterances: {num_in_total}")
- logging.info(f"After removing short and long utterances: {num_left}")
- logging.info(f"Removed {num_removed} utterances ({removed_percent:.5f}%)")
-
train_dl = librispeech.train_dataloaders(train_cuts)
valid_cuts = librispeech.dev_clean_cuts()
diff --git a/egs/librispeech/ASR/zipformer_mmi/README.md b/egs/librispeech/ASR/zipformer_mmi/README.md
new file mode 100644
index 000000000..e9a37a52a
--- /dev/null
+++ b/egs/librispeech/ASR/zipformer_mmi/README.md
@@ -0,0 +1,26 @@
+This recipe implements Zipformer-MMI model.
+
+See https://k2-fsa.github.io/icefall/recipes/Non-streaming-ASR/librispeech/zipformer_mmi.html for detailed tutorials.
+
+It uses **CTC loss for warm-up** and then switches to MMI loss during training.
+
+For decoding, it uses HP (H is ctc_topo, P is token-level bi-gram) as decoding graph. Supported decoding methods are:
+- **1best**. Extract the best path from the decoding lattice as the decoding result.
+- **nbest**. Extract n paths from the decoding lattice; the path with the highest score is the decoding result.
+- **nbest-rescoring-LG**. Extract n paths from the decoding lattice, rescore them with an word-level 3-gram LM, the path with the highest score is the decoding result.
+- **nbest-rescoring-3-gram**. Extract n paths from the decoding lattice, rescore them with an token-level 3-gram LM, the path with the highest score is the decoding result.
+- **nbest-rescoring-4-gram**. Extract n paths from the decoding lattice, rescore them with an token-level 4-gram LM, the path with the highest score is the decoding result.
+
+Experimental results training on train-clean-100 (epoch-30-avg-10):
+- 1best. 6.43 & 17.44
+- nbest, nbest-scale=1.2, 6.43 & 17.45
+- nbest-rescoring-LG, nbest-scale=1.2, 5.87 & 16.35
+- nbest-rescoring-3-gram, nbest-scale=1.2, 6.19 & 16.57
+- nbest-rescoring-4-gram, nbest-scale=1.2, 5.87 & 16.07
+
+Experimental results training on full librispeech (epoch-30-avg-10):
+- 1best. 2.54 & 5.65
+- nbest, nbest-scale=1.2, 2.54 & 5.66
+- nbest-rescoring-LG, nbest-scale=1.2, 2.49 & 5.42
+- nbest-rescoring-3-gram, nbest-scale=1.2, 2.52 & 5.62
+- nbest-rescoring-4-gram, nbest-scale=1.2, 2.5 & 5.51
diff --git a/egs/librispeech/ASR/zipformer_mmi/__init__.py b/egs/librispeech/ASR/zipformer_mmi/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/egs/librispeech/ASR/zipformer_mmi/asr_datamodule.py b/egs/librispeech/ASR/zipformer_mmi/asr_datamodule.py
new file mode 120000
index 000000000..a074d6085
--- /dev/null
+++ b/egs/librispeech/ASR/zipformer_mmi/asr_datamodule.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/asr_datamodule.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/zipformer_mmi/decode.py b/egs/librispeech/ASR/zipformer_mmi/decode.py
new file mode 100755
index 000000000..7d0ea78bb
--- /dev/null
+++ b/egs/librispeech/ASR/zipformer_mmi/decode.py
@@ -0,0 +1,736 @@
+#!/usr/bin/env python3
+#
+# Copyright 2021-2022 Xiaomi Corporation (Author: Fangjun Kuang,
+# Liyong Guo,
+# Zengwei Yao)
+#
+# 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.
+"""
+Usage:
+(1) 1best
+./zipformer_mmi/mmi_decode.py \
+ --epoch 30 \
+ --avg 15 \
+ --exp-dir ./zipformer_mmi/exp \
+ --max-duration 100 \
+ --decoding-method 1best
+(2) nbest
+./zipformer_mmi/mmi_decode.py \
+ --epoch 30 \
+ --avg 15 \
+ --exp-dir ./zipformer_mmi/exp \
+ --max-duration 100 \
+ --nbest-scale 1.0 \
+ --decoding-method nbest
+(3) nbest-rescoring-LG
+./zipformer_mmi/mmi_decode.py \
+ --epoch 30 \
+ --avg 15 \
+ --exp-dir ./zipformer_mmi/exp \
+ --max-duration 100 \
+ --nbest-scale 1.0 \
+ --decoding-method nbest-rescoring-LG
+(4) nbest-rescoring-3-gram
+./zipformer_mmi/mmi_decode.py \
+ --epoch 30 \
+ --avg 15 \
+ --exp-dir ./zipformer_mmi/exp \
+ --max-duration 100 \
+ --nbest-scale 1.0 \
+ --decoding-method nbest-rescoring-3-gram
+(5) nbest-rescoring-4-gram
+./zipformer_mmi/mmi_decode.py \
+ --epoch 30 \
+ --avg 15 \
+ --exp-dir ./zipformer_mmi/exp \
+ --max-duration 100 \
+ --nbest-scale 1.0 \
+ --decoding-method nbest-rescoring-4-gram
+"""
+
+
+import argparse
+import logging
+import math
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import k2
+import sentencepiece as spm
+import torch
+import torch.nn as nn
+from asr_datamodule import LibriSpeechAsrDataModule
+from train import add_model_arguments, get_ctc_model, get_params
+
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from icefall.decode import (
+ get_lattice,
+ nbest_decoding,
+ nbest_rescore_with_LM,
+ one_best_decoding,
+)
+from icefall.lexicon import Lexicon
+from icefall.mmi_graph_compiler import MmiTrainingGraphCompiler
+from icefall.utils import (
+ AttributeDict,
+ get_texts,
+ setup_logger,
+ store_transcripts,
+ str2bool,
+ write_error_stats,
+)
+
+LOG_EPS = math.log(1e-10)
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=30,
+ help="""It specifies the checkpoint to use for decoding.
+ Note: Epoch counts from 1.
+ You can specify --avg to use more checkpoints for model averaging.""",
+ )
+
+ parser.add_argument(
+ "--iter",
+ type=int,
+ default=0,
+ help="""If positive, --epoch is ignored and it
+ will use the checkpoint exp_dir/checkpoint-iter.pt.
+ You can specify --avg to use more checkpoints for model averaging.
+ """,
+ )
+
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=15,
+ help="Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch' and '--iter'",
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=True,
+ help="Whether to load averaged model. Currently it only supports "
+ "using --epoch. If True, it would decode with the averaged model "
+ "over the epoch range from `epoch-avg` (excluded) to `epoch`."
+ "Actually only the models with epoch number of `epoch-avg` and "
+ "`epoch` are loaded for averaging. ",
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="zipformer_mmi/exp",
+ help="The experiment dir",
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ default="data/lang_bpe_500/bpe.model",
+ help="Path to the BPE model",
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=Path,
+ default="data/lang_bpe_500",
+ help="The lang dir containing word table and LG graph",
+ )
+
+ parser.add_argument(
+ "--decoding-method",
+ type=str,
+ default="1best",
+ help="""Decoding method. Use HP as decoding graph, where H is
+ ctc_topo and P is token-level bi-gram lm.
+ Supported values are:
+ - (1) 1best. Extract the best path from the decoding lattice as the
+ decoding result.
+ - (2) nbest. Extract n paths from the decoding lattice; the path
+ with the highest score is the decoding result.
+ - (4) nbest-rescoring-LG. Extract n paths from the decoding lattice,
+ rescore them with an word-level 3-gram LM, the path with the
+ highest score is the decoding result.
+ - (5) nbest-rescoring-3-gram. Extract n paths from the decoding
+ lattice, rescore them with an token-level 3-gram LM, the path with
+ the highest score is the decoding result.
+ - (6) nbest-rescoring-4-gram. Extract n paths from the decoding
+ lattice, rescore them with an token-level 4-gram LM, the path with
+ the highest score is the decoding result.
+ """,
+ )
+
+ parser.add_argument(
+ "--num-paths",
+ type=int,
+ default=100,
+ help="""Number of paths for n-best based decoding method.
+ Used only when "method" is one of the following values:
+ nbest, nbest-rescoring, and nbest-oracle
+ """,
+ )
+
+ parser.add_argument(
+ "--nbest-scale",
+ type=float,
+ default=1.0,
+ help="""The scale to be applied to `lattice.scores`.
+ It's needed if you use any kinds of n-best based rescoring.
+ Used only when "method" is one of the following values:
+ nbest, nbest-rescoring, and nbest-oracle
+ A smaller value results in more unique paths.
+ """,
+ )
+
+ parser.add_argument(
+ "--hp-scale",
+ type=float,
+ default=1.0,
+ help="""The scale to be applied to `ctc_topo_P.scores`.
+ """,
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def get_decoding_params() -> AttributeDict:
+ """Parameters for decoding."""
+ params = AttributeDict(
+ {
+ "frame_shift_ms": 10,
+ "search_beam": 20,
+ "output_beam": 8,
+ "min_active_states": 30,
+ "max_active_states": 10000,
+ "use_double_scores": True,
+ }
+ )
+ return params
+
+
+def decode_one_batch(
+ params: AttributeDict,
+ model: nn.Module,
+ HP: Optional[k2.Fsa],
+ bpe_model: Optional[spm.SentencePieceProcessor],
+ batch: dict,
+ G: Optional[k2.Fsa] = None,
+ LG: Optional[k2.Fsa] = None,
+) -> Dict[str, List[List[str]]]:
+ """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.decoding_method is "1best", it uses 1best decoding without LM rescoring.
+ - params.decoding_method is "nbest", it uses nbest decoding without LM rescoring.
+ - params.decoding_method is "nbest-rescoring-LG", it uses nbest rescoring with word-level 3-gram LM.
+ - params.decoding_method is "nbest-rescoring-3-gram", it uses nbest rescoring with token-level 3-gram LM.
+ - params.decoding_method is "nbest-rescoring-4-gram", it uses nbest rescoring with token-level 4-gram LM.
+
+ model:
+ The neural model.
+ HP:
+ The decoding graph. H is ctc_topo, P is token-level bi-gram LM.
+ bpe_model:
+ The BPE model.
+ batch:
+ It is the return value from iterating
+ `lhotse.dataset.K2SpeechRecognitionDataset`. See its documentation
+ for the format of the `batch`.
+ LG:
+ An LM. L is the lexicon, G is a word-level 3-gram LM.
+ It is used when params.decoding_method is "nbest-rescoring-LG".
+ G:
+ An LM. L is the lexicon, G is a token-level 3-gram or 4-gram LM.
+ It is used when params.decoding_method is "nbest-rescoring-3-gram"
+ or "nbest-rescoring-4-gram".
+ Returns:
+ Return the decoding result. See above description for the format of
+ the returned dict. Note: If it decodes to nothing, then return None.
+ """
+ device = HP.device
+ feature = batch["inputs"]
+ assert feature.ndim == 3, feature.shape
+ feature = feature.to(device)
+
+ # at entry, feature is (N, T, C)
+
+ supervisions = batch["supervisions"]
+ feature_lens = supervisions["num_frames"].to(device)
+
+ nnet_output, encoder_out_lens = model(x=feature, x_lens=feature_lens)
+ # 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,
+ decoding_graph=HP,
+ 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,
+ )
+
+ method = params.decoding_method
+
+ if method in ["1best", "nbest"]:
+ if 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,
+ nbest_scale=params.nbest_scale,
+ )
+ key = f"no_rescore-nbest-scale-{params.nbest_scale}-{params.num_paths}" # noqa
+
+ # Note: `best_path.aux_labels` contains token IDs, not word IDs
+ # since we are using HP, not HLG here.
+ #
+ # token_ids is a lit-of-list of IDs
+ token_ids = get_texts(best_path)
+ # hyps is a list of str, e.g., ['xxx yyy zzz', ...]
+ hyps = bpe_model.decode(token_ids)
+ # hyps is a list of list of str, e.g., [['xxx', 'yyy', 'zzz'], ... ]
+ hyps = [s.split() for s in hyps]
+ return {key: hyps}
+
+ assert method in [
+ "nbest-rescoring-LG", # word-level 3-gram lm
+ "nbest-rescoring-3-gram", # token-level 3-gram lm
+ "nbest-rescoring-4-gram", # token-level 4-gram lm
+ ]
+
+ lm_scale_list = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]
+ 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 method == "nbest-rescoring-LG":
+ assert LG is not None
+ LM = LG
+ else:
+ assert G is not None
+ LM = G
+ best_path_dict = nbest_rescore_with_LM(
+ lattice=lattice,
+ LM=LM,
+ num_paths=params.num_paths,
+ lm_scale_list=lm_scale_list,
+ nbest_scale=params.nbest_scale,
+ )
+
+ ans = dict()
+ suffix = f"-nbest-scale-{params.nbest_scale}-{params.num_paths}"
+ for lm_scale_str, best_path in best_path_dict.items():
+ token_ids = get_texts(best_path)
+ # hyps is a list of str, e.g., ['xxx yyy zzz', ...]
+ hyps = bpe_model.decode(token_ids)
+ # hyps is a list of list of str, e.g., [['xxx', 'yyy', 'zzz'], ... ]
+ hyps = [s.split() for s in hyps]
+ ans[lm_scale_str + suffix] = hyps
+ return ans
+
+
+def decode_dataset(
+ dl: torch.utils.data.DataLoader,
+ params: AttributeDict,
+ model: nn.Module,
+ HP: k2.Fsa,
+ bpe_model: spm.SentencePieceProcessor,
+ G: Optional[k2.Fsa] = None,
+ LG: Optional[k2.Fsa] = None,
+) -> Dict[str, List[Tuple[str, List[str], List[str]]]]:
+ """Decode dataset.
+
+ Args:
+ dl:
+ PyTorch's dataloader containing the dataset to decode.
+ params:
+ It is returned by :func:`get_params`.
+ model:
+ The neural model.
+ HP:
+ The decoding graph. H is ctc_topo, P is token-level bi-gram LM.
+ bpe_model:
+ The BPE model.
+ LG:
+ An LM. L is the lexicon, G is a word-level 3-gram LM.
+ It is used when params.decoding_method is "nbest-rescoring-LG".
+ G:
+ An LM. L is the lexicon, G is a token-level 3-gram or 4-gram LM.
+ It is used when params.decoding_method is "nbest-rescoring-3-gram"
+ or "nbest-rescoring-4-gram".
+
+ 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.
+ """
+ num_cuts = 0
+
+ try:
+ num_batches = len(dl)
+ except TypeError:
+ num_batches = "?"
+
+ results = defaultdict(list)
+ for batch_idx, batch in enumerate(dl):
+ texts = batch["supervisions"]["text"]
+ cut_ids = [cut.id for cut in batch["supervisions"]["cut"]]
+
+ hyps_dict = decode_one_batch(
+ params=params,
+ model=model,
+ HP=HP,
+ bpe_model=bpe_model,
+ batch=batch,
+ G=G,
+ LG=LG,
+ )
+
+ for name, hyps in hyps_dict.items():
+ this_batch = []
+ assert len(hyps) == len(texts)
+ for cut_id, hyp_words, ref_text in zip(cut_ids, hyps, texts):
+ ref_words = ref_text.split()
+ this_batch.append((cut_id, ref_words, hyp_words))
+
+ results[name].extend(this_batch)
+
+ num_cuts += len(texts)
+
+ if batch_idx % 100 == 0:
+ batch_str = f"{batch_idx}/{num_batches}"
+
+ logging.info(f"batch {batch_str}, cuts processed until now is {num_cuts}")
+ return results
+
+
+def save_results(
+ params: AttributeDict,
+ test_set_name: str,
+ results_dict: Dict[str, List[Tuple[str, List[str], List[str]]]],
+):
+ test_set_wers = dict()
+ for key, results in results_dict.items():
+ recog_path = (
+ params.res_dir / f"recogs-{test_set_name}-{key}-{params.suffix}.txt"
+ )
+ results = sorted(results)
+ 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.res_dir / f"errs-{test_set_name}-{key}-{params.suffix}.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.res_dir / f"wer-summary-{test_set_name}-{key}-{params.suffix}.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()
+ args.exp_dir = Path(args.exp_dir)
+ args.lang_dir = Path(args.lang_dir)
+
+ params = get_params()
+ # add decoding params
+ params.update(get_decoding_params())
+ params.update(vars(args))
+
+ assert params.decoding_method in (
+ "1best",
+ "nbest",
+ "nbest-rescoring-LG", # word-level 3-gram lm
+ "nbest-rescoring-3-gram", # token-level 3-gram lm
+ "nbest-rescoring-4-gram", # token-level 4-gram lm
+ ), params.decoding_method
+ params.res_dir = params.exp_dir / params.decoding_method
+
+ if params.iter > 0:
+ params.suffix = f"iter-{params.iter}-avg-{params.avg}"
+ else:
+ params.suffix = f"epoch-{params.epoch}-avg-{params.avg}"
+
+ if params.use_averaged_model:
+ params.suffix += "-use-averaged-model"
+
+ setup_logger(f"{params.res_dir}/log-decode-{params.suffix}")
+ logging.info("decoding started")
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"device: {device}")
+ logging.info(params)
+
+ lexicon = Lexicon(params.lang_dir)
+ max_token_id = max(lexicon.tokens)
+ num_classes = max_token_id + 1 # +1 for the blank
+
+ params.vocab_size = num_classes
+ # and are defined in local/train_bpe_model.py
+ params.blank_id = 0
+
+ bpe_model = spm.SentencePieceProcessor()
+ bpe_model.load(str(params.lang_dir / "bpe.model"))
+ mmi_graph_compiler = MmiTrainingGraphCompiler(
+ params.lang_dir,
+ uniq_filename="lexicon.txt",
+ device=device,
+ oov="",
+ sos_id=1,
+ eos_id=1,
+ )
+ HP = mmi_graph_compiler.ctc_topo_P
+ HP.scores *= params.hp_scale
+ if not hasattr(HP, "lm_scores"):
+ HP.lm_scores = HP.scores.clone()
+
+ LG = None
+ G = None
+
+ if params.decoding_method == "nbest-rescoring-LG":
+ lg_filename = params.lang_dir / "LG.pt"
+ logging.info(f"Loading {lg_filename}")
+ LG = k2.Fsa.from_dict(torch.load(lg_filename, map_location=device))
+ LG = k2.Fsa.from_fsas([LG]).to(device)
+ LG.lm_scores = LG.scores.clone()
+
+ elif params.decoding_method in ["nbest-rescoring-3-gram", "nbest-rescoring-4-gram"]:
+ order = params.decoding_method[-6]
+ assert order in ("3", "4"), (params.decoding_method, order)
+ order = int(order)
+ if not (params.lang_dir / f"{order}gram.pt").is_file():
+ logging.info(f"Loading {order}gram.fst.txt")
+ logging.warning("It may take a few minutes.")
+ with open(params.lang_dir / f"{order}gram.fst.txt") as f:
+ first_token_disambig_id = lexicon.token_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_token_disambig_id] = 0
+ G = k2.Fsa.from_fsas([G]).to(device)
+ # G = k2.remove_epsilon(G)
+ G = k2.arc_sort(G)
+ # Save a dummy value so that it can be loaded in C++.
+ # See https://github.com/pytorch/pytorch/issues/67902
+ # for why we need to do this.
+ G.dummy = 1
+
+ torch.save(G.as_dict(), params.lang_dir / f"{order}gram.pt")
+ else:
+ logging.info(f"Loading pre-compiled {order}gram.pt")
+ d = torch.load(params.lang_dir / f"{order}gram.pt", map_location=device)
+ G = k2.Fsa.from_dict(d)
+
+ G.lm_scores = G.scores.clone()
+
+ logging.info("About to create model")
+ model = get_ctc_model(params)
+
+ if not params.use_averaged_model:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ elif 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 i >= 1:
+ filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg + 1
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg + 1:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ filename_start = filenames[-1]
+ filename_end = filenames[0]
+ logging.info(
+ "Calculating the averaged model over iteration checkpoints"
+ f" from {filename_start} (excluded) to {filename_end}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+ else:
+ assert params.avg > 0, params.avg
+ start = params.epoch - params.avg
+ assert start >= 1, start
+ filename_start = f"{params.exp_dir}/epoch-{start}.pt"
+ filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
+ logging.info(
+ f"Calculating the averaged model over epoch range from "
+ f"{start} (excluded) to {params.epoch}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+
+ model.to(device)
+ model.eval()
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ # we need cut ids to display recognition results.
+ args.return_cuts = True
+ librispeech = LibriSpeechAsrDataModule(args)
+
+ test_clean_cuts = librispeech.test_clean_cuts()
+ test_other_cuts = librispeech.test_other_cuts()
+
+ test_clean_dl = librispeech.test_dataloaders(test_clean_cuts)
+ test_other_dl = librispeech.test_dataloaders(test_other_cuts)
+
+ test_sets = ["test-clean", "test-other"]
+ test_dl = [test_clean_dl, test_other_dl]
+
+ for test_set, test_dl in zip(test_sets, test_dl):
+ results_dict = decode_dataset(
+ dl=test_dl,
+ params=params,
+ model=model,
+ HP=HP,
+ bpe_model=bpe_model,
+ G=G,
+ LG=LG,
+ )
+
+ 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/zipformer_mmi/encoder_interface.py b/egs/librispeech/ASR/zipformer_mmi/encoder_interface.py
new file mode 120000
index 000000000..b9aa0ae08
--- /dev/null
+++ b/egs/librispeech/ASR/zipformer_mmi/encoder_interface.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/encoder_interface.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/zipformer_mmi/export.py b/egs/librispeech/ASR/zipformer_mmi/export.py
new file mode 100755
index 000000000..0af7bd367
--- /dev/null
+++ b/egs/librispeech/ASR/zipformer_mmi/export.py
@@ -0,0 +1,307 @@
+#!/usr/bin/env python3
+#
+# Copyright 2021 Xiaomi Corporation (Author: Fangjun Kuang)
+#
+# 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.
+
+# This script converts several saved checkpoints
+# to a single one using model averaging.
+"""
+
+Usage:
+
+(1) Export to torchscript model using torch.jit.script()
+
+./zipformer_mmi/export.py \
+ --exp-dir ./zipformer_mmi/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 9 \
+ --jit 1
+
+It will generate a file `cpu_jit.pt` in the given `exp_dir`. You can later
+load it by `torch.jit.load("cpu_jit.pt")`.
+
+Note `cpu` in the name `cpu_jit.pt` means the parameters when loaded into Python
+are on CPU. You can use `to("cuda")` to move them to a CUDA device.
+
+Check
+https://github.com/k2-fsa/sherpa
+for how to use the exported models outside of icefall.
+
+(2) Export `model.state_dict()`
+
+./zipformer_mmi/export.py \
+ --exp-dir ./zipformer_mmi/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 20 \
+ --avg 10
+
+It will generate a file `pretrained.pt` in the given `exp_dir`. You can later
+load it by `icefall.checkpoint.load_checkpoint()`.
+
+To use the generated file with `zipformer_mmi/decode.py`,
+you can do:
+
+ cd /path/to/exp_dir
+ ln -s pretrained.pt epoch-9999.pt
+
+ cd /path/to/egs/librispeech/ASR
+ ./zipformer_mmi/decode.py \
+ --exp-dir ./zipformer_mmi/exp \
+ --epoch 9999 \
+ --avg 1 \
+ --max-duration 600 \
+ --decoding-method greedy_search \
+ --bpe-model data/lang_bpe_500/bpe.model
+
+Check ./pretrained.py for its usage.
+
+Note: If you don't want to train a model from scratch, we have
+provided one for you. You can get it at
+
+https://huggingface.co/Zengwei/icefall-asr-librispeech-zipformer-mmi-2022-12-08
+
+with the following commands:
+
+ sudo apt-get install git-lfs
+ git lfs install
+ git clone https://huggingface.co/Zengwei/icefall-asr-librispeech-zipformer-mmi-2022-12-08
+ # You will find the pre-trained model in icefall-asr-librispeech-zipformer-mmi-2022-12-08/exp
+"""
+
+import argparse
+import logging
+from pathlib import Path
+
+import sentencepiece as spm
+import torch
+from scaling_converter import convert_scaled_to_non_scaled
+from train import add_model_arguments, get_ctc_model, get_params
+
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from icefall.utils import str2bool
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=30,
+ help="""It specifies the checkpoint to use for decoding.
+ Note: Epoch counts from 1.
+ You can specify --avg to use more checkpoints for model averaging.""",
+ )
+
+ parser.add_argument(
+ "--iter",
+ type=int,
+ default=0,
+ help="""If positive, --epoch is ignored and it
+ will use the checkpoint exp_dir/checkpoint-iter.pt.
+ You can specify --avg to use more checkpoints for model averaging.
+ """,
+ )
+
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=9,
+ help="Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch' and '--iter'",
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=True,
+ help="Whether to load averaged model. Currently it only supports "
+ "using --epoch. If True, it would decode with the averaged model "
+ "over the epoch range from `epoch-avg` (excluded) to `epoch`."
+ "Actually only the models with epoch number of `epoch-avg` and "
+ "`epoch` are loaded for averaging. ",
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="zipformer_mmi/exp",
+ help="""It specifies the directory where all training related
+ files, e.g., checkpoints, log, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ default="data/lang_bpe_500/bpe.model",
+ help="Path to the BPE model",
+ )
+
+ parser.add_argument(
+ "--jit",
+ type=str2bool,
+ default=False,
+ help="""True to save a model after applying torch.jit.script.
+ It will generate a file named cpu_jit.pt
+
+ Check ./jit_pretrained.py for how to use it.
+ """,
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+@torch.no_grad()
+def main():
+ args = get_parser().parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ params = get_params()
+ params.update(vars(args))
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"device: {device}")
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(params.bpe_model)
+
+ # is defined in local/train_bpe_model.py
+ params.blank_id = sp.piece_to_id("")
+ params.vocab_size = sp.get_piece_size()
+
+ logging.info(params)
+
+ logging.info("About to create model")
+ model = get_ctc_model(params)
+
+ model.to(device)
+
+ if not params.use_averaged_model:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ elif 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 i >= 1:
+ filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg + 1
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg + 1:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ filename_start = filenames[-1]
+ filename_end = filenames[0]
+ logging.info(
+ "Calculating the averaged model over iteration checkpoints"
+ f" from {filename_start} (excluded) to {filename_end}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+ else:
+ assert params.avg > 0, params.avg
+ start = params.epoch - params.avg
+ assert start >= 1, start
+ filename_start = f"{params.exp_dir}/epoch-{start}.pt"
+ filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
+ logging.info(
+ f"Calculating the averaged model over epoch range from "
+ f"{start} (excluded) to {params.epoch}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+
+ model.to("cpu")
+ model.eval()
+
+ if params.jit is True:
+ convert_scaled_to_non_scaled(model, inplace=True)
+ logging.info("Using torch.jit.script()")
+ model = torch.jit.script(model)
+ filename = params.exp_dir / "cpu_jit.pt"
+ model.save(str(filename))
+ logging.info(f"Saved to {filename}")
+ else:
+ logging.info("Not using torchscript. Export model.state_dict()")
+ # Save it using a format so that it can be loaded
+ # by :func:`load_checkpoint`
+ filename = params.exp_dir / "pretrained.pt"
+ torch.save({"model": model.state_dict()}, str(filename))
+ logging.info(f"Saved to {filename}")
+
+
+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/zipformer_mmi/jit_pretrained.py b/egs/librispeech/ASR/zipformer_mmi/jit_pretrained.py
new file mode 100755
index 000000000..c9ef16ffa
--- /dev/null
+++ b/egs/librispeech/ASR/zipformer_mmi/jit_pretrained.py
@@ -0,0 +1,391 @@
+#!/usr/bin/env python3
+# Copyright 2021-2022 Xiaomi Corp. (authors: Fangjun Kuang,
+# Zengwei)
+#
+# 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.
+"""
+This script loads torchscript models, exported by `torch.jit.script()`
+and uses them to decode waves.
+You can use the following command to get the exported models:
+
+./zipformer_mmi/export.py \
+ --exp-dir ./zipformer_mmi/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 20 \
+ --avg 10 \
+ --jit 1
+
+Usage of this script:
+
+(1) 1best
+./zipformer_mmi/jit_pretrained.py \
+ --nn-model-filename ./zipformer_mmi/exp/cpu_jit.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method 1best \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+(2) nbest
+./zipformer_mmi/jit_pretrained.py \
+ --nn-model-filename ./zipformer_mmi/exp/cpu_jit.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --nbest-scale 1.2 \
+ --method nbest \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+(3) nbest-rescoring-LG
+./zipformer_mmi/jit_pretrained.py \
+ --nn-model-filename ./zipformer_mmi/exp/cpu_jit.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --nbest-scale 1.2 \
+ --method nbest-rescoring-LG \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+(4) nbest-rescoring-3-gram
+./zipformer_mmi/jit_pretrained.py \
+ --nn-model-filename ./zipformer_mmi/exp/cpu_jit.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --nbest-scale 1.2 \
+ --method nbest-rescoring-3-gram \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+(5) nbest-rescoring-4-gram
+./zipformer_mmi/jit_pretrained.py \
+ --nn-model-filename ./zipformer_mmi/exp/cpu_jit.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --nbest-scale 1.2 \
+ --method nbest-rescoring-4-gram \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+"""
+
+import argparse
+import logging
+import math
+from pathlib import Path
+from typing import List
+
+import k2
+import kaldifeat
+import sentencepiece as spm
+import torch
+import torchaudio
+from decode import get_decoding_params
+from torch.nn.utils.rnn import pad_sequence
+from train import get_params
+
+from icefall.decode import (
+ get_lattice,
+ nbest_decoding,
+ nbest_rescore_with_LM,
+ one_best_decoding,
+)
+from icefall.mmi_graph_compiler import MmiTrainingGraphCompiler
+from icefall.utils import get_texts
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--nn-model-filename",
+ type=str,
+ required=True,
+ help="Path to the torchscript model cpu_jit.pt",
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ help="""Path to bpe.model.""",
+ )
+
+ parser.add_argument(
+ "--method",
+ type=str,
+ default="1best",
+ help="""Decoding method. Use HP as decoding graph, where H is
+ ctc_topo and P is token-level bi-gram lm.
+ Supported values are:
+ - (1) 1best. Extract the best path from the decoding lattice as the
+ decoding result.
+ - (2) nbest. Extract n paths from the decoding lattice; the path
+ with the highest score is the decoding result.
+ - (4) nbest-rescoring-LG. Extract n paths from the decoding lattice,
+ rescore them with an word-level 3-gram LM, the path with the
+ highest score is the decoding result.
+ - (5) nbest-rescoring-3-gram. Extract n paths from the decoding
+ lattice, rescore them with an token-level 3-gram LM, the path with
+ the highest score is the decoding result.
+ - (6) nbest-rescoring-4-gram. Extract n paths from the decoding
+ lattice, rescore them with an token-level 4-gram LM, the path with
+ the highest score is the decoding result.
+ """,
+ )
+
+ parser.add_argument(
+ "--sample-rate",
+ type=int,
+ default=16000,
+ help="The sample rate of the input sound file",
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=Path,
+ default="data/lang_bpe_500",
+ help="The lang dir containing word table and LG graph",
+ )
+
+ parser.add_argument(
+ "--num-paths",
+ type=int,
+ default=100,
+ help="""Number of paths for n-best based decoding method.
+ Used only when "method" is one of the following values:
+ nbest, nbest-rescoring, and nbest-oracle
+ """,
+ )
+
+ parser.add_argument(
+ "--nbest-scale",
+ type=float,
+ default=1.2,
+ help="""The scale to be applied to `lattice.scores`.
+ It's needed if you use any kinds of n-best based rescoring.
+ Used only when "method" is one of the following values:
+ nbest, nbest-rescoring, and nbest-oracle
+ A smaller value results in more unique paths.
+ """,
+ )
+
+ parser.add_argument(
+ "--ngram-lm-scale",
+ type=float,
+ default=0.1,
+ help="""
+ Used when method is nbest-rescoring-LG, nbest-rescoring-3-gram,
+ and nbest-rescoring-4-gram.
+ It specifies the scale for n-gram LM scores.
+ (Note: You need to tune it on a dataset.)
+ """,
+ )
+
+ parser.add_argument(
+ "--hp-scale",
+ type=float,
+ default=1.0,
+ help="""The scale to be applied to `ctc_topo_P.scores`.
+ """,
+ )
+
+ parser.add_argument(
+ "sound_files",
+ type=str,
+ nargs="+",
+ help="The input sound file(s) to transcribe. "
+ "Supported formats are those supported by torchaudio.load(). "
+ "For example, wav and flac are supported. "
+ "The sample rate has to be 16kHz.",
+ )
+
+ return parser
+
+
+def read_sound_files(
+ filenames: List[str], expected_sample_rate: float = 16000
+) -> List[torch.Tensor]:
+ """Read a list of sound files into a list 1-D float32 torch tensors.
+ Args:
+ filenames:
+ A list of sound filenames.
+ expected_sample_rate:
+ The expected sample rate of the sound files.
+ Returns:
+ Return a list of 1-D float32 torch tensors.
+ """
+ ans = []
+ for f in filenames:
+ wave, sample_rate = torchaudio.load(f)
+ assert (
+ sample_rate == expected_sample_rate
+ ), f"expected sample rate: {expected_sample_rate}. Given: {sample_rate}"
+ # We use only the first channel
+ ans.append(wave[0])
+ return ans
+
+
+@torch.no_grad()
+def main():
+ parser = get_parser()
+ args = parser.parse_args()
+ logging.info(vars(args))
+
+ params = get_params()
+ # add decoding params
+ params.update(get_decoding_params())
+ params.update(vars(args))
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"device: {device}")
+
+ model = torch.jit.load(params.nn_model_filename)
+ model.eval()
+ model.to(device)
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(args.bpe_model)
+
+ logging.info("Constructing Fbank computer")
+ opts = kaldifeat.FbankOptions()
+ opts.device = device
+ opts.frame_opts.dither = 0
+ opts.frame_opts.snip_edges = False
+ opts.frame_opts.samp_freq = 16000
+ opts.mel_opts.num_bins = 80
+
+ fbank = kaldifeat.Fbank(opts)
+
+ logging.info(f"Reading sound files: {args.sound_files}")
+ waves = read_sound_files(
+ filenames=params.sound_files, expected_sample_rate=params.sample_rate
+ )
+ waves = [w.to(device) for w in waves]
+
+ logging.info("Decoding started")
+ features = fbank(waves)
+ feature_lengths = [f.size(0) for f in features]
+
+ features = pad_sequence(
+ features,
+ batch_first=True,
+ padding_value=math.log(1e-10),
+ )
+ feature_lengths = torch.tensor(feature_lengths, device=device)
+
+ bpe_model = spm.SentencePieceProcessor()
+ bpe_model.load(str(params.lang_dir / "bpe.model"))
+ mmi_graph_compiler = MmiTrainingGraphCompiler(
+ params.lang_dir,
+ uniq_filename="lexicon.txt",
+ device=device,
+ oov="",
+ sos_id=1,
+ eos_id=1,
+ )
+ HP = mmi_graph_compiler.ctc_topo_P
+ HP.scores *= params.hp_scale
+ if not hasattr(HP, "lm_scores"):
+ HP.lm_scores = HP.scores.clone()
+
+ method = params.method
+ assert method in (
+ "1best",
+ "nbest",
+ "nbest-rescoring-LG", # word-level 3-gram lm
+ "nbest-rescoring-3-gram", # token-level 3-gram lm
+ "nbest-rescoring-4-gram", # token-level 4-gram lm
+ )
+ # loading language model for rescoring
+ LM = None
+ if method == "nbest-rescoring-LG":
+ lg_filename = params.lang_dir / "LG.pt"
+ logging.info(f"Loading {lg_filename}")
+ LG = k2.Fsa.from_dict(torch.load(lg_filename, map_location=device))
+ LG = k2.Fsa.from_fsas([LG]).to(device)
+ LG.lm_scores = LG.scores.clone()
+ LM = LG
+ elif method in ["nbest-rescoring-3-gram", "nbest-rescoring-4-gram"]:
+ order = method[-6]
+ assert order in ("3", "4")
+ order = int(order)
+ logging.info(f"Loading pre-compiled {order}gram.pt")
+ d = torch.load(params.lang_dir / f"{order}gram.pt", map_location=device)
+ G = k2.Fsa.from_dict(d)
+ G.lm_scores = G.scores.clone()
+ LM = G
+
+ # Encoder forward
+ nnet_output, encoder_out_lens = model(x=features, x_lens=feature_lengths)
+
+ batch_size = nnet_output.shape[0]
+ supervision_segments = torch.tensor(
+ [
+ [i, 0, feature_lengths[i] // params.subsampling_factor]
+ for i in range(batch_size)
+ ],
+ dtype=torch.int32,
+ )
+
+ lattice = get_lattice(
+ nnet_output=nnet_output,
+ decoding_graph=HP,
+ 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 method in ["1best", "nbest"]:
+ if method == "1best":
+ best_path = one_best_decoding(
+ lattice=lattice, use_double_scores=params.use_double_scores
+ )
+ else:
+ best_path = nbest_decoding(
+ lattice=lattice,
+ num_paths=params.num_paths,
+ use_double_scores=params.use_double_scores,
+ nbest_scale=params.nbest_scale,
+ )
+ else:
+ best_path_dict = nbest_rescore_with_LM(
+ lattice=lattice,
+ LM=LM,
+ num_paths=params.num_paths,
+ lm_scale_list=[params.ngram_lm_scale],
+ nbest_scale=params.nbest_scale,
+ )
+ best_path = next(iter(best_path_dict.values()))
+
+ # Note: `best_path.aux_labels` contains token IDs, not word IDs
+ # since we are using HP, not HLG here.
+ #
+ # token_ids is a lit-of-list of IDs
+ token_ids = get_texts(best_path)
+ # hyps is a list of str, e.g., ['xxx yyy zzz', ...]
+ hyps = bpe_model.decode(token_ids)
+ # hyps is a list of list of str, e.g., [['xxx', 'yyy', 'zzz'], ... ]
+ hyps = [s.split() for s in hyps]
+ s = "\n"
+ for filename, hyp in zip(params.sound_files, hyps):
+ words = " ".join(hyp)
+ s += f"{filename}:\n{words}\n\n"
+ logging.info(s)
+
+ logging.info("Decoding Done")
+
+
+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/zipformer_mmi/model.py b/egs/librispeech/ASR/zipformer_mmi/model.py
new file mode 100644
index 000000000..4045c8b64
--- /dev/null
+++ b/egs/librispeech/ASR/zipformer_mmi/model.py
@@ -0,0 +1,75 @@
+# Copyright 2022 Xiaomi Corp. (authors: Zengwei Yao)
+#
+# 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.
+
+
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+from encoder_interface import EncoderInterface
+
+
+class CTCModel(nn.Module):
+ def __init__(
+ self,
+ encoder: EncoderInterface,
+ encoder_dim: int,
+ vocab_size: int,
+ ):
+ """
+ Args:
+ encoder:
+ It is the transcription network in the paper. Its accepts
+ two inputs: `x` of (N, T, encoder_dim) and `x_lens` of shape (N,).
+ It returns two tensors: `logits` of shape (N, T, encoder_dm) and
+ `logit_lens` of shape (N,).
+ """
+ super().__init__()
+ assert isinstance(encoder, EncoderInterface), type(encoder)
+
+ self.encoder = encoder
+
+ self.ctc_output = nn.Sequential(
+ nn.Dropout(p=0.1),
+ nn.Linear(encoder_dim, vocab_size),
+ nn.LogSoftmax(dim=-1),
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ x_lens: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Args:
+ x:
+ A 3-D tensor of shape (N, T, C).
+ x_lens:
+ A 1-D tensor of shape (N,). It contains the number of frames in `x`
+ before padding.
+ Returns:
+ Return the ctc outputs and encoder output lengths.
+ """
+ assert x.ndim == 3, x.shape
+ assert x_lens.ndim == 1, x_lens.shape
+
+ encoder_out, x_lens = self.encoder(x, x_lens)
+ assert torch.all(x_lens > 0)
+
+ # compute ctc log-probs
+ ctc_output = self.ctc_output(encoder_out)
+
+ return ctc_output, x_lens
diff --git a/egs/librispeech/ASR/zipformer_mmi/optim.py b/egs/librispeech/ASR/zipformer_mmi/optim.py
new file mode 120000
index 000000000..81ac4a89a
--- /dev/null
+++ b/egs/librispeech/ASR/zipformer_mmi/optim.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/optim.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/zipformer_mmi/pretrained.py b/egs/librispeech/ASR/zipformer_mmi/pretrained.py
new file mode 100755
index 000000000..0e7fd0daf
--- /dev/null
+++ b/egs/librispeech/ASR/zipformer_mmi/pretrained.py
@@ -0,0 +1,410 @@
+#!/usr/bin/env python3
+# Copyright 2021-2022 Xiaomi Corp. (authors: Fangjun Kuang,
+# Zengwei)
+#
+# 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.
+"""
+This script loads a checkpoint and uses it to decode waves.
+You can generate the checkpoint with the following command:
+
+./zipformer_mmi/export.py \
+ --exp-dir ./zipformer_mmi/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 20 \
+ --avg 10
+
+Usage of this script:
+
+(1) 1best
+./zipformer_mmi/pretrained.py \
+ --checkpoint ./zipformer_mmi/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --method 1best \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+(2) nbest
+./zipformer_mmi/pretrained.py \
+ --checkpoint ./zipformer_mmi/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --nbest-scale 1.2 \
+ --method nbest \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+(3) nbest-rescoring-LG
+./zipformer_mmi/pretrained.py \
+ --checkpoint ./zipformer_mmi/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --nbest-scale 1.2 \
+ --method nbest-rescoring-LG \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+(4) nbest-rescoring-3-gram
+./zipformer_mmi/pretrained.py \
+ --checkpoint ./zipformer_mmi/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --nbest-scale 1.2 \
+ --method nbest-rescoring-3-gram \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+(5) nbest-rescoring-4-gram
+./zipformer_mmi/pretrained.py \
+ --checkpoint ./zipformer_mmi/exp/pretrained.pt \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --nbest-scale 1.2 \
+ --method nbest-rescoring-4-gram \
+ /path/to/foo.wav \
+ /path/to/bar.wav
+
+
+You can also use `./zipformer_mmi/exp/epoch-xx.pt`.
+
+Note: ./zipformer_mmi/exp/pretrained.pt is generated by
+./zipformer_mmi/export.py
+"""
+
+
+import argparse
+import logging
+import math
+from pathlib import Path
+from typing import List
+
+import k2
+import kaldifeat
+import sentencepiece as spm
+import torch
+import torchaudio
+from decode import get_decoding_params
+from torch.nn.utils.rnn import pad_sequence
+from train import add_model_arguments, get_ctc_model, get_params
+
+from icefall.decode import (
+ get_lattice,
+ nbest_decoding,
+ nbest_rescore_with_LM,
+ one_best_decoding,
+)
+from icefall.mmi_graph_compiler import MmiTrainingGraphCompiler
+from icefall.utils import get_texts
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--checkpoint",
+ type=str,
+ required=True,
+ help="Path to the checkpoint. "
+ "The checkpoint is assumed to be saved by "
+ "icefall.checkpoint.save_checkpoint().",
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ help="""Path to bpe.model.""",
+ )
+
+ parser.add_argument(
+ "--method",
+ type=str,
+ default="1best",
+ help="""Decoding method. Use HP as decoding graph, where H is
+ ctc_topo and P is token-level bi-gram lm.
+ Supported values are:
+ - (1) 1best. Extract the best path from the decoding lattice as the
+ decoding result.
+ - (2) nbest. Extract n paths from the decoding lattice; the path
+ with the highest score is the decoding result.
+ - (4) nbest-rescoring-LG. Extract n paths from the decoding lattice,
+ rescore them with an word-level 3-gram LM, the path with the
+ highest score is the decoding result.
+ - (5) nbest-rescoring-3-gram. Extract n paths from the decoding
+ lattice, rescore them with an token-level 3-gram LM, the path with
+ the highest score is the decoding result.
+ - (6) nbest-rescoring-4-gram. Extract n paths from the decoding
+ lattice, rescore them with an token-level 4-gram LM, the path with
+ the highest score is the decoding result.
+ """,
+ )
+
+ parser.add_argument(
+ "--sample-rate",
+ type=int,
+ default=16000,
+ help="The sample rate of the input sound file",
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=Path,
+ default="data/lang_bpe_500",
+ help="The lang dir containing word table and LG graph",
+ )
+
+ parser.add_argument(
+ "--num-paths",
+ type=int,
+ default=100,
+ help="""Number of paths for n-best based decoding method.
+ Used only when "method" is one of the following values:
+ nbest, nbest-rescoring, and nbest-oracle
+ """,
+ )
+
+ parser.add_argument(
+ "--nbest-scale",
+ type=float,
+ default=1.2,
+ help="""The scale to be applied to `lattice.scores`.
+ It's needed if you use any kinds of n-best based rescoring.
+ Used only when "method" is one of the following values:
+ nbest, nbest-rescoring, and nbest-oracle
+ A smaller value results in more unique paths.
+ """,
+ )
+
+ parser.add_argument(
+ "--ngram-lm-scale",
+ type=float,
+ default=0.1,
+ help="""
+ Used when method is nbest-rescoring-LG, nbest-rescoring-3-gram,
+ and nbest-rescoring-4-gram.
+ It specifies the scale for n-gram LM scores.
+ (Note: You need to tune it on a dataset.)
+ """,
+ )
+
+ parser.add_argument(
+ "--hp-scale",
+ type=float,
+ default=1.0,
+ help="""The scale to be applied to `ctc_topo_P.scores`.
+ """,
+ )
+
+ parser.add_argument(
+ "sound_files",
+ type=str,
+ nargs="+",
+ help="The input sound file(s) to transcribe. "
+ "Supported formats are those supported by torchaudio.load(). "
+ "For example, wav and flac are supported. "
+ "The sample rate has to be 16kHz.",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def read_sound_files(
+ filenames: List[str], expected_sample_rate: float
+) -> List[torch.Tensor]:
+ """Read a list of sound files into a list 1-D float32 torch tensors.
+ Args:
+ filenames:
+ A list of sound filenames.
+ expected_sample_rate:
+ The expected sample rate of the sound files.
+ Returns:
+ Return a list of 1-D float32 torch tensors.
+ """
+ ans = []
+ for f in filenames:
+ wave, sample_rate = torchaudio.load(f)
+ assert (
+ sample_rate == expected_sample_rate
+ ), f"expected sample rate: {expected_sample_rate}. Given: {sample_rate}"
+ # We use only the first channel
+ ans.append(wave[0])
+ return ans
+
+
+@torch.no_grad()
+def main():
+ parser = get_parser()
+ args = parser.parse_args()
+
+ params = get_params()
+ # add decoding params
+ params.update(get_decoding_params())
+ params.update(vars(args))
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(params.bpe_model)
+
+ # is defined in local/train_bpe_model.py
+ params.blank_id = sp.piece_to_id("")
+ params.unk_id = sp.piece_to_id("")
+ params.vocab_size = sp.get_piece_size()
+
+ logging.info(f"{params}")
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"device: {device}")
+
+ logging.info("Creating model")
+ model = get_ctc_model(params)
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ checkpoint = torch.load(args.checkpoint, map_location="cpu")
+ model.load_state_dict(checkpoint["model"], strict=False)
+ model.to(device)
+ model.eval()
+ model.device = device
+
+ logging.info("Constructing Fbank computer")
+ opts = kaldifeat.FbankOptions()
+ opts.device = device
+ opts.frame_opts.dither = 0
+ opts.frame_opts.snip_edges = False
+ opts.frame_opts.samp_freq = params.sample_rate
+ opts.mel_opts.num_bins = params.feature_dim
+
+ fbank = kaldifeat.Fbank(opts)
+
+ logging.info(f"Reading sound files: {params.sound_files}")
+ waves = read_sound_files(
+ filenames=params.sound_files, expected_sample_rate=params.sample_rate
+ )
+ waves = [w.to(device) for w in waves]
+
+ logging.info("Decoding started")
+ features = fbank(waves)
+ feature_lengths = [f.size(0) for f in features]
+
+ features = pad_sequence(features, batch_first=True, padding_value=math.log(1e-10))
+ feature_lengths = torch.tensor(feature_lengths, device=device)
+
+ bpe_model = spm.SentencePieceProcessor()
+ bpe_model.load(str(params.lang_dir / "bpe.model"))
+ mmi_graph_compiler = MmiTrainingGraphCompiler(
+ params.lang_dir,
+ uniq_filename="lexicon.txt",
+ device=device,
+ oov="",
+ sos_id=1,
+ eos_id=1,
+ )
+ HP = mmi_graph_compiler.ctc_topo_P
+ HP.scores *= params.hp_scale
+ if not hasattr(HP, "lm_scores"):
+ HP.lm_scores = HP.scores.clone()
+
+ method = params.method
+ assert method in (
+ "1best",
+ "nbest",
+ "nbest-rescoring-LG", # word-level 3-gram lm
+ "nbest-rescoring-3-gram", # token-level 3-gram lm
+ "nbest-rescoring-4-gram", # token-level 4-gram lm
+ )
+ # loading language model for rescoring
+ LM = None
+ if method == "nbest-rescoring-LG":
+ lg_filename = params.lang_dir / "LG.pt"
+ logging.info(f"Loading {lg_filename}")
+ LG = k2.Fsa.from_dict(torch.load(lg_filename, map_location=device))
+ LG = k2.Fsa.from_fsas([LG]).to(device)
+ LG.lm_scores = LG.scores.clone()
+ LM = LG
+ elif method in ["nbest-rescoring-3-gram", "nbest-rescoring-4-gram"]:
+ order = method[-6]
+ assert order in ("3", "4")
+ order = int(order)
+ logging.info(f"Loading pre-compiled {order}gram.pt")
+ d = torch.load(params.lang_dir / f"{order}gram.pt", map_location=device)
+ G = k2.Fsa.from_dict(d)
+ G.lm_scores = G.scores.clone()
+ LM = G
+
+ # Encoder forward
+ nnet_output, encoder_out_lens = model(x=features, x_lens=feature_lengths)
+
+ batch_size = nnet_output.shape[0]
+ supervision_segments = torch.tensor(
+ [
+ [i, 0, feature_lengths[i] // params.subsampling_factor]
+ for i in range(batch_size)
+ ],
+ dtype=torch.int32,
+ )
+
+ lattice = get_lattice(
+ nnet_output=nnet_output,
+ decoding_graph=HP,
+ 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 method in ["1best", "nbest"]:
+ if method == "1best":
+ best_path = one_best_decoding(
+ lattice=lattice, use_double_scores=params.use_double_scores
+ )
+ else:
+ best_path = nbest_decoding(
+ lattice=lattice,
+ num_paths=params.num_paths,
+ use_double_scores=params.use_double_scores,
+ nbest_scale=params.nbest_scale,
+ )
+ else:
+ best_path_dict = nbest_rescore_with_LM(
+ lattice=lattice,
+ LM=LM,
+ num_paths=params.num_paths,
+ lm_scale_list=[params.ngram_lm_scale],
+ nbest_scale=params.nbest_scale,
+ )
+ best_path = next(iter(best_path_dict.values()))
+
+ # Note: `best_path.aux_labels` contains token IDs, not word IDs
+ # since we are using HP, not HLG here.
+ #
+ # token_ids is a lit-of-list of IDs
+ token_ids = get_texts(best_path)
+ # hyps is a list of str, e.g., ['xxx yyy zzz', ...]
+ hyps = bpe_model.decode(token_ids)
+ # hyps is a list of list of str, e.g., [['xxx', 'yyy', 'zzz'], ... ]
+ hyps = [s.split() for s in hyps]
+ s = "\n"
+ for filename, hyp in zip(params.sound_files, hyps):
+ words = " ".join(hyp)
+ s += f"{filename}:\n{words}\n\n"
+ logging.info(s)
+
+ logging.info("Decoding Done")
+
+
+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/zipformer_mmi/scaling.py b/egs/librispeech/ASR/zipformer_mmi/scaling.py
new file mode 120000
index 000000000..2428b74b9
--- /dev/null
+++ b/egs/librispeech/ASR/zipformer_mmi/scaling.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/scaling.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/zipformer_mmi/scaling_converter.py b/egs/librispeech/ASR/zipformer_mmi/scaling_converter.py
new file mode 120000
index 000000000..b8b8ba432
--- /dev/null
+++ b/egs/librispeech/ASR/zipformer_mmi/scaling_converter.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/scaling_converter.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/zipformer_mmi/test_model.py b/egs/librispeech/ASR/zipformer_mmi/test_model.py
new file mode 100755
index 000000000..7782845f4
--- /dev/null
+++ b/egs/librispeech/ASR/zipformer_mmi/test_model.py
@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+# Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang)
+#
+# 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.
+
+
+"""
+To run this file, do:
+
+ cd icefall/egs/librispeech/ASR
+ python ./zipformer_mmi/test_model.py
+"""
+
+import torch
+from train import get_ctc_model, get_params
+
+
+def test_model():
+ params = get_params()
+ params.vocab_size = 500
+ params.num_encoder_layers = "2,4,3,2,4"
+ # params.feedforward_dims = "1024,1024,1536,1536,1024"
+ params.feedforward_dims = "1024,1024,2048,2048,1024"
+ params.nhead = "8,8,8,8,8"
+ params.encoder_dims = "384,384,384,384,384"
+ params.attention_dims = "192,192,192,192,192"
+ params.encoder_unmasked_dims = "256,256,256,256,256"
+ params.zipformer_downsampling_factors = "1,2,4,8,2"
+ params.cnn_module_kernels = "31,31,31,31,31"
+ model = get_ctc_model(params)
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ print(f"Number of model parameters: {num_param}")
+
+ features = torch.randn(2, 100, 80)
+ feature_lengths = torch.full((2,), 100)
+ model(x=features, x_lens=feature_lengths)
+
+
+def main():
+ test_model()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/librispeech/ASR/zipformer_mmi/train.py b/egs/librispeech/ASR/zipformer_mmi/train.py
new file mode 100755
index 000000000..b2784e47c
--- /dev/null
+++ b/egs/librispeech/ASR/zipformer_mmi/train.py
@@ -0,0 +1,1198 @@
+#!/usr/bin/env python3
+# Copyright 2021-2022 Xiaomi Corp. (authors: Fangjun Kuang,
+# Wei Kang,
+# Mingshuang Luo,)
+# Zengwei Yao)
+#
+# 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.
+"""
+Usage:
+
+export CUDA_VISIBLE_DEVICES="0,1,2,3"
+
+./zipformer_mmi/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --exp-dir zipformer_mmi/exp \
+ --full-libri 1 \
+ --max-duration 300
+
+# For mix precision training:
+
+./zipformer_mmi/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --use-fp16 1 \
+ --exp-dir zipformer_mmi/exp \
+ --full-libri 1 \
+ --max-duration 500
+
+"""
+
+
+import argparse
+import copy
+import logging
+import warnings
+from pathlib import Path
+from shutil import copyfile
+from typing import Any, Dict, Optional, Tuple, Union
+
+import k2
+import optim
+import torch
+import torch.multiprocessing as mp
+import torch.nn as nn
+from asr_datamodule import LibriSpeechAsrDataModule
+from lhotse.cut import Cut
+from lhotse.dataset.sampling.base import CutSampler
+from lhotse.utils import fix_random_seed
+from model import CTCModel
+from optim import Eden, ScaledAdam
+from torch import Tensor
+from torch.cuda.amp import GradScaler
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.tensorboard import SummaryWriter
+from zipformer import Zipformer
+
+from icefall import diagnostics
+from icefall.bpe_graph_compiler import BpeCtcTrainingGraphCompiler
+from icefall.checkpoint import load_checkpoint, remove_checkpoints
+from icefall.checkpoint import save_checkpoint as save_checkpoint_impl
+from icefall.checkpoint import (
+ save_checkpoint_with_global_batch_idx,
+ update_averaged_model,
+)
+from icefall.dist import cleanup_dist, setup_dist
+from icefall.env import get_env_info
+from icefall.hooks import register_inf_check_hooks
+from icefall.lexicon import Lexicon, UniqLexicon
+from icefall.mmi import LFMMILoss
+from icefall.mmi_graph_compiler import MmiTrainingGraphCompiler
+from icefall.utils import (
+ AttributeDict,
+ MetricsTracker,
+ encode_supervisions,
+ setup_logger,
+ str2bool,
+)
+
+LRSchedulerType = Union[torch.optim.lr_scheduler._LRScheduler, optim.LRScheduler]
+
+
+def set_batch_count(model: Union[nn.Module, DDP], batch_count: float) -> None:
+ if isinstance(model, DDP):
+ # get underlying nn.Module
+ model = model.module
+ for module in model.modules():
+ if hasattr(module, "batch_count"):
+ module.batch_count = batch_count
+
+
+def add_model_arguments(parser: argparse.ArgumentParser):
+ parser.add_argument(
+ "--num-encoder-layers",
+ type=str,
+ default="2,4,3,2,4",
+ help="Number of zipformer encoder layers, comma separated.",
+ )
+
+ parser.add_argument(
+ "--feedforward-dims",
+ type=str,
+ default="1024,1024,2048,2048,1024",
+ help="Feedforward dimension of the zipformer encoder layers, comma separated.",
+ )
+
+ parser.add_argument(
+ "--nhead",
+ type=str,
+ default="8,8,8,8,8",
+ help="Number of attention heads in the zipformer encoder layers.",
+ )
+
+ parser.add_argument(
+ "--encoder-dims",
+ type=str,
+ default="384,384,384,384,384",
+ help="Embedding dimension in the 2 blocks of zipformer encoder layers, comma separated",
+ )
+
+ parser.add_argument(
+ "--attention-dims",
+ type=str,
+ default="192,192,192,192,192",
+ help="""Attention dimension in the 2 blocks of zipformer encoder layers, comma separated;
+ not the same as embedding dimension.""",
+ )
+
+ parser.add_argument(
+ "--encoder-unmasked-dims",
+ type=str,
+ default="256,256,256,256,256",
+ help="Unmasked dimensions in the encoders, relates to augmentation during training. "
+ "Must be <= each of encoder_dims. Empirically, less than 256 seems to make performance "
+ " worse.",
+ )
+
+ parser.add_argument(
+ "--zipformer-downsampling-factors",
+ type=str,
+ default="1,2,4,8,2",
+ help="Downsampling factor for each stack of encoder layers.",
+ )
+
+ parser.add_argument(
+ "--cnn-module-kernels",
+ type=str,
+ default="31,31,31,31,31",
+ help="Sizes of kernels in convolution modules",
+ )
+
+
+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.",
+ )
+
+ parser.add_argument(
+ "--num-epochs",
+ type=int,
+ default=30,
+ help="Number of epochs to train.",
+ )
+
+ parser.add_argument(
+ "--start-epoch",
+ type=int,
+ default=1,
+ help="""Resume training from this epoch. It should be positive.
+ If larger than 1, it will load checkpoint from
+ exp-dir/epoch-{start_epoch-1}.pt
+ """,
+ )
+
+ parser.add_argument(
+ "--start-batch",
+ type=int,
+ default=0,
+ help="""If positive, --start-epoch is ignored and
+ it loads the checkpoint from exp-dir/checkpoint-{start_batch}.pt
+ """,
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="zipformer_mmi/exp",
+ help="""The experiment dir.
+ It specifies the directory where all training related
+ files, e.g., checkpoints, log, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=str,
+ default="data/lang_bpe_500",
+ help="""The lang dir
+ It contains language related input files such as
+ "lexicon.txt"
+ """,
+ )
+
+ parser.add_argument(
+ "--base-lr", type=float, default=0.05, help="The base learning rate."
+ )
+
+ parser.add_argument(
+ "--lr-batches",
+ type=float,
+ default=5000,
+ help="""Number of steps that affects how rapidly the learning rate
+ decreases. We suggest not to change this.""",
+ )
+
+ parser.add_argument(
+ "--lr-epochs",
+ type=float,
+ default=3.5,
+ help="""Number of epochs that affects how rapidly the learning rate decreases.
+ """,
+ )
+
+ parser.add_argument(
+ "--seed",
+ type=int,
+ default=42,
+ help="The seed for random generators intended for reproducibility",
+ )
+
+ parser.add_argument(
+ "--use-pruned-intersect",
+ type=str2bool,
+ default=False,
+ help="""Whether to use `intersect_dense_pruned` to get denominator
+ lattice.""",
+ )
+
+ parser.add_argument(
+ "--print-diagnostics",
+ type=str2bool,
+ default=False,
+ help="Accumulate stats on activations, print them and exit.",
+ )
+
+ parser.add_argument(
+ "--inf-check",
+ type=str2bool,
+ default=False,
+ help="Add hooks to check for infinite module outputs and gradients.",
+ )
+
+ parser.add_argument(
+ "--save-every-n",
+ type=int,
+ default=2000,
+ help="""Save checkpoint after processing this number of batches"
+ periodically. We save checkpoint to exp-dir/ whenever
+ params.batch_idx_train % save_every_n == 0. The checkpoint filename
+ has the form: f'exp-dir/checkpoint-{params.batch_idx_train}.pt'
+ Note: It also saves checkpoint to `exp-dir/epoch-xxx.pt` at the
+ end of each epoch where `xxx` is the epoch number counting from 0.
+ """,
+ )
+
+ parser.add_argument(
+ "--keep-last-k",
+ type=int,
+ default=30,
+ help="""Only keep this number of checkpoints on disk.
+ For instance, if it is 3, there are only 3 checkpoints
+ in the exp-dir with filenames `checkpoint-xxx.pt`.
+ It does not affect checkpoints with name `epoch-xxx.pt`.
+ """,
+ )
+
+ parser.add_argument(
+ "--average-period",
+ type=int,
+ default=200,
+ help="""Update the averaged model, namely `model_avg`, after processing
+ this number of batches. `model_avg` is a separate version of model,
+ in which each floating-point parameter is the average of all the
+ parameters from the start of training. Each time we take the average,
+ we do: `model_avg = model * (average_period / batch_idx_train) +
+ model_avg * ((batch_idx_train - average_period) / batch_idx_train)`.
+ """,
+ )
+
+ parser.add_argument(
+ "--use-fp16",
+ type=str2bool,
+ default=False,
+ help="Whether to use half precision training.",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def get_params() -> AttributeDict:
+ """Return a dict containing training parameters.
+
+ All training related parameters that are not passed from the commandline
+ are 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`:
+
+ - 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
+
+ - reset_interval: Reset statistics if batch_idx % reset_interval is 0
+
+ - valid_interval: Run validation if batch_idx % valid_interval is 0
+
+ - feature_dim: The model input dim. It has to match the one used
+ in computing features.
+
+ - subsampling_factor: The subsampling factor for the model.
+
+ - encoder_dim: Hidden dim for multi-head attention model.
+
+ - num_decoder_layers: Number of decoder layer of transformer decoder.
+
+ - warm_step: The warmup period that dictates the decay of the
+ scale on "simple" (un-pruned) loss.
+ """
+ params = AttributeDict(
+ {
+ "best_train_loss": float("inf"),
+ "best_valid_loss": float("inf"),
+ "best_train_epoch": -1,
+ "best_valid_epoch": -1,
+ "batch_idx_train": 0,
+ "log_interval": 50,
+ "reset_interval": 200,
+ "valid_interval": 3000, # For the 100h subset, use 800
+ # parameters for zipformer
+ "feature_dim": 80,
+ "subsampling_factor": 4, # not passed in, this is fixed.
+ # parameters for mmi loss
+ "mmi_beam_size": 6,
+ "den_scale": 1.0,
+ # parameters for mmi loss
+ "ctc_beam_size": 10,
+ "reduction": "sum",
+ "use_double_scores": True,
+ "warm_step": 2000,
+ "env_info": get_env_info(),
+ }
+ )
+
+ return params
+
+
+def get_encoder_model(params: AttributeDict) -> nn.Module:
+ # TODO: We can add an option to switch between Zipformer and Transformer
+ def to_int_tuple(s: str):
+ return tuple(map(int, s.split(",")))
+
+ encoder = Zipformer(
+ num_features=params.feature_dim,
+ output_downsampling_factor=2,
+ zipformer_downsampling_factors=to_int_tuple(
+ params.zipformer_downsampling_factors
+ ),
+ encoder_dims=to_int_tuple(params.encoder_dims),
+ attention_dim=to_int_tuple(params.attention_dims),
+ encoder_unmasked_dims=to_int_tuple(params.encoder_unmasked_dims),
+ nhead=to_int_tuple(params.nhead),
+ feedforward_dim=to_int_tuple(params.feedforward_dims),
+ cnn_module_kernels=to_int_tuple(params.cnn_module_kernels),
+ num_encoder_layers=to_int_tuple(params.num_encoder_layers),
+ )
+ return encoder
+
+
+def get_ctc_model(params: AttributeDict) -> nn.Module:
+ encoder = get_encoder_model(params)
+
+ model = CTCModel(
+ encoder=encoder,
+ encoder_dim=int(params.encoder_dims.split(",")[-1]),
+ vocab_size=params.vocab_size,
+ )
+ return model
+
+
+def load_checkpoint_if_available(
+ params: AttributeDict,
+ model: nn.Module,
+ model_avg: nn.Module = None,
+ optimizer: Optional[torch.optim.Optimizer] = None,
+ scheduler: Optional[LRSchedulerType] = None,
+) -> Optional[Dict[str, Any]]:
+ """Load checkpoint from file.
+
+ If params.start_batch is positive, it will load the checkpoint from
+ `params.exp_dir/checkpoint-{params.start_batch}.pt`. Otherwise, if
+ params.start_epoch is larger than 1, it will load the checkpoint from
+ `params.start_epoch - 1`.
+
+ Apart from loading state dict for `model` and `optimizer` 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.
+ model_avg:
+ The stored model averaged from the start of training.
+ optimizer:
+ The optimizer that we are using.
+ scheduler:
+ The scheduler that we are using.
+ Returns:
+ Return a dict containing previously saved training info.
+ """
+ if params.start_batch > 0:
+ filename = params.exp_dir / f"checkpoint-{params.start_batch}.pt"
+ elif params.start_epoch > 1:
+ filename = params.exp_dir / f"epoch-{params.start_epoch-1}.pt"
+ else:
+ return None
+
+ assert filename.is_file(), f"{filename} does not exist!"
+
+ saved_params = load_checkpoint(
+ filename,
+ model=model,
+ model_avg=model_avg,
+ 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]
+
+ if params.start_batch > 0:
+ if "cur_epoch" in saved_params:
+ params["start_epoch"] = saved_params["cur_epoch"]
+
+ if "cur_batch_idx" in saved_params:
+ params["cur_batch_idx"] = saved_params["cur_batch_idx"]
+
+ return saved_params
+
+
+def save_checkpoint(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ model_avg: Optional[nn.Module] = None,
+ optimizer: Optional[torch.optim.Optimizer] = None,
+ scheduler: Optional[LRSchedulerType] = None,
+ sampler: Optional[CutSampler] = None,
+ scaler: Optional[GradScaler] = 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.
+ model_avg:
+ The stored model averaged from the start of training.
+ optimizer:
+ The optimizer used in the training.
+ sampler:
+ The sampler for the training dataset.
+ scaler:
+ The scaler used for mix precision training.
+ """
+ if rank != 0:
+ return
+ filename = params.exp_dir / f"epoch-{params.cur_epoch}.pt"
+ save_checkpoint_impl(
+ filename=filename,
+ model=model,
+ model_avg=model_avg,
+ params=params,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=sampler,
+ scaler=scaler,
+ 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: Union[nn.Module, DDP],
+ ctc_graph_compiler: BpeCtcTrainingGraphCompiler,
+ mmi_graph_compiler: MmiTrainingGraphCompiler,
+ batch: dict,
+ is_training: bool,
+) -> Tuple[Tensor, MetricsTracker]:
+ """
+ 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 Zipformer in our case.
+ 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.
+ batch:
+ A batch of data. See `lhotse.dataset.K2SpeechRecognitionDataset()`
+ for the content in it.
+ 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 = model.device if isinstance(model, DDP) else next(model.parameters()).device
+ feature = batch["inputs"]
+ # at entry, feature is (N, T, C)
+ assert feature.ndim == 3
+ feature = feature.to(device)
+
+ supervisions = batch["supervisions"]
+ feature_lens = supervisions["num_frames"].to(device)
+
+ batch_idx_train = params.batch_idx_train
+ warm_step = params.warm_step
+
+ with torch.set_grad_enabled(is_training):
+ nnet_output, encoder_out_lens = model(x=feature, x_lens=feature_lens)
+
+ # NOTE: We need `encode_supervisions` to sort sequences with
+ # different duration in decreasing order, required by
+ # `k2.intersect_dense` called in `LFMMILoss.forward()`
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ supervision_segments, texts = encode_supervisions(
+ supervisions, subsampling_factor=params.subsampling_factor
+ )
+
+ dense_fsa_vec = k2.DenseFsaVec(
+ nnet_output,
+ supervision_segments,
+ allow_truncate=params.subsampling_factor - 1,
+ )
+
+ info = MetricsTracker()
+ if batch_idx_train < warm_step:
+ # Training with ctc loss
+ # Works with a BPE model
+ token_ids = ctc_graph_compiler.texts_to_ids(texts)
+ decoding_graph = ctc_graph_compiler.compile(token_ids)
+ loss = k2.ctc_loss(
+ decoding_graph=decoding_graph,
+ dense_fsa_vec=dense_fsa_vec,
+ output_beam=params.ctc_beam_size,
+ reduction=params.reduction,
+ use_double_scores=params.use_double_scores,
+ )
+ info["ctc_loss"] = loss.detach().cpu().item()
+ info["mmi_loss"] = 0
+ else:
+ # Training with mmi loss
+ loss_fn = LFMMILoss(
+ graph_compiler=mmi_graph_compiler,
+ use_pruned_intersect=params.use_pruned_intersect,
+ den_scale=params.den_scale,
+ beam_size=params.mmi_beam_size,
+ )
+ loss = loss_fn(dense_fsa_vec=dense_fsa_vec, texts=texts)
+ info["ctc_loss"] = 0
+ info["mmi_loss"] = loss.detach().cpu().item()
+
+ assert loss.requires_grad == is_training
+
+ info["frames"] = encoder_out_lens.sum().cpu().item()
+ # Note: We use reduction=sum while computing the loss.
+ info["loss"] = loss.detach().cpu().item()
+
+ return loss, info
+
+
+def compute_validation_loss(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ ctc_graph_compiler: BpeCtcTrainingGraphCompiler,
+ mmi_graph_compiler: MmiTrainingGraphCompiler,
+ valid_dl: torch.utils.data.DataLoader,
+ world_size: int = 1,
+) -> MetricsTracker:
+ """Run the validation process."""
+ model.eval()
+
+ tot_loss = MetricsTracker()
+
+ for batch_idx, batch in enumerate(valid_dl):
+ loss, loss_info = compute_loss(
+ params=params,
+ model=model,
+ ctc_graph_compiler=ctc_graph_compiler,
+ mmi_graph_compiler=mmi_graph_compiler,
+ batch=batch,
+ is_training=False,
+ )
+ assert loss.requires_grad is False
+ tot_loss = tot_loss + loss_info
+
+ if world_size > 1:
+ tot_loss.reduce(loss.device)
+
+ loss_value = tot_loss["loss"] / tot_loss["frames"]
+ if loss_value < params.best_valid_loss:
+ params.best_valid_epoch = params.cur_epoch
+ params.best_valid_loss = loss_value
+
+ return tot_loss
+
+
+def train_one_epoch(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ optimizer: torch.optim.Optimizer,
+ scheduler: LRSchedulerType,
+ ctc_graph_compiler: BpeCtcTrainingGraphCompiler,
+ mmi_graph_compiler: MmiTrainingGraphCompiler,
+ train_dl: torch.utils.data.DataLoader,
+ valid_dl: torch.utils.data.DataLoader,
+ scaler: GradScaler,
+ model_avg: Optional[nn.Module] = None,
+ tb_writer: Optional[SummaryWriter] = None,
+ world_size: int = 1,
+ rank: int = 0,
+) -> 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.
+ scheduler:
+ The learning rate scheduler, we call step() every step.
+ graph_compiler:
+ It is used to convert transcripts to FSAs.
+ train_dl:
+ Dataloader for the training dataset.
+ valid_dl:
+ Dataloader for the validation dataset.
+ scaler:
+ The scaler used for mix precision training.
+ model_avg:
+ The stored model averaged from the start of training.
+ tb_writer:
+ Writer to write log messages to tensorboard.
+ world_size:
+ Number of nodes in DDP training. If it is 1, DDP is disabled.
+ rank:
+ The rank of the node in DDP training. If no DDP is used, it should
+ be set to 0.
+ """
+ model.train()
+
+ tot_loss = MetricsTracker()
+
+ cur_batch_idx = params.get("cur_batch_idx", 0)
+
+ for batch_idx, batch in enumerate(train_dl):
+ if batch_idx < cur_batch_idx:
+ continue
+ cur_batch_idx = batch_idx
+
+ params.batch_idx_train += 1
+ batch_size = len(batch["supervisions"]["text"])
+
+ try:
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, loss_info = compute_loss(
+ params=params,
+ model=model,
+ ctc_graph_compiler=ctc_graph_compiler,
+ mmi_graph_compiler=mmi_graph_compiler,
+ batch=batch,
+ is_training=True,
+ )
+ # summary stats
+ tot_loss = (tot_loss * (1 - 1 / params.reset_interval)) + loss_info
+
+ # NOTE: We use reduction==sum and loss is computed over utterances
+ # in the batch and there is no normalization to it so far.
+ scaler.scale(loss).backward()
+ set_batch_count(model, params.batch_idx_train)
+ scheduler.step_batch(params.batch_idx_train)
+
+ scaler.step(optimizer)
+ scaler.update()
+ optimizer.zero_grad()
+ except: # noqa
+ display_and_save_batch(
+ batch, params=params, graph_compiler=mmi_graph_compiler
+ )
+ raise
+
+ if params.print_diagnostics and batch_idx == 5:
+ return
+
+ if (
+ rank == 0
+ and params.batch_idx_train > 0
+ and params.batch_idx_train % params.average_period == 0
+ ):
+ update_averaged_model(
+ params=params,
+ model_cur=model,
+ model_avg=model_avg,
+ )
+
+ if (
+ params.batch_idx_train > 0
+ and params.batch_idx_train % params.save_every_n == 0
+ ):
+ params.cur_batch_idx = batch_idx
+ save_checkpoint_with_global_batch_idx(
+ out_dir=params.exp_dir,
+ global_batch_idx=params.batch_idx_train,
+ model=model,
+ model_avg=model_avg,
+ params=params,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=train_dl.sampler,
+ scaler=scaler,
+ rank=rank,
+ )
+ del params.cur_batch_idx
+ remove_checkpoints(
+ out_dir=params.exp_dir,
+ topk=params.keep_last_k,
+ rank=rank,
+ )
+
+ if batch_idx % 100 == 0 and params.use_fp16:
+ # If the grad scale was less than 1, try increasing it. The _growth_interval
+ # of the grad scaler is configurable, but we can't configure it to have different
+ # behavior depending on the current grad scale.
+ cur_grad_scale = scaler._scale.item()
+ if cur_grad_scale < 1.0 or (cur_grad_scale < 8.0 and batch_idx % 400 == 0):
+ scaler.update(cur_grad_scale * 2.0)
+ if cur_grad_scale < 0.01:
+ logging.warning(f"Grad scale is small: {cur_grad_scale}")
+ if cur_grad_scale < 1.0e-05:
+ raise RuntimeError(
+ f"grad_scale is too small, exiting: {cur_grad_scale}"
+ )
+
+ if batch_idx % params.log_interval == 0:
+ cur_lr = scheduler.get_last_lr()[0]
+ cur_grad_scale = scaler._scale.item() if params.use_fp16 else 1.0
+
+ logging.info(
+ f"Epoch {params.cur_epoch}, "
+ f"batch {batch_idx}, loss[{loss_info}], "
+ f"tot_loss[{tot_loss}], batch size: {batch_size}, "
+ f"lr: {cur_lr:.2e}, "
+ + (f"grad_scale: {scaler._scale.item()}" if params.use_fp16 else "")
+ )
+
+ if tb_writer is not None:
+ tb_writer.add_scalar(
+ "train/learning_rate", cur_lr, params.batch_idx_train
+ )
+
+ loss_info.write_summary(
+ tb_writer, "train/current_", params.batch_idx_train
+ )
+ tot_loss.write_summary(tb_writer, "train/tot_", params.batch_idx_train)
+ if params.use_fp16:
+ tb_writer.add_scalar(
+ "train/grad_scale",
+ cur_grad_scale,
+ params.batch_idx_train,
+ )
+
+ if batch_idx % params.valid_interval == 0 and not params.print_diagnostics:
+ logging.info("Computing validation loss")
+ valid_info = compute_validation_loss(
+ params=params,
+ model=model,
+ ctc_graph_compiler=ctc_graph_compiler,
+ mmi_graph_compiler=mmi_graph_compiler,
+ valid_dl=valid_dl,
+ world_size=world_size,
+ )
+ model.train()
+ logging.info(f"Epoch {params.cur_epoch}, validation: {valid_info}")
+ logging.info(
+ f"Maximum memory allocated so far is {torch.cuda.max_memory_allocated()//1000000}MB"
+ )
+ if tb_writer is not None:
+ valid_info.write_summary(
+ tb_writer, "train/valid_", params.batch_idx_train
+ )
+
+ loss_value = tot_loss["loss"] / tot_loss["frames"]
+ params.train_loss = loss_value
+ 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))
+ if params.full_libri is False:
+ params.valid_interval = 1600
+
+ fix_random_seed(params.seed)
+ 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")
+
+ 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
+ params.vocab_size = num_classes
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", rank)
+ logging.info(f"Device: {device}")
+
+ assert "lang_bpe" in str(params.lang_dir)
+ ctc_graph_compiler = BpeCtcTrainingGraphCompiler(
+ params.lang_dir,
+ device=device,
+ sos_token="",
+ eos_token="",
+ )
+ mmi_graph_compiler = MmiTrainingGraphCompiler(
+ params.lang_dir,
+ uniq_filename="lexicon.txt",
+ device=device,
+ oov="",
+ sos_id=1,
+ eos_id=1,
+ )
+
+ logging.info(params)
+
+ logging.info("About to create model")
+ model = get_ctc_model(params)
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ assert params.save_every_n >= params.average_period
+ model_avg: Optional[nn.Module] = None
+ if rank == 0:
+ # model_avg is only used with rank 0
+ model_avg = copy.deepcopy(model).to(torch.float64)
+
+ assert params.start_epoch > 0, params.start_epoch
+ checkpoints = load_checkpoint_if_available(
+ params=params, model=model, model_avg=model_avg
+ )
+
+ model.to(device)
+ if world_size > 1:
+ logging.info("Using DDP")
+ model = DDP(model, device_ids=[rank], find_unused_parameters=True)
+
+ parameters_names = []
+ parameters_names.append(
+ [name_param_pair[0] for name_param_pair in model.named_parameters()]
+ )
+ optimizer = ScaledAdam(
+ model.parameters(),
+ lr=params.base_lr,
+ clipping_scale=2.0,
+ parameters_names=parameters_names,
+ )
+
+ scheduler = Eden(optimizer, params.lr_batches, params.lr_epochs)
+
+ if checkpoints and "optimizer" in checkpoints:
+ logging.info("Loading optimizer state dict")
+ optimizer.load_state_dict(checkpoints["optimizer"])
+
+ if (
+ checkpoints
+ and "scheduler" in checkpoints
+ and checkpoints["scheduler"] is not None
+ ):
+ logging.info("Loading scheduler state dict")
+ scheduler.load_state_dict(checkpoints["scheduler"])
+
+ if params.print_diagnostics:
+ opts = diagnostics.TensorDiagnosticOptions(
+ 2**22
+ ) # allow 4 megabytes per sub-module
+ diagnostic = diagnostics.attach_diagnostics(model, opts)
+
+ if params.inf_check:
+ register_inf_check_hooks(model)
+
+ librispeech = LibriSpeechAsrDataModule(args)
+
+ # train_cuts = librispeech.train_clean_100_cuts()
+ if params.full_libri:
+ # train_cuts += librispeech.train_clean_360_cuts()
+ # train_cuts += librispeech.train_other_500_cuts()
+ train_cuts = librispeech.train_all_shuf_cuts()
+ else:
+ train_cuts = librispeech.train_clean_100_cuts()
+
+ def remove_short_and_long_utt(c: Cut):
+ # Keep only utterances with duration between 1 second and 20 seconds
+ #
+ # Caution: There is a reason to select 20.0 here. Please see
+ # ../local/display_manifest_statistics.py
+ #
+ # You should use ../local/display_manifest_statistics.py to get
+ # an utterance duration distribution for your dataset to select
+ # the threshold
+ return 1.0 <= c.duration <= 20.0
+
+ train_cuts = train_cuts.filter(remove_short_and_long_utt)
+
+ if params.start_batch > 0 and checkpoints and "sampler" in checkpoints:
+ # We only load the sampler's state dict when it loads a checkpoint
+ # saved in the middle of an epoch
+ sampler_state_dict = checkpoints["sampler"]
+ else:
+ sampler_state_dict = None
+
+ train_dl = librispeech.train_dataloaders(
+ train_cuts, sampler_state_dict=sampler_state_dict
+ )
+
+ valid_cuts = librispeech.dev_clean_cuts()
+ valid_cuts += librispeech.dev_other_cuts()
+ valid_dl = librispeech.valid_dataloaders(valid_cuts)
+
+ if not params.print_diagnostics:
+ scan_pessimistic_batches_for_oom(
+ model=model,
+ train_dl=train_dl,
+ optimizer=optimizer,
+ ctc_graph_compiler=ctc_graph_compiler,
+ mmi_graph_compiler=mmi_graph_compiler,
+ params=params,
+ )
+
+ scaler = GradScaler(enabled=params.use_fp16, init_scale=1.0)
+ if checkpoints and "grad_scaler" in checkpoints:
+ logging.info("Loading grad scaler state dict")
+ scaler.load_state_dict(checkpoints["grad_scaler"])
+
+ for epoch in range(params.start_epoch, params.num_epochs + 1):
+ scheduler.step_epoch(epoch - 1)
+ fix_random_seed(params.seed + epoch - 1)
+ train_dl.sampler.set_epoch(epoch - 1)
+
+ if tb_writer is not None:
+ tb_writer.add_scalar("train/epoch", epoch, params.batch_idx_train)
+
+ params.cur_epoch = epoch
+
+ train_one_epoch(
+ params=params,
+ model=model,
+ model_avg=model_avg,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ ctc_graph_compiler=ctc_graph_compiler,
+ mmi_graph_compiler=mmi_graph_compiler,
+ train_dl=train_dl,
+ valid_dl=valid_dl,
+ scaler=scaler,
+ tb_writer=tb_writer,
+ world_size=world_size,
+ rank=rank,
+ )
+
+ if params.print_diagnostics:
+ diagnostic.print_diagnostics()
+ break
+
+ save_checkpoint(
+ params=params,
+ model=model,
+ model_avg=model_avg,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=train_dl.sampler,
+ scaler=scaler,
+ rank=rank,
+ )
+
+ logging.info("Done!")
+
+ if world_size > 1:
+ torch.distributed.barrier()
+ cleanup_dist()
+
+
+def display_and_save_batch(
+ batch: dict,
+ params: AttributeDict,
+ graph_compiler: MmiTrainingGraphCompiler,
+) -> None:
+ """Display the batch statistics and save the batch into disk.
+
+ Args:
+ batch:
+ A batch of data. See `lhotse.dataset.K2SpeechRecognitionDataset()`
+ for the content in it.
+ params:
+ Parameters for training. See :func:`get_params`.
+ sp:
+ The BPE model.
+ """
+ from lhotse.utils import uuid4
+
+ filename = f"{params.exp_dir}/batch-{uuid4()}.pt"
+ logging.info(f"Saving batch to {filename}")
+ torch.save(batch, filename)
+
+ supervisions = batch["supervisions"]
+ features = batch["inputs"]
+
+ logging.info(f"features shape: {features.shape}")
+ y = graph_compiler.texts_to_ids(supervisions["text"])
+ num_tokens = sum(len(i) for i in y)
+ logging.info(f"num tokens: {num_tokens}")
+
+
+def scan_pessimistic_batches_for_oom(
+ model: Union[nn.Module, DDP],
+ train_dl: torch.utils.data.DataLoader,
+ optimizer: torch.optim.Optimizer,
+ ctc_graph_compiler: BpeCtcTrainingGraphCompiler,
+ mmi_graph_compiler: MmiTrainingGraphCompiler,
+ params: AttributeDict,
+):
+ from lhotse.dataset import find_pessimistic_batches
+
+ logging.info(
+ "Sanity check -- see if any of the batches in epoch 1 would cause OOM."
+ )
+ batches, crit_values = find_pessimistic_batches(train_dl.sampler)
+ for criterion, cuts in batches.items():
+ batch = train_dl.dataset[cuts]
+ try:
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, _ = compute_loss(
+ params=params,
+ model=model,
+ ctc_graph_compiler=ctc_graph_compiler,
+ mmi_graph_compiler=mmi_graph_compiler,
+ batch=batch,
+ is_training=True,
+ )
+ loss.backward()
+ optimizer.zero_grad()
+ except Exception as e:
+ if "CUDA out of memory" in str(e):
+ logging.error(
+ "Your GPU ran out of memory with the current "
+ "max_duration setting. We recommend decreasing "
+ "max_duration and trying again.\n"
+ f"Failing criterion: {criterion} "
+ f"(={crit_values[criterion]}) ..."
+ )
+ display_and_save_batch(
+ batch, params=params, graph_compiler=mmi_graph_compiler
+ )
+ raise
+ logging.info(
+ f"Maximum memory allocated so far is {torch.cuda.max_memory_allocated()//1000000}MB"
+ )
+
+
+def main():
+ parser = get_parser()
+ LibriSpeechAsrDataModule.add_arguments(parser)
+ args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ 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)
+
+
+torch.set_num_threads(1)
+torch.set_num_interop_threads(1)
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/librispeech/ASR/zipformer_mmi/zipformer.py b/egs/librispeech/ASR/zipformer_mmi/zipformer.py
new file mode 120000
index 000000000..79b076556
--- /dev/null
+++ b/egs/librispeech/ASR/zipformer_mmi/zipformer.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless7/zipformer.py
\ No newline at end of file
diff --git a/egs/tedlium3/ASR/RESULTS.md b/egs/tedlium3/ASR/RESULTS.md
index 511b19f73..38eaa8f44 100644
--- a/egs/tedlium3/ASR/RESULTS.md
+++ b/egs/tedlium3/ASR/RESULTS.md
@@ -1,5 +1,88 @@
## Results
+### TedLium3 BPE training results (Conformer-CTC 2)
+
+#### [conformer_ctc2](./conformer_ctc2)
+
+See for more details.
+
+The tensorboard log can be found at
+
+
+You can find a pretrained model and decoding results at:
+
+
+Number of model parameters: 101141699, i.e., 101.14 M
+
+The WERs are
+
+| | dev | test | comment |
+|--------------------------|------------|-------------|---------------------|
+| ctc decoding | 6.45 | 5.96 | --epoch 38 --avg 26 |
+| 1best | 5.92 | 5.51 | --epoch 38 --avg 26 |
+| whole lattice rescoring | 5.96 | 5.47 | --epoch 38 --avg 26 |
+| attention decoder | 5.60 | 5.33 | --epoch 38 --avg 26 |
+
+The training command for reproducing is given below:
+
+```
+export CUDA_VISIBLE_DEVICES="0,1,2,3"
+
+./conformer_ctc2/train.py \
+ --world-size 4 \
+ --num-epochs 40 \
+ --exp-dir conformer_ctc2/exp \
+ --max-duration 350 \
+ --use-fp16 true
+```
+
+The decoding command is:
+```
+epoch=38
+avg=26
+
+## ctc decoding
+./conformer_ctc2/decode.py \
+ --method ctc-decoding \
+ --exp-dir conformer_ctc2/exp \
+ --lang-dir data/lang_bpe_500 \
+ --result-dir conformer_ctc2/exp \
+ --max-duration 500 \
+ --epoch $epoch \
+ --avg $avg
+
+## 1best
+./conformer_ctc2/decode.py \
+ --method 1best \
+ --exp-dir conformer_ctc2/exp \
+ --lang-dir data/lang_bpe_500 \
+ --result-dir conformer_ctc2/exp \
+ --max-duration 500 \
+ --epoch $epoch \
+ --avg $avg
+
+## whole lattice rescoring
+./conformer_ctc2/decode.py \
+ --method whole-lattice-rescoring \
+ --exp-dir conformer_ctc2/exp \
+ --lm-path data/lm/G_4_gram_big.pt \
+ --lang-dir data/lang_bpe_500 \
+ --result-dir conformer_ctc2/exp \
+ --max-duration 500 \
+ --epoch $epoch \
+ --avg $avg
+
+## attention decoder
+./conformer_ctc2/decode.py \
+ --method attention-decoder \
+ --exp-dir conformer_ctc2/exp \
+ --lang-dir data/lang_bpe_500 \
+ --result-dir conformer_ctc2/exp \
+ --max-duration 500 \
+ --epoch $epoch \
+ --avg $avg
+```
+
### TedLium3 BPE training results (Pruned Transducer)
#### 2022-03-21
diff --git a/egs/tedlium3/ASR/conformer_ctc2/__init__.py b/egs/tedlium3/ASR/conformer_ctc2/__init__.py
new file mode 100755
index 000000000..e69de29bb
diff --git a/egs/tedlium3/ASR/conformer_ctc2/asr_datamodule.py b/egs/tedlium3/ASR/conformer_ctc2/asr_datamodule.py
new file mode 120000
index 000000000..49b2ee483
--- /dev/null
+++ b/egs/tedlium3/ASR/conformer_ctc2/asr_datamodule.py
@@ -0,0 +1 @@
+../transducer_stateless/asr_datamodule.py
\ No newline at end of file
diff --git a/egs/tedlium3/ASR/conformer_ctc2/attention.py b/egs/tedlium3/ASR/conformer_ctc2/attention.py
new file mode 100644
index 000000000..178cd7e62
--- /dev/null
+++ b/egs/tedlium3/ASR/conformer_ctc2/attention.py
@@ -0,0 +1,201 @@
+# Copyright 2022 Behavox LLC. (author: Daniil Kulko)
+#
+# 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.
+
+from typing import Optional, Tuple, Union
+
+import torch
+from scaling import ScaledLinear
+
+
+class MultiheadAttention(torch.nn.Module):
+ """Allows the model to jointly attend to information
+ from different representation subspaces. This is a modified
+ version of the original version of multihead attention
+ (see Attention Is All You Need )
+ with replacement of input / output projection layers
+ with newly introduced ScaleLinear layer
+ (see https://github.com/k2-fsa/icefall/blob/master/egs/librispeech/ASR/pruned_transducer_stateless2/scaling.py).
+
+ Args:
+ embed_dim:
+ total dimension of the model.
+ num_heads:
+ number of parallel attention heads. Note that embed_dim will be split
+ across num_heads, i.e. each head will have dimension (embed_dim // num_heads).
+ dropout:
+ dropout probability on attn_output_weights. (default=0.0).
+ bias:
+ if specified, adds bias to input / output projection layers (default=True).
+ add_bias_kv:
+ if specified, adds bias to the key and value sequences at dim=0 (default=False).
+ add_zero_attn:
+ if specified, adds a new batch of zeros to the key and value sequences
+ at dim=1 (default=False).
+ batch_first:
+ if True, then the input and output tensors are provided as
+ (batch, seq, feature), otherwise (seq, batch, feature) (default=False).
+
+ Examples::
+ >>> multihead_attn = MultiheadAttention(embed_dim, num_heads)
+ >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
+ """
+
+ def __init__(
+ self,
+ embed_dim: int,
+ num_heads: int,
+ dropout: float = 0.0,
+ bias: bool = True,
+ add_bias_kv: bool = False,
+ add_zero_attn: bool = False,
+ batch_first: bool = False,
+ device: Union[torch.device, str, None] = None,
+ dtype: Union[torch.dtype, str, None] = None,
+ ) -> None:
+
+ super().__init__()
+
+ self.embed_dim = embed_dim
+ self.num_heads = num_heads
+ self.dropout = dropout
+ self.batch_first = batch_first
+
+ if embed_dim % num_heads != 0:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads. "
+ "Got embedding dim vs number 0f heads: "
+ f"{embed_dim} vs {num_heads}"
+ )
+
+ self.head_dim = embed_dim // num_heads
+
+ self.in_proj = ScaledLinear(
+ embed_dim,
+ 3 * embed_dim,
+ bias=bias,
+ device=device,
+ dtype=dtype,
+ )
+ self.out_proj = ScaledLinear(
+ embed_dim,
+ embed_dim,
+ bias=bias,
+ initial_scale=0.25,
+ device=device,
+ dtype=dtype,
+ )
+
+ if add_bias_kv:
+ self.bias_k = torch.nn.Parameter(
+ torch.empty((1, 1, embed_dim), device=device, dtype=dtype)
+ )
+ self.bias_v = torch.nn.Parameter(
+ torch.empty((1, 1, embed_dim), device=device, dtype=dtype)
+ )
+ else:
+ self.register_parameter("bias_k", None)
+ self.register_parameter("bias_v", None)
+
+ self.add_zero_attn = add_zero_attn
+
+ self._reset_parameters()
+
+ def _reset_parameters(self) -> None:
+ if self.bias_k is not None:
+ torch.nn.init.xavier_normal_(self.bias_k)
+ if self.bias_v is not None:
+ torch.nn.init.xavier_normal_(self.bias_v)
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ key_padding_mask: Optional[torch.Tensor] = None,
+ need_weights: bool = True,
+ attn_mask: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ """
+ Args:
+ query:
+ Query embeddings of shape (L, N, E_q) when batch_first=False or (N, L, E_q)
+ when batch_first=True, where L is the target sequence length, N is the batch size,
+ and E_q is the query embedding dimension embed_dim. Queries are compared against
+ key-value pairs to produce the output. See "Attention Is All You Need" for more details.
+ key:
+ Key embeddings of shape (S, N, E_k) when batch_first=False or (N, S, E_k) when
+ batch_first=True, where S is the source sequence length, N is the batch size, and
+ E_k is the key embedding dimension kdim. See "Attention Is All You Need" for more details.
+ value:
+ Value embeddings of shape (S, N, E_v) when batch_first=False or (N, S, E_v) when
+ batch_first=True, where S is the source sequence length, N is the batch size, and
+ E_v is the value embedding dimension vdim. See "Attention Is All You Need" for more details.
+ key_padding_mask:
+ If specified, a mask of shape (N, S) indicating which elements within key
+ to ignore for the purpose of attention (i.e. treat as "padding").
+ Binary and byte masks are supported. For a binary mask, a True value indicates
+ that the corresponding key value will be ignored for the purpose of attention.
+ For a byte mask, a non-zero value indicates that the corresponding key value will be ignored.
+ need_weights:
+ If specifid, returns attn_output_weights in addition to attn_outputs (default=True).
+ attn_mask:
+ If specified, a 2D or 3D mask preventing attention to certain positions. Must be of shape
+ (L, S) or (N * num_heads, L, S), where N is the batch size, L is the target sequence length,
+ and S is the source sequence length. A 2D mask will be broadcasted across the batch while
+ a 3D mask allows for a different mask for each entry in the batch.
+ Binary, byte, and float masks are supported. For a binary mask, a True value indicates
+ that the corresponding position is not allowed to attend. For a byte mask, a non-zero
+ value indicates that the corresponding position is not allowed to attend. For a float mask,
+ the mask values will be added to the attention weight.
+
+ Returns:
+ attn_output:
+ Attention outputs of shape (L, N, E) when batch_first=False or (N, L, E) when batch_first=True,
+ where L is the target sequence length, N is the batch size, and E is the embedding dimension
+ embed_dim.
+ attn_output_weights:
+ Attention output weights of shape (N, L, S), where N is the batch size, L is the target sequence
+ length, and S is the source sequence length. Only returned when need_weights=True.
+ """
+ if self.batch_first:
+ query, key, value = [x.transpose(1, 0) for x in (query, key, value)]
+
+ (
+ attn_output,
+ attn_output_weights,
+ ) = torch.nn.functional.multi_head_attention_forward(
+ query,
+ key,
+ value,
+ self.embed_dim,
+ self.num_heads,
+ in_proj_weight=self.in_proj.get_weight(),
+ in_proj_bias=self.in_proj.get_bias(),
+ bias_k=self.bias_k,
+ bias_v=self.bias_v,
+ add_zero_attn=self.add_zero_attn,
+ dropout_p=self.dropout,
+ out_proj_weight=self.out_proj.get_weight(),
+ out_proj_bias=self.out_proj.get_bias(),
+ training=self.training,
+ key_padding_mask=key_padding_mask,
+ need_weights=need_weights,
+ attn_mask=attn_mask,
+ )
+
+ if self.batch_first:
+ return attn_output.transpose(1, 0), attn_output_weights
+ return attn_output, attn_output_weights
diff --git a/egs/tedlium3/ASR/conformer_ctc2/combiner.py b/egs/tedlium3/ASR/conformer_ctc2/combiner.py
new file mode 100644
index 000000000..ff526029d
--- /dev/null
+++ b/egs/tedlium3/ASR/conformer_ctc2/combiner.py
@@ -0,0 +1,244 @@
+# Copyright 2022 Behavox LLC. (author: Daniil Kulko)
+#
+# 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.
+
+from typing import List
+
+import torch
+
+
+class RandomCombine(torch.nn.Module):
+ """
+ This module combines a list of Tensors, all with the same shape, to
+ produce a single output of that same shape which, in training time,
+ is a random combination of all the inputs; but which in test time
+ will be just the last input.
+ The idea is that the list of Tensors will be a list of outputs of multiple
+ conformer layers. This has a similar effect as iterated loss. (See:
+ DEJA-VU: DOUBLE FEATURE PRESENTATION AND ITERATED LOSS IN DEEP TRANSFORMER
+ NETWORKS).
+ """
+
+ def __init__(
+ self,
+ num_inputs: int,
+ final_weight: float = 0.5,
+ pure_prob: float = 0.5,
+ stddev: float = 2.0,
+ ) -> None:
+ """
+ Args:
+ num_inputs:
+ The number of tensor inputs, which equals the number of layers'
+ outputs that are fed into this module. E.g. in an 18-layer neural
+ net if we output layers 16, 12, 18, num_inputs would be 3.
+ final_weight:
+ The amount of weight or probability we assign to the
+ final layer when randomly choosing layers or when choosing
+ continuous layer weights.
+ pure_prob:
+ The probability, on each frame, with which we choose
+ only a single layer to output (rather than an interpolation)
+ stddev:
+ A standard deviation that we add to log-probs for computing
+ randomized weights.
+ The method of choosing which layers, or combinations of layers, to use,
+ is conceptually as follows::
+ With probability `pure_prob`::
+ With probability `final_weight`: choose final layer,
+ Else: choose random non-final layer.
+ Else::
+ Choose initial log-weights that correspond to assigning
+ weight `final_weight` to the final layer and equal
+ weights to other layers; then add Gaussian noise
+ with variance `stddev` to these log-weights, and normalize
+ to weights (note: the average weight assigned to the
+ final layer here will not be `final_weight` if stddev>0).
+ """
+ super().__init__()
+ assert 0 <= pure_prob <= 1, pure_prob
+ assert 0 < final_weight < 1, final_weight
+ assert num_inputs >= 1, num_inputs
+
+ self.num_inputs = num_inputs
+ self.final_weight = final_weight
+ self.pure_prob = pure_prob
+ self.stddev = stddev
+
+ self.final_log_weight = (
+ torch.tensor((final_weight / (1 - final_weight)) * (self.num_inputs - 1))
+ .log()
+ .item()
+ )
+
+ def forward(self, inputs: List[torch.Tensor]) -> torch.Tensor:
+ """Forward function.
+ Args:
+ inputs:
+ A list of Tensor, e.g. from various layers of a transformer.
+ All must be the same shape, of (*, num_channels)
+ Returns:
+ A Tensor of shape (*, num_channels). In test mode
+ this is just the final input.
+ """
+ num_inputs = self.num_inputs
+ assert len(inputs) == num_inputs, f"{len(inputs)}, {num_inputs}"
+ if not self.training or torch.jit.is_scripting() or len(inputs) == 1:
+ return inputs[-1]
+
+ # Shape of weights: (*, num_inputs)
+ num_channels = inputs[0].shape[-1]
+ num_frames = inputs[0].numel() // num_channels
+
+ ndim = inputs[0].ndim
+ # stacked_inputs: (num_frames, num_channels, num_inputs)
+ stacked_inputs = torch.stack(inputs, dim=ndim).reshape(
+ (num_frames, num_channels, num_inputs)
+ )
+
+ # weights: (num_frames, num_inputs)
+ weights = self._get_random_weights(
+ inputs[0].dtype, inputs[0].device, num_frames
+ )
+
+ weights = weights.reshape(num_frames, num_inputs, 1)
+ # ans: (num_frames, num_channels, 1)
+ ans = torch.matmul(stacked_inputs, weights)
+ # ans: (*, num_channels)
+
+ ans = ans.reshape(inputs[0].shape[:-1] + (num_channels,))
+
+ return ans
+
+ def _get_random_weights(
+ self, dtype: torch.dtype, device: torch.device, num_frames: int
+ ) -> torch.Tensor:
+ """Return a tensor of random weights, of shape
+ `(num_frames, self.num_inputs)`,
+ Args:
+ dtype:
+ The data-type desired for the answer, e.g. float, double.
+ device:
+ The device needed for the answer.
+ num_frames:
+ The number of sets of weights desired
+ Returns:
+ A tensor of shape (num_frames, self.num_inputs), such that
+ `ans.sum(dim=1)` is all ones.
+ """
+ pure_prob = self.pure_prob
+ if pure_prob == 0.0:
+ return self._get_random_mixed_weights(dtype, device, num_frames)
+ elif pure_prob == 1.0:
+ return self._get_random_pure_weights(dtype, device, num_frames)
+ else:
+ p = self._get_random_pure_weights(dtype, device, num_frames)
+ m = self._get_random_mixed_weights(dtype, device, num_frames)
+ return torch.where(
+ torch.rand(num_frames, 1, device=device) < self.pure_prob, p, m
+ )
+
+ def _get_random_pure_weights(
+ self, dtype: torch.dtype, device: torch.device, num_frames: int
+ ) -> torch.Tensor:
+ """Return a tensor of random one-hot weights, of shape
+ `(num_frames, self.num_inputs)`,
+ Args:
+ dtype:
+ The data-type desired for the answer, e.g. float, double.
+ device:
+ The device needed for the answer.
+ num_frames:
+ The number of sets of weights desired.
+ Returns:
+ A one-hot tensor of shape `(num_frames, self.num_inputs)`, with
+ exactly one weight equal to 1.0 on each frame.
+ """
+ final_prob = self.final_weight
+
+ # final contains self.num_inputs - 1 in all elements
+ final = torch.full((num_frames,), self.num_inputs - 1, device=device)
+ # nonfinal contains random integers in [0..num_inputs - 2], these are for non-final weights.
+ nonfinal = torch.randint(self.num_inputs - 1, (num_frames,), device=device)
+
+ indexes = torch.where(
+ torch.rand(num_frames, device=device) < final_prob, final, nonfinal
+ )
+ ans = torch.nn.functional.one_hot(indexes, num_classes=self.num_inputs).to(
+ dtype=dtype
+ )
+ return ans
+
+ def _get_random_mixed_weights(
+ self, dtype: torch.dtype, device: torch.device, num_frames: int
+ ) -> torch.Tensor:
+ """Return a tensor of random one-hot weights, of shape
+ `(num_frames, self.num_inputs)`,
+ Args:
+ dtype:
+ The data-type desired for the answer, e.g. float, double.
+ device:
+ The device needed for the answer.
+ num_frames:
+ The number of sets of weights desired.
+ Returns:
+ A tensor of shape (num_frames, self.num_inputs), which elements
+ in [0..1] that sum to one over the second axis, i.e.
+ `ans.sum(dim=1)` is all ones.
+ """
+ logprobs = (
+ torch.randn(num_frames, self.num_inputs, dtype=dtype, device=device)
+ * self.stddev
+ )
+ logprobs[:, -1] += self.final_log_weight
+ return logprobs.softmax(dim=1)
+
+
+def _test_random_combine(
+ final_weight: float,
+ pure_prob: float,
+ stddev: float,
+) -> None:
+ print(
+ f"_test_random_combine: final_weight={final_weight}, "
+ f"pure_prob={pure_prob}, stddev={stddev}"
+ )
+ num_inputs = 3
+ num_channels = 50
+ m = RandomCombine(
+ num_inputs=num_inputs,
+ final_weight=final_weight,
+ pure_prob=pure_prob,
+ stddev=stddev,
+ )
+
+ x = [torch.ones(3, 4, num_channels) for _ in range(num_inputs)]
+
+ y = m(x)
+ assert y.shape == x[0].shape
+ assert torch.allclose(y, x[0]) # .. since actually all ones.
+
+
+def _test_random_combine_main() -> None:
+ _test_random_combine(0.999, 0, 0.0)
+ _test_random_combine(0.5, 0, 0.0)
+ _test_random_combine(0.999, 0, 0.0)
+ _test_random_combine(0.5, 0, 0.3)
+ _test_random_combine(0.5, 1, 0.3)
+ _test_random_combine(0.5, 0.5, 0.3)
+
+
+if __name__ == "__main__":
+ _test_random_combine_main()
diff --git a/egs/tedlium3/ASR/conformer_ctc2/conformer.py b/egs/tedlium3/ASR/conformer_ctc2/conformer.py
new file mode 100644
index 000000000..fad2f371f
--- /dev/null
+++ b/egs/tedlium3/ASR/conformer_ctc2/conformer.py
@@ -0,0 +1,1033 @@
+#!/usr/bin/env python3
+# Copyright (c) 2021 University of Chinese Academy of Sciences (author: Han Zhu)
+# 2022 Xiaomi Corp. (author: Quandong Wang)
+#
+# 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.
+
+import copy
+import math
+import warnings
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from combiner import RandomCombine
+from scaling import (
+ ActivationBalancer,
+ BasicNorm,
+ DoubleSwish,
+ ScaledConv1d,
+ ScaledLinear,
+)
+from subsampling import Conv2dSubsampling
+from transformer import Supervisions, Transformer, encoder_padding_mask
+
+
+class Conformer(Transformer):
+ 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,
+ layer_dropout: float = 0.075,
+ cnn_module_kernel: int = 31,
+ aux_layer_period: int = 3,
+ ) -> None:
+ """
+ Args:
+ num_features (int):
+ number of input features.
+ num_classes (int):
+ number of output classes.
+ subsampling_factor (int):
+ subsampling factor of encoder;
+ currently, subsampling_factor MUST be 4.
+ d_model (int):
+ attention dimension, also the output dimension.
+ nhead (int):
+ number of heads in multi-head attention;
+ must satisfy d_model // nhead == 0.
+ 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.
+ layer_dropout (float):
+ layer-dropout rate.
+ cnn_module_kernel (int):
+ kernel size of convolution module.
+ aux_layer_period (int):
+ determines the auxiliary encoder layers.
+ """
+
+ super().__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,
+ layer_dropout=layer_dropout,
+ )
+
+ self.num_features = num_features
+ self.subsampling_factor = subsampling_factor
+ if subsampling_factor != 4:
+ raise NotImplementedError("Support only 'subsampling_factor=4'.")
+
+ # self.encoder_embed converts the input of shape (N, T, num_features)
+ # to the shape (N, T//subsampling_factor, d_model).
+ # That is, it does two things simultaneously:
+ # (1) subsampling: T -> T//subsampling_factor
+ # (2) embedding: num_features -> d_model
+ self.encoder_embed = Conv2dSubsampling(num_features, d_model)
+
+ self.encoder_pos = RelPositionalEncoding(d_model, dropout)
+
+ encoder_layer = ConformerEncoderLayer(
+ d_model=d_model,
+ nhead=nhead,
+ dim_feedforward=dim_feedforward,
+ dropout=dropout,
+ layer_dropout=layer_dropout,
+ cnn_module_kernel=cnn_module_kernel,
+ )
+
+ # aux_layers from 1/3
+ self.encoder = ConformerEncoder(
+ encoder_layer=encoder_layer,
+ num_layers=num_encoder_layers,
+ aux_layers=list(
+ range(
+ num_encoder_layers // 3,
+ num_encoder_layers - 1,
+ aux_layer_period,
+ )
+ ),
+ )
+
+ def run_encoder(
+ self,
+ x: torch.Tensor,
+ supervisions: Optional[Supervisions] = None,
+ warmup: float = 1.0,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ """
+ Args:
+ x:
+ the input tensor. Its shape is (batch_size, seq_len, feature_dim).
+ supervisions:
+ Supervision in lhotse format.
+ See https://github.com/lhotse-speech/lhotse/blob/master/lhotse/dataset/speech_recognition.py#L32 # noqa
+ CAUTION: It contains length information, i.e., start and number of
+ frames, before subsampling
+ It is read directly from the batch, without any sorting. It is used
+ to compute encoder padding mask, which is used as memory key padding
+ mask for the decoder.
+ warmup:
+ a floating point value that gradually increases from 0 throughout
+ training; when it is >= 1.0 we are "fully warmed up". It is used
+ to turn modules on sequentially.
+
+ Returns:
+ torch.Tensor: Predictor tensor of dimension (S, N, C).
+ torch.Tensor: Mask tensor of dimension (N, S)
+ """
+ x = self.encoder_embed(x)
+ x, pos_emb = self.encoder_pos(x)
+ x = x.permute(1, 0, 2) # (N, S, C) -> (S, N, C)
+ mask = encoder_padding_mask(x.size(0), supervisions)
+ mask = mask.to(x.device) if mask is not None else None
+
+ x = self.encoder(
+ x, pos_emb, src_key_padding_mask=mask, warmup=warmup
+ ) # (S, N, C)
+
+ 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"
+
+ 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,
+ bypass_scale: float = 0.1,
+ layer_dropout: float = 0.075,
+ cnn_module_kernel: int = 31,
+ ) -> None:
+ """
+ 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).
+ bypass_scale:
+ a scale on the layer's output, used in bypass (resnet-type) skip-connection;
+ when the layer is bypassed the final output will be a
+ weighted sum of the layer's input and layer's output with weights
+ (1.0-bypass_scale) and bypass_scale correspondingly (default=0.1).
+ layer_dropout:
+ the probability to bypass the layer (default=0.075).
+ cnn_module_kernel (int):
+ kernel size of convolution module (default=31).
+ """
+ super().__init__()
+
+ if bypass_scale < 0.0 or bypass_scale > 1.0:
+ raise ValueError("bypass_scale should be between 0.0 and 1.0")
+
+ if layer_dropout < 0.0 or layer_dropout > 1.0:
+ raise ValueError("layer_dropout should be between 0.0 and 1.0")
+
+ self.bypass_scale = bypass_scale
+ self.layer_dropout = layer_dropout
+
+ self.self_attn = RelPositionMultiheadAttention(d_model, nhead, dropout=0.0)
+
+ self.feed_forward = nn.Sequential(
+ ScaledLinear(d_model, dim_feedforward),
+ ActivationBalancer(channel_dim=-1),
+ DoubleSwish(),
+ nn.Dropout(dropout),
+ ScaledLinear(dim_feedforward, d_model, initial_scale=0.25),
+ )
+
+ self.feed_forward_macaron = nn.Sequential(
+ ScaledLinear(d_model, dim_feedforward),
+ ActivationBalancer(channel_dim=-1),
+ DoubleSwish(),
+ nn.Dropout(dropout),
+ ScaledLinear(dim_feedforward, d_model, initial_scale=0.25),
+ )
+
+ self.conv_module = ConvolutionModule(d_model, cnn_module_kernel)
+
+ self.norm_final = BasicNorm(d_model)
+
+ # try to ensure the output is close to zero-mean (or at least, zero-median).
+ self.balancer = ActivationBalancer(
+ channel_dim=-1, min_positive=0.45, max_positive=0.55, max_abs=6.0
+ )
+
+ self.dropout = nn.Dropout(dropout)
+
+ def forward(
+ self,
+ src: torch.Tensor,
+ pos_emb: torch.Tensor,
+ src_mask: Optional[torch.Tensor] = None,
+ src_key_padding_mask: Optional[torch.Tensor] = None,
+ warmup: float = 1.0,
+ ) -> torch.Tensor:
+ """
+ Pass the input through the encoder layer.
+
+ Args:
+ src:
+ the sequence to the encoder layer of shape (S, N, C) (required).
+ pos_emb:
+ positional embedding tensor of shape (N, 2*S-1, C) (required).
+ src_mask:
+ the mask for the src sequence of shape (S, S) (optional).
+ src_key_padding_mask:
+ the mask for the src keys per batch of shape (N, S) (optional).
+ warmup:
+ controls selective bypass of of layers; if < 1.0, we will
+ bypass layers more frequently.
+
+ Returns:
+ Output tensor of the shape (S, N, C), where
+ S is the source sequence length,
+ N is the batch size,
+ C is the feature number
+ """
+ src_orig = src
+
+ warmup_scale = min(self.bypass_scale + warmup, 1.0)
+ # alpha = 1.0 means fully use this encoder layer, 0.0 would mean
+ # completely bypass it.
+ if self.training:
+ alpha = (
+ warmup_scale
+ if torch.rand(()).item() <= (1.0 - self.layer_dropout)
+ else self.bypass_scale
+ )
+ else:
+ alpha = 1.0
+
+ # macaron style feed forward module
+ src = src + self.dropout(self.feed_forward_macaron(src))
+
+ # multi-headed self-attention module
+ 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 = src + self.dropout(src_att)
+
+ # convolution module
+ src = src + self.dropout(self.conv_module(src))
+
+ # feed forward module
+ src = src + self.dropout(self.feed_forward(src))
+
+ src = self.norm_final(self.balancer(src))
+
+ if alpha != 1.0:
+ src = alpha * src + (1 - alpha) * src_orig
+
+ return src
+
+
+class ConformerEncoder(nn.Module):
+ """
+ ConformerEncoder is a stack of N encoder layers
+
+ 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,
+ aux_layers: List[int],
+ ) -> None:
+
+ """
+ Args:
+ encoder_layer:
+ an instance of the ConformerEncoderLayer() class (required).
+ num_layers:
+ the number of sub-encoder-layers in the encoder (required).
+ aux_layers:
+ list of indexes of sub-encoder-layers outputs to be combined (required).
+ """
+
+ super().__init__()
+ self.layers = nn.ModuleList(
+ [copy.deepcopy(encoder_layer) for i in range(num_layers)]
+ )
+ self.num_layers = num_layers
+
+ assert len(set(aux_layers)) == len(aux_layers)
+
+ assert num_layers - 1 not in aux_layers
+ self.aux_layers = aux_layers + [num_layers - 1]
+
+ self.combiner = RandomCombine(
+ num_inputs=len(self.aux_layers),
+ final_weight=0.5,
+ pure_prob=0.333,
+ stddev=2.0,
+ )
+
+ def forward(
+ self,
+ src: torch.Tensor,
+ pos_emb: torch.Tensor,
+ mask: Optional[torch.Tensor] = None,
+ src_key_padding_mask: Optional[torch.Tensor] = None,
+ warmup: float = 1.0,
+ ) -> torch.Tensor:
+ """
+ Pass the input through the encoder layers in turn.
+
+ Args:
+ src:
+ the sequence to the encoder of shape (S, N, C) (required).
+ pos_emb:
+ positional embedding tensor of shape (N, 2*S-1, C) (required).
+ mask:
+ the mask for the src sequence of shape (S, S) (optional).
+ src_key_padding_mask:
+ the mask for the src keys per batch of shape (N, S) (optional).
+ warmup:
+ controls selective bypass of layer; if < 1.0, we will
+ bypass the layer more frequently (default=1.0).
+
+ Returns:
+ Output tensor of the shape (S, N, C), where
+ S is the source sequence length,
+ N is the batch size,
+ C is the feature number.
+
+ """
+ output = src
+
+ outputs = []
+ for i, mod in enumerate(self.layers):
+ output = mod(
+ output,
+ pos_emb,
+ src_mask=mask,
+ src_key_padding_mask=src_key_padding_mask,
+ warmup=warmup,
+ )
+
+ if i in self.aux_layers:
+ outputs.append(output)
+
+ output = self.combiner(outputs)
+
+ 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
+
+ """
+
+ def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000) -> None:
+ """
+ Construct an PositionalEncoding object.
+
+ Args:
+ d_model: Embedding dimension.
+ dropout_rate: Dropout rate.
+ max_len: Maximum input length.
+
+ """
+ super().__init__()
+ self.d_model = 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: torch.Tensor) -> None:
+ """
+ Reset the positional encodings.
+
+ Args:
+ x:
+ input tensor (N, T, C), where
+ T is the source sequence length,
+ N is the batch size.
+ C is the feature number.
+
+ """
+ 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[torch.Tensor, torch.Tensor]:
+ """
+ Add positional encoding.
+
+ Args:
+ x:
+ input tensor (N, T, C).
+
+ Returns:
+ torch.Tensor: Encoded tensor (N, T, C).
+ torch.Tensor: Encoded tensor (N, 2*T-1, C), where
+ T is the source sequence length,
+ N is the batch size.
+ C is the feature number.
+
+ """
+ self.extend_pe(x)
+ pos_emb = self.pe[
+ :,
+ self.pe.size(1) // 2
+ - x.size(1)
+ + 1 : self.pe.size(1) // 2 # noqa E203
+ + x.size(1),
+ ]
+ return self.dropout(x), self.dropout(pos_emb)
+
+
+class RelPositionMultiheadAttention(nn.Module):
+ """
+ Multi-Head Attention layer with relative position encoding
+ See reference: "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context".
+
+ """
+
+ def __init__(
+ self,
+ embed_dim: int,
+ num_heads: int,
+ dropout: float = 0.0,
+ ) -> None:
+ """
+ Args:
+ embed_dim:
+ total dimension of the model.
+ num_heads:
+ parallel attention heads.
+ dropout:
+ a Dropout layer on attn_output_weights. Default: 0.0.
+ """
+ super().__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 = ScaledLinear(embed_dim, 3 * embed_dim, bias=True)
+ self.out_proj = ScaledLinear(
+ embed_dim, embed_dim, bias=True, initial_scale=0.25
+ )
+
+ # linear transformation for positional encoding.
+ self.linear_pos = ScaledLinear(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.pos_bias_u_scale = nn.Parameter(torch.zeros(()).detach())
+ self.pos_bias_v_scale = nn.Parameter(torch.zeros(()).detach())
+ self._reset_parameters()
+
+ def _pos_bias_u(self):
+ return self.pos_bias_u * self.pos_bias_u_scale.exp()
+
+ def _pos_bias_v(self):
+ return self.pos_bias_v * self.pos_bias_v_scale.exp()
+
+ def _reset_parameters(self) -> None:
+ nn.init.normal_(self.pos_bias_u, std=0.01)
+ nn.init.normal_(self.pos_bias_v, std=0.01)
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ pos_emb: torch.Tensor,
+ key_padding_mask: Optional[torch.Tensor] = None,
+ need_weights: bool = False,
+ attn_mask: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ """
+ 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.get_weight(),
+ self.in_proj.get_bias(),
+ self.dropout,
+ self.out_proj.get_weight(),
+ self.out_proj.get_bias(),
+ training=self.training,
+ key_padding_mask=key_padding_mask,
+ need_weights=need_weights,
+ attn_mask=attn_mask,
+ )
+
+ def rel_shift(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Compute relative positional encoding.
+
+ Args:
+ x:
+ input tensor (batch, head, time1, 2*time1-1).
+ time1 means the length of query vector.
+
+ Returns:
+ torch.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: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ pos_emb: torch.Tensor,
+ embed_dim_to_check: int,
+ num_heads: int,
+ in_proj_weight: torch.Tensor,
+ in_proj_bias: torch.Tensor,
+ dropout_p: float,
+ out_proj_weight: torch.Tensor,
+ out_proj_bias: torch.Tensor,
+ training: bool = True,
+ key_padding_mask: Optional[torch.Tensor] = None,
+ need_weights: bool = False,
+ attn_mask: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ """
+ 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 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(
+ f"attn_mask's dimension {attn_mask.dim()} is not supported"
+ )
+ # 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 * scaling).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)
+
+ attn_output_weights = matrix_ac + matrix_bd # (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):
+ def __init__(self, channels: int, kernel_size: int, bias: bool = True) -> None:
+ """
+ ConvolutionModule in Conformer model.
+ Modified from https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/conformer/convolution.py
+ Construct a ConvolutionModule object.
+
+ 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).
+ """
+ super().__init__()
+ # kernerl_size should be a odd number for 'SAME' padding
+ assert (kernel_size - 1) % 2 == 0
+
+ self.pointwise_conv1 = ScaledConv1d(
+ channels,
+ 2 * channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ bias=bias,
+ )
+
+ # after pointwise_conv1 we put x through a gated linear unit (nn.functional.glu).
+ # For most layers the normal rms value of channels of x seems to be in the range 1 to 4,
+ # but sometimes, for some reason, for layer 0 the rms ends up being very large,
+ # between 50 and 100 for different channels. This will cause very peaky and
+ # sparse derivatives for the sigmoid gating function, which will tend to make
+ # the loss function not learn effectively. (for most layers the average absolute values
+ # are in the range 0.5..9.0, and the average p(x>0), i.e. positive proportion,
+ # at the output of pointwise_conv1.output is around 0.35 to 0.45 for different
+ # layers, which likely breaks down as 0.5 for the "linear" half and
+ # 0.2 to 0.3 for the part that goes into the sigmoid. The idea is that if we
+ # constrain the rms values to a reasonable range via a constraint of max_abs=10.0,
+ # it will be in a better position to start learning something, i.e. to latch onto
+ # the correct range.
+ self.deriv_balancer1 = ActivationBalancer(
+ channel_dim=1, max_abs=10.0, min_positive=0.05, max_positive=1.0
+ )
+
+ self.depthwise_conv = ScaledConv1d(
+ channels,
+ channels,
+ kernel_size,
+ stride=1,
+ padding=(kernel_size - 1) // 2,
+ groups=channels,
+ bias=bias,
+ )
+
+ self.deriv_balancer2 = ActivationBalancer(
+ channel_dim=1, min_positive=0.05, max_positive=1.0
+ )
+
+ self.activation = DoubleSwish()
+
+ self.pointwise_conv2 = ScaledConv1d(
+ channels,
+ channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ bias=bias,
+ initial_scale=0.25,
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """Compute convolution module.
+
+ Args:
+ x:
+ input tensor of shape (T, N, C).
+
+ Returns:
+ torch.Tensor: Output tensor (T, N, C), where
+ T is the source sequence length,
+ N is the batch size,
+ C is the feature number.
+
+ """
+ # 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 = self.deriv_balancer1(x)
+ x = nn.functional.glu(x, dim=1) # (batch, channels, time)
+
+ # 1D Depthwise Conv
+ x = self.depthwise_conv(x)
+
+ x = self.deriv_balancer2(x)
+ x = self.activation(x)
+
+ x = self.pointwise_conv2(x) # (batch, channel, time)
+
+ return x.permute(2, 0, 1)
diff --git a/egs/tedlium3/ASR/conformer_ctc2/decode.py b/egs/tedlium3/ASR/conformer_ctc2/decode.py
new file mode 100755
index 000000000..28d39de70
--- /dev/null
+++ b/egs/tedlium3/ASR/conformer_ctc2/decode.py
@@ -0,0 +1,896 @@
+#!/usr/bin/env python3
+# Copyright 2021 Xiaomi Corporation (Author: Liyong Guo,
+# Fangjun Kuang,
+# Quandong Wang)
+#
+# 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.
+
+
+import argparse
+import logging
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import k2
+import sentencepiece as spm
+import torch
+import torch.nn as nn
+from asr_datamodule import TedLiumAsrDataModule
+from conformer import Conformer
+from train import add_model_arguments
+
+from icefall.bpe_graph_compiler import BpeCtcTrainingGraphCompiler
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from icefall.decode import (
+ get_lattice,
+ nbest_decoding,
+ nbest_oracle,
+ one_best_decoding,
+ rescore_with_attention_decoder,
+ rescore_with_n_best_list,
+ rescore_with_whole_lattice,
+)
+from icefall.env import get_env_info
+from icefall.lexicon import Lexicon
+from icefall.utils import (
+ AttributeDict,
+ get_texts,
+ load_averaged_model,
+ setup_logger,
+ store_transcripts,
+ str2bool,
+ write_error_stats,
+)
+
+
+def get_parser() -> argparse.ArgumentParser:
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=30,
+ help="""It specifies the checkpoint to use for decoding.
+ Note: Epoch counts from 1.
+ You can specify --avg to use more checkpoints for model averaging.""",
+ )
+
+ parser.add_argument(
+ "--iter",
+ type=int,
+ default=0,
+ help="""If positive, --epoch is ignored and it
+ will use the checkpoint exp_dir/checkpoint-iter.pt.
+ You can specify --avg to use more checkpoints for model averaging.
+ """,
+ )
+
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=15,
+ help="Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch' and '--iter'",
+ )
+
+ parser.add_argument(
+ "--method",
+ type=str,
+ default="attention-decoder",
+ help="""Decoding method.
+ Supported values are:
+ - (0) ctc-decoding. Use CTC decoding. It uses a sentence piece
+ model, i.e., lang_dir/bpe.model, to convert word pieces to words.
+ It needs neither a lexicon nor an n-gram LM.
+ - (1) ctc-greedy-search. It only use CTC output and a sentence piece
+ model for decoding. It produces the same results with ctc-decoding.
+ - (2) 1best. Extract the best path from the decoding lattice as the
+ decoding result.
+ - (3) nbest. Extract n paths from the decoding lattice; the path
+ with the highest score is the decoding result.
+ - (4) nbest-rescoring. Extract n paths from the decoding lattice,
+ rescore them with an n-gram LM (e.g., a 4-gram LM), the path with
+ the highest score is the decoding result.
+ - (5) whole-lattice-rescoring. Rescore the decoding lattice with an
+ n-gram LM (e.g., a 4-gram LM), the best path of rescored lattice
+ is the decoding result.
+ - (6) attention-decoder. Extract n paths from the LM rescored
+ lattice, the path with the highest score is the decoding result.
+ - (7) nbest-oracle. Its WER is the lower bound of any n-best
+ rescoring method can achieve. Useful for debugging n-best
+ rescoring method.
+ """,
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=True,
+ help="Whether to load averaged model. Currently it only supports "
+ "using --epoch. If True, it would decode with the averaged model "
+ "over the epoch range from `epoch-avg` (excluded) to `epoch`."
+ "Actually only the models with epoch number of `epoch-avg` and "
+ "`epoch` are loaded for averaging. ",
+ )
+
+ parser.add_argument(
+ "--num-paths",
+ type=int,
+ default=100,
+ help="""Number of paths for n-best based decoding method.
+ Used only when "method" is one of the following values:
+ nbest, nbest-rescoring, attention-decoder, and nbest-oracle
+ """,
+ )
+
+ parser.add_argument(
+ "--nbest-scale",
+ type=float,
+ default=0.5,
+ help="""The scale to be applied to `lattice.scores`.
+ It's needed if you use any kinds of n-best based rescoring.
+ Used only when "method" is one of the following values:
+ nbest, nbest-rescoring, attention-decoder, and nbest-oracle
+ A smaller value results in more unique paths.
+ """,
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="conformer_ctc2/exp",
+ help="The experiment dir",
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=str,
+ default="data/lang_bpe_500",
+ help="The lang dir",
+ )
+
+ parser.add_argument(
+ "--lm-path",
+ type=str,
+ default="data/lm/G_4_gram.pt",
+ help="""The n-gram LM dir for rescoring.
+ It should contain either lm_fname.pt or lm_fname.fst.txt
+ """,
+ )
+
+ parser.add_argument(
+ "--result-dir",
+ type=str,
+ default="conformer_ctc2/exp/results",
+ help="Directory to store results.",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def get_params() -> AttributeDict:
+ """Return a dict containing training parameters.
+
+ All training related parameters that are not passed from the commandline
+ are 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`:
+
+ - feature_dim: The model input dim. It has to match the one used
+ in computing features.
+
+ - subsampling_factor: The subsampling factor for the model.
+ """
+ params = AttributeDict(
+ {
+ # parameters for conformer
+ "subsampling_factor": 4,
+ "feature_dim": 80,
+ # parameters for decoding
+ "search_beam": 15,
+ "output_beam": 8,
+ "min_active_states": 10,
+ "max_active_states": 7000,
+ "use_double_scores": True,
+ "env_info": get_env_info(),
+ }
+ )
+ return params
+
+
+def ctc_greedy_search(
+ ctc_probs: torch.Tensor,
+ mask: torch.Tensor,
+) -> List[List[int]]:
+ """Apply CTC greedy search
+ Args:
+ ctc_probs (torch.Tensor): (batch, max_len, num_bpe)
+ mask (torch.Tensor): (batch, max_len)
+ Returns:
+ best path result
+ """
+
+ _, max_index = ctc_probs.max(2) # (B, maxlen)
+ max_index = max_index.masked_fill_(mask, 0) # (B, maxlen)
+
+ ret_hyps = []
+ for hyp in max_index:
+ hyp = torch.unique_consecutive(hyp)
+ hyp = hyp[hyp > 0].tolist()
+ ret_hyps.append(hyp)
+ return ret_hyps
+
+
+def decode_one_batch(
+ params: AttributeDict,
+ model: nn.Module,
+ HLG: Optional[k2.Fsa],
+ H: Optional[k2.Fsa],
+ bpe_model: Optional[spm.SentencePieceProcessor],
+ batch: dict,
+ word_table: k2.SymbolTable,
+ sos_id: int,
+ eos_id: int,
+ G: Optional[k2.Fsa] = None,
+) -> Dict[str, List[List[str]]]:
+ """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. Used only when params.method is NOT ctc-decoding.
+ H:
+ The ctc topo. Used only when params.method is ctc-decoding.
+ bpe_model:
+ The BPE model. Used only when params.method is ctc-decoding.
+ batch:
+ It is the return value from iterating
+ `lhotse.dataset.K2SpeechRecognitionDataset`. See its documentation
+ for the format of the `batch`.
+ word_table:
+ The word symbol table.
+ sos_id:
+ The token ID of the SOS.
+ eos_id:
+ The token ID of the EOS.
+ 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. Note: If it decodes to nothing, then return None.
+ """
+ if HLG is not None:
+ device = HLG.device
+ else:
+ device = H.device
+ feature = batch["inputs"]
+ assert feature.ndim == 3
+ feature = feature.to(device)
+ # at entry, feature is (N, T, C)
+
+ supervisions = batch["supervisions"]
+
+ nnet_output, memory, memory_key_padding_mask = model(feature, supervisions)
+ # nnet_output is (N, T, C)
+
+ supervision_segments = torch.stack(
+ (
+ supervisions["sequence_idx"],
+ torch.div(
+ supervisions["start_frame"],
+ params.subsampling_factor,
+ rounding_mode="floor",
+ ),
+ torch.div(
+ supervisions["num_frames"],
+ params.subsampling_factor,
+ rounding_mode="floor",
+ ),
+ ),
+ 1,
+ ).to(torch.int32)
+
+ if H is None:
+ assert HLG is not None
+ decoding_graph = HLG
+ else:
+ assert HLG is None
+ assert bpe_model is not None
+ decoding_graph = H
+
+ lattice = get_lattice(
+ nnet_output=nnet_output,
+ decoding_graph=decoding_graph,
+ 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 == "ctc-decoding":
+ best_path = one_best_decoding(
+ lattice=lattice, use_double_scores=params.use_double_scores
+ )
+ # Note: `best_path.aux_labels` contains token IDs, not word IDs
+ # since we are using H, not HLG here.
+ #
+ # token_ids is a lit-of-list of IDs
+ token_ids = get_texts(best_path)
+
+ # hyps is a list of str, e.g., ['xxx yyy zzz', ...]
+ hyps = bpe_model.decode(token_ids)
+
+ # hyps is a list of list of str, e.g., [['xxx', 'yyy', 'zzz'], ... ]
+ unk = bpe_model.decode(bpe_model.unk_id()).strip()
+ hyps = [[w for w in s.split() if w != unk] for s in hyps]
+ key = "ctc-decoding"
+
+ return {key: hyps}
+
+ if params.method == "ctc-greedy-search":
+ hyps = ctc_greedy_search(nnet_output, memory_key_padding_mask)
+
+ # hyps is a list of str, e.g., ['xxx yyy zzz', ...]
+ hyps = bpe_model.decode(hyps)
+
+ # hyps is a list of list of str, e.g., [['xxx', 'yyy', 'zzz'], ... ]
+ unk = bpe_model.decode(bpe_model.unk_id()).strip()
+ hyps = [[w for w in s.split() if w != unk] for s in hyps]
+ key = "ctc-greedy-search"
+
+ return {key: hyps}
+
+ if params.method == "nbest-oracle":
+ # Note: You can also pass rescored lattices to it.
+ # We choose the HLG decoded lattice for speed reasons
+ # as HLG decoding is faster and the oracle WER
+ # is only slightly worse than that of rescored lattices.
+ best_path = nbest_oracle(
+ lattice=lattice,
+ num_paths=params.num_paths,
+ ref_texts=supervisions["text"],
+ word_table=word_table,
+ nbest_scale=params.nbest_scale,
+ oov="",
+ )
+ hyps = get_texts(best_path)
+ hyps = [
+ [word_table[i] for i in ids if word_table[i] != ""] for ids in hyps
+ ]
+ key = f"oracle_{params.num_paths}_nbest_scale_{params.nbest_scale}" # noqa
+ return {key: hyps}
+
+ if params.method == "nbest":
+ best_path = nbest_decoding(
+ lattice=lattice,
+ num_paths=params.num_paths,
+ use_double_scores=params.use_double_scores,
+ nbest_scale=params.nbest_scale,
+ )
+ key = f"no_rescore-nbest-scale-{params.nbest_scale}-{params.num_paths}" # noqa
+
+ hyps = get_texts(best_path)
+ hyps = [
+ [word_table[i] for i in ids if word_table[i] != ""] for ids in hyps
+ ]
+ return {key: hyps}
+
+ assert params.method in [
+ "1best",
+ "nbest-rescoring",
+ "whole-lattice-rescoring",
+ "attention-decoder",
+ ]
+
+ lm_scale_list = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]
+ 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 == "1best":
+ best_path_dict = one_best_decoding(
+ lattice=lattice,
+ lm_scale_list=lm_scale_list,
+ )
+ elif 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,
+ nbest_scale=params.nbest_scale,
+ )
+ 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":
+ best_path_dict = rescore_with_attention_decoder(
+ lattice=lattice,
+ num_paths=params.num_paths,
+ model=model,
+ memory=memory,
+ memory_key_padding_mask=memory_key_padding_mask,
+ sos_id=sos_id,
+ eos_id=eos_id,
+ nbest_scale=params.nbest_scale,
+ )
+ else:
+ raise ValueError(f"Unsupported decoding method: {params.method}")
+
+ ans = dict()
+ if best_path_dict is not None:
+ for lm_scale_str, best_path in best_path_dict.items():
+ hyps = get_texts(best_path)
+ hyps = [
+ [word_table[i] for i in ids if word_table[i] != ""] for ids in hyps
+ ]
+ ans[lm_scale_str] = hyps
+ else:
+ ans = None
+ return ans
+
+
+def decode_dataset(
+ dl: torch.utils.data.DataLoader,
+ params: AttributeDict,
+ model: nn.Module,
+ HLG: Optional[k2.Fsa],
+ H: Optional[k2.Fsa],
+ bpe_model: Optional[spm.SentencePieceProcessor],
+ word_table: k2.SymbolTable,
+ sos_id: int,
+ eos_id: int,
+ G: Optional[k2.Fsa] = None,
+) -> Dict[str, List[Tuple[str, List[str], List[str]]]]:
+ """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. Used only when params.method is NOT ctc-decoding.
+ H:
+ The ctc topo. Used only when params.method is ctc-decoding.
+ bpe_model:
+ The BPE model. Used only when params.method is ctc-decoding.
+ word_table:
+ It is the word symbol table.
+ sos_id:
+ The token ID for SOS.
+ eos_id:
+ The token ID for EOS.
+ 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.
+ """
+ num_cuts = 0
+
+ try:
+ num_batches = len(dl)
+ except TypeError:
+ num_batches = "?"
+
+ results = defaultdict(list)
+ for batch_idx, batch in enumerate(dl):
+ texts = batch["supervisions"]["text"]
+ cut_ids = [cut.id for cut in batch["supervisions"]["cut"]]
+
+ hyps_dict = decode_one_batch(
+ params=params,
+ model=model,
+ HLG=HLG,
+ H=H,
+ bpe_model=bpe_model,
+ batch=batch,
+ word_table=word_table,
+ G=G,
+ sos_id=sos_id,
+ eos_id=eos_id,
+ )
+
+ if hyps_dict is not None:
+ for lm_scale, hyps in hyps_dict.items():
+ this_batch = []
+ assert len(hyps) == len(texts)
+ for cut_id, hyp_words, ref_text in zip(cut_ids, hyps, texts):
+ ref_words = ref_text.split()
+ this_batch.append((cut_id, ref_words, hyp_words))
+
+ results[lm_scale].extend(this_batch)
+ else:
+ assert len(results) > 0, "It should not decode to empty in the first batch!"
+ this_batch = []
+ hyp_words = []
+ for ref_text in texts:
+ ref_words = ref_text.split()
+ this_batch.append((ref_words, hyp_words))
+
+ for lm_scale in results.keys():
+ results[lm_scale].extend(this_batch)
+
+ num_cuts += len(texts)
+
+ if batch_idx % 100 == 0:
+ batch_str = f"{batch_idx}/{num_batches}"
+
+ logging.info(f"batch {batch_str}, cuts processed until now is {num_cuts}")
+ return results
+
+
+def save_results(
+ params: AttributeDict,
+ test_set_name: str,
+ results_dict: Dict[str, List[Tuple[str, List[str], List[str]]]],
+) -> None:
+ if params.method == "attention-decoder":
+ # Set it to False since there are too many logs.
+ enable_log = False
+ else:
+ enable_log = True
+ test_set_wers = dict()
+ for key, results in results_dict.items():
+ recog_path = params.result_dir / f"recogs-{test_set_name}-{key}.txt"
+ results = sorted(results)
+ store_transcripts(filename=recog_path, texts=results)
+ if enable_log:
+ 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.result_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, enable_log=enable_log
+ )
+ test_set_wers[key] = wer
+
+ if enable_log:
+ 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.result_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() -> None:
+ parser = get_parser()
+ TedLiumAsrDataModule.add_arguments(parser)
+ args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
+ args.lang_dir = Path(args.lang_dir)
+ args.lm_path = Path(args.lm_path)
+ args.result_dir = Path(args.result_dir)
+
+ args.result_dir.mkdir(exist_ok=True)
+
+ params = get_params()
+ params.update(vars(args))
+
+ setup_logger(f"{params.exp_dir}/log-{params.method}/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}")
+
+ graph_compiler = BpeCtcTrainingGraphCompiler(
+ params.lang_dir,
+ device=device,
+ sos_token="",
+ eos_token="",
+ )
+ sos_id = graph_compiler.sos_id
+ eos_id = graph_compiler.eos_id
+
+ if params.method in ("ctc-decoding", "ctc-greedy-search"):
+ HLG = None
+ H = k2.ctc_topo(
+ max_token=max_token_id,
+ modified=False,
+ device=device,
+ )
+ bpe_model = spm.SentencePieceProcessor()
+ bpe_model.load(str(params.lang_dir / "bpe.model"))
+ else:
+ H = None
+ bpe_model = None
+ HLG = k2.Fsa.from_dict(
+ torch.load(f"{params.lang_dir}/HLG.pt", map_location=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"):
+ assert params.lm_path.suffix in (".pt", ".txt")
+
+ if params.lm_path.is_file() and params.lm_path.suffix == ".pt":
+ logging.info(f"Loading pre-compiled {params.lm_path.name}")
+ d = torch.load(params.lm_path, map_location=device)
+ G = k2.Fsa.from_dict(d)
+ elif not params.lm_path.is_file() and params.lm_path.suffix == ".txt":
+ raise FileNotFoundError(f"No such language model file: '{params.lm_path}'")
+ else:
+ # here we pass only if LM filename ends with '.pt' and doesn't exist
+ # or if LM filename ends '.txt' and exists.
+ if (
+ not params.lm_path.is_file()
+ and params.lm_path.suffix == ".pt"
+ and not (
+ params.lm_path.parent / f"{params.lm_path.stem}.fst.txt"
+ ).is_file()
+ ):
+ raise FileNotFoundError(
+ f"No such language model file: '{params.lm_path}'\n"
+ "'.fst.txt' representation of the language model was "
+ "not found either."
+ )
+ else:
+ # whatever params.lm_path.name we got lm_name.pt or lm_name.fst.txt
+ # we are going to load lm_name.fst.txt here
+ params.lm_path = params.lm_path.parent / params.lm_path.name.replace(
+ ".pt", ".fst.txt"
+ )
+ logging.info(f"Loading {params.lm_path.name}")
+ logging.warning("It may take 8 minutes.")
+ with open(params.lm_path) 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
+ # See https://github.com/k2-fsa/k2/issues/874
+ # for why we need to set G.properties to None
+ G.__dict__["_properties"] = None
+ G = k2.Fsa.from_fsas([G]).to(device)
+ G = k2.arc_sort(G)
+ # Save a dummy value so that it can be loaded in C++.
+ # See https://github.com/pytorch/pytorch/issues/67902
+ # for why we need to do this.
+ G.dummy = 1
+
+ torch.save(
+ G.as_dict(),
+ params.lm_path.parent
+ / params.lm_path.name.replace(".fst.txt", ".pt"),
+ )
+
+ 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 = Conformer(
+ num_features=params.feature_dim,
+ num_classes=num_classes,
+ subsampling_factor=params.subsampling_factor,
+ d_model=params.dim_model,
+ nhead=params.nhead,
+ dim_feedforward=params.dim_feedforward,
+ num_encoder_layers=params.num_encoder_layers,
+ num_decoder_layers=params.num_decoder_layers,
+ )
+
+ if not params.use_averaged_model:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ elif 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 i >= 1:
+ filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg + 1
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg + 1:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ filename_start = filenames[-1]
+ filename_end = filenames[0]
+ logging.info(
+ "Calculating the averaged model over iteration checkpoints"
+ f" from {filename_start} (excluded) to {filename_end}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+ else:
+ assert params.avg > 0, params.avg
+ start = params.epoch - params.avg
+ assert start >= 1, start
+ filename_start = f"{params.exp_dir}/epoch-{start}.pt"
+ filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
+ logging.info(
+ f"Calculating the averaged model over epoch range from "
+ f"{start} (excluded) to {params.epoch}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+
+ model.to(device)
+ model.eval()
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ # we need cut ids to display recognition results.
+ args.return_cuts = True
+ tedlium = TedLiumAsrDataModule(args)
+
+ valid_cuts = tedlium.dev_cuts()
+ test_cuts = tedlium.test_cuts()
+
+ valid_dl = tedlium.valid_dataloaders(valid_cuts)
+ test_dl = tedlium.test_dataloaders(test_cuts)
+
+ test_sets = ["dev", "test"]
+ test_dls = [valid_dl, test_dl]
+
+ for test_set, test_dl in zip(test_sets, test_dls):
+ results_dict = decode_dataset(
+ dl=test_dl,
+ params=params,
+ model=model,
+ HLG=HLG,
+ H=H,
+ bpe_model=bpe_model,
+ word_table=lexicon.word_table,
+ G=G,
+ sos_id=sos_id,
+ eos_id=eos_id,
+ )
+
+ save_results(params=params, test_set_name=test_set, results_dict=results_dict)
+
+ logging.info("Done!")
+
+
+torch.set_num_threads(1)
+# when we import add_model_arguments from train.py
+# we enforce torch.set_num_interop_threads(1) in it,
+# so we ended up with setting num_interop_threads to one
+# two times: in train.py and decode.py which cause an error,
+# that is why added an additional if statement.
+if torch.get_num_interop_threads() != 1:
+ torch.set_num_interop_threads(1)
+
+# The flag below controls whether to allow TF32 on matmul. This flag defaults to False
+# in PyTorch 1.12 and later.
+torch.backends.cuda.matmul.allow_tf32 = True
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/tedlium3/ASR/conformer_ctc2/export.py b/egs/tedlium3/ASR/conformer_ctc2/export.py
new file mode 100755
index 000000000..009bea230
--- /dev/null
+++ b/egs/tedlium3/ASR/conformer_ctc2/export.py
@@ -0,0 +1,294 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022 Behavox LLC (Author: Daniil Kulko)
+#
+# 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.
+
+# This script converts several saved checkpoints
+# to a single one using model averaging.
+"""
+Usage:
+./conformer_ctc2/export.py \
+ --exp-dir ./conformer_ctc2/exp \
+ --epoch 20 \
+ --avg 10
+
+It will generate a file exp_dir/pretrained.pt
+
+To use the generated file with `conformer_ctc2/decode.py`,
+you can do:
+
+ cd /path/to/exp_dir
+ ln -s pretrained.pt epoch-9999.pt
+
+ cd /path/to/egs/tedlium3/ASR
+ ./conformer_ctc2/decode.py \
+ --exp-dir ./conformer_ctc2/exp \
+ --epoch 9999 \
+ --avg 1 \
+ --max-duration 100
+"""
+
+import argparse
+import logging
+from pathlib import Path
+
+import torch
+from conformer import Conformer
+from scaling_converter import convert_scaled_to_non_scaled
+from train import add_model_arguments
+
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from icefall.lexicon import Lexicon
+from icefall.utils import AttributeDict, str2bool
+
+
+def get_parser() -> argparse.ArgumentParser:
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=30,
+ help="""It specifies the checkpoint to use for averaging.
+ Note: Epoch counts from 0.
+ You can specify --avg to use more checkpoints for model averaging.""",
+ )
+
+ parser.add_argument(
+ "--iter",
+ type=int,
+ default=0,
+ help="""If positive, --epoch is ignored and it
+ will use the checkpoint exp_dir/checkpoint-iter.pt.
+ You can specify --avg to use more checkpoints for model averaging.
+ """,
+ )
+
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=15,
+ help=(
+ "Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch' and '--iter'"
+ ),
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=True,
+ help=(
+ "Whether to load averaged model. Currently it only supports "
+ "using --epoch. If True, it would decode with the averaged model "
+ "over the epoch range from `epoch-avg` (excluded) to `epoch`."
+ "Actually only the models with epoch number of `epoch-avg` and "
+ "`epoch` are loaded for averaging. "
+ ),
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="conformer_ctc2/exp",
+ help="""It specifies the directory where all training related
+ files, e.g., checkpoints, log, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=str,
+ default="data/lang_bpe_500",
+ help="The lang dir",
+ )
+
+ parser.add_argument(
+ "--jit",
+ type=str2bool,
+ default=True,
+ help="""True to save a model after applying torch.jit.script.
+ """,
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def get_params() -> AttributeDict:
+ """Return a dict containing training parameters.
+
+ All training related parameters that are not passed from the commandline
+ are 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`:
+
+ - feature_dim: The model input dim. It has to match the one used
+ in computing features.
+
+ - subsampling_factor: The subsampling factor for the model.
+ """
+ # parameters for conformer
+ params = AttributeDict({"subsampling_factor": 4, "feature_dim": 80})
+ return params
+
+
+def main():
+ args = get_parser().parse_args()
+ args.exp_dir = Path(args.exp_dir)
+ args.lang_dir = Path(args.lang_dir)
+
+ params = get_params()
+ params.update(vars(args))
+
+ 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}")
+
+ logging.info(params)
+
+ logging.info("About to create model")
+
+ model = Conformer(
+ num_features=params.feature_dim,
+ num_classes=num_classes,
+ subsampling_factor=params.subsampling_factor,
+ d_model=params.dim_model,
+ nhead=params.nhead,
+ dim_feedforward=params.dim_feedforward,
+ num_encoder_layers=params.num_encoder_layers,
+ num_decoder_layers=params.num_decoder_layers,
+ )
+
+ model.to(device)
+
+ if not params.use_averaged_model:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ elif 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 i >= 1:
+ filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg + 1
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg + 1:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ filename_start = filenames[-1]
+ filename_end = filenames[0]
+ logging.info(
+ "Calculating the averaged model over iteration checkpoints"
+ f" from {filename_start} (excluded) to {filename_end}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+ else:
+ assert params.avg > 0, params.avg
+ start = params.epoch - params.avg
+ assert start >= 1, start
+ filename_start = f"{params.exp_dir}/epoch-{start}.pt"
+ filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
+ logging.info(
+ "Calculating the averaged model over epoch range from "
+ f"{start} (excluded) to {params.epoch}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+
+ model.to("cpu")
+ model.eval()
+
+ if params.jit:
+ convert_scaled_to_non_scaled(model, inplace=True)
+ logging.info("Using torch.jit.script")
+ model = torch.jit.script(model)
+ filename = params.exp_dir / "cpu_jit.pt"
+ model.save(str(filename))
+ logging.info(f"Saved to {filename}")
+ else:
+ logging.info("Not using torch.jit.script")
+ # Save it using a format so that it can be loaded
+ # by :func:`load_checkpoint`
+ filename = params.exp_dir / "pretrained.pt"
+ torch.save({"model": model.state_dict()}, str(filename))
+ logging.info(f"Saved to {filename}")
+
+
+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/tedlium3/ASR/conformer_ctc2/label_smoothing.py b/egs/tedlium3/ASR/conformer_ctc2/label_smoothing.py
new file mode 120000
index 000000000..e9d239fff
--- /dev/null
+++ b/egs/tedlium3/ASR/conformer_ctc2/label_smoothing.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/conformer_ctc/label_smoothing.py
\ No newline at end of file
diff --git a/egs/tedlium3/ASR/conformer_ctc2/lstmp.py b/egs/tedlium3/ASR/conformer_ctc2/lstmp.py
new file mode 120000
index 000000000..b82e115fc
--- /dev/null
+++ b/egs/tedlium3/ASR/conformer_ctc2/lstmp.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/lstm_transducer_stateless2/lstmp.py
\ No newline at end of file
diff --git a/egs/tedlium3/ASR/conformer_ctc2/optim.py b/egs/tedlium3/ASR/conformer_ctc2/optim.py
new file mode 120000
index 000000000..0a2f285aa
--- /dev/null
+++ b/egs/tedlium3/ASR/conformer_ctc2/optim.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless2/optim.py
\ No newline at end of file
diff --git a/egs/tedlium3/ASR/conformer_ctc2/scaling.py b/egs/tedlium3/ASR/conformer_ctc2/scaling.py
new file mode 120000
index 000000000..c10cdfe12
--- /dev/null
+++ b/egs/tedlium3/ASR/conformer_ctc2/scaling.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless2/scaling.py
\ No newline at end of file
diff --git a/egs/tedlium3/ASR/conformer_ctc2/scaling_converter.py b/egs/tedlium3/ASR/conformer_ctc2/scaling_converter.py
new file mode 120000
index 000000000..db93d155b
--- /dev/null
+++ b/egs/tedlium3/ASR/conformer_ctc2/scaling_converter.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless3/scaling_converter.py
\ No newline at end of file
diff --git a/egs/tedlium3/ASR/conformer_ctc2/subsampling.py b/egs/tedlium3/ASR/conformer_ctc2/subsampling.py
new file mode 120000
index 000000000..8c91f2336
--- /dev/null
+++ b/egs/tedlium3/ASR/conformer_ctc2/subsampling.py
@@ -0,0 +1 @@
+../../../librispeech/ASR/conformer_ctc2/subsampling.py
\ No newline at end of file
diff --git a/egs/tedlium3/ASR/conformer_ctc2/train.py b/egs/tedlium3/ASR/conformer_ctc2/train.py
new file mode 100755
index 000000000..42e4c010a
--- /dev/null
+++ b/egs/tedlium3/ASR/conformer_ctc2/train.py
@@ -0,0 +1,1061 @@
+#!/usr/bin/env python3
+# Copyright 2022 Behavox LLC. (authors: Daniil Kulko)
+#
+# 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.
+"""
+Usage:
+
+export CUDA_VISIBLE_DEVICES="0,1,2,3"
+
+./conformer_ctc/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --exp-dir conformer_ctc/exp \
+ --max-duration 300
+
+# For mix precision training:
+
+./conformer_ctc/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --use-fp16 1 \
+ --exp-dir conformer_ctc/exp \
+ --max-duration 550
+
+"""
+
+
+import argparse
+import copy
+import logging
+from pathlib import Path
+from shutil import copyfile
+from typing import Any, Dict, Optional, Tuple, Union
+
+import k2
+import optim
+import sentencepiece as spm
+import torch
+import torch.multiprocessing as mp
+from asr_datamodule import TedLiumAsrDataModule
+from conformer import Conformer
+from lhotse.dataset.sampling.base import CutSampler
+from lhotse.utils import fix_random_seed
+from local.convert_transcript_words_to_bpe_ids import convert_texts_into_ids
+from torch import Tensor
+from torch.cuda.amp import GradScaler
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.tensorboard import SummaryWriter
+
+from icefall import diagnostics
+from icefall.bpe_graph_compiler import BpeCtcTrainingGraphCompiler
+from icefall.checkpoint import load_checkpoint, remove_checkpoints
+from icefall.checkpoint import save_checkpoint as save_checkpoint_impl
+from icefall.checkpoint import (
+ save_checkpoint_with_global_batch_idx,
+ update_averaged_model,
+)
+from icefall.dist import cleanup_dist, setup_dist
+from icefall.env import get_env_info
+from icefall.lexicon import Lexicon
+from icefall.utils import (
+ AttributeDict,
+ MetricsTracker,
+ display_and_save_batch,
+ encode_supervisions,
+ setup_logger,
+ str2bool,
+)
+
+LRSchedulerType = Union[torch.optim.lr_scheduler._LRScheduler, optim.LRScheduler]
+
+
+def add_model_arguments(parser: argparse.ArgumentParser) -> None:
+ parser.add_argument(
+ "--num-encoder-layers",
+ type=int,
+ default=24,
+ help="Number of conformer encoder layers..",
+ )
+
+ parser.add_argument(
+ "--num-decoder-layers",
+ type=int,
+ default=6,
+ help="""Number of decoder layer of transformer decoder.
+ Setting this to 0 will not create the decoder at all (pure CTC model)
+ """,
+ )
+
+ parser.add_argument(
+ "--att-rate",
+ type=float,
+ default=0.8,
+ help="""The attention rate.
+ The total loss is (1 - att_rate) * ctc_loss + att_rate * att_loss
+ """,
+ )
+
+ parser.add_argument(
+ "--dim-feedforward",
+ type=int,
+ default=1536,
+ help="Feedforward module dimension of the conformer model.",
+ )
+
+ parser.add_argument(
+ "--nhead",
+ type=int,
+ default=8,
+ help="Number of attention heads in the conformer multiheadattention modules.",
+ )
+
+ parser.add_argument(
+ "--dim-model",
+ type=int,
+ default=384,
+ help="Attention dimension in the conformer model.",
+ )
+
+
+def get_parser() -> argparse.ArgumentParser:
+ 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.",
+ )
+
+ parser.add_argument(
+ "--num-epochs",
+ type=int,
+ default=30,
+ help="Number of epochs to train.",
+ )
+
+ parser.add_argument(
+ "--start-epoch",
+ type=int,
+ default=1,
+ help="""Resume training from this epoch. It should be positive.
+ If larger than 1, it will load checkpoint from
+ exp-dir/epoch-{start_epoch-1}.pt
+ """,
+ )
+
+ parser.add_argument(
+ "--start-batch",
+ type=int,
+ default=0,
+ help="""If positive, --start-epoch is ignored and
+ it loads the checkpoint from exp-dir/checkpoint-{start_batch}.pt
+ """,
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="conformer_ctc/exp",
+ help="""The experiment dir.
+ It specifies the directory where all training related
+ files, e.g., checkpoints, log, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=str,
+ default="data/lang_bpe_500",
+ help="""The lang dir
+ It contains language related input files such as
+ "lexicon.txt" and "bpe.model"
+ """,
+ )
+
+ parser.add_argument(
+ "--initial-lr",
+ type=float,
+ default=0.003,
+ help="The initial learning rate. This value should not need to be changed.",
+ )
+
+ parser.add_argument(
+ "--lr-batches",
+ type=float,
+ default=5000,
+ help="""Number of steps that affects how rapidly the learning rate
+ decreases. We suggest not to change this.""",
+ )
+
+ parser.add_argument(
+ "--lr-epochs",
+ type=float,
+ default=6,
+ help="Number of epochs that affects how rapidly the learning rate decreases.",
+ )
+
+ parser.add_argument(
+ "--seed",
+ type=int,
+ default=42,
+ help="The seed for random generators intended for reproducibility",
+ )
+
+ parser.add_argument(
+ "--print-diagnostics",
+ type=str2bool,
+ default=False,
+ help="Accumulate stats on activations, print them and exit.",
+ )
+
+ parser.add_argument(
+ "--save-every-n",
+ type=int,
+ default=4000,
+ help="""Save checkpoint after processing this number of batches"
+ periodically. We save checkpoint to exp-dir/ whenever
+ params.batch_idx_train % save_every_n == 0. The checkpoint filename
+ has the form: f'exp-dir/checkpoint-{params.batch_idx_train}.pt'
+ Note: It also saves checkpoint to `exp-dir/epoch-xxx.pt` at the
+ end of each epoch where `xxx` is the epoch number counting from 0.
+ """,
+ )
+
+ parser.add_argument(
+ "--keep-last-k",
+ type=int,
+ default=30,
+ help="""Only keep this number of checkpoints on disk.
+ For instance, if it is 3, there are only 3 checkpoints
+ in the exp-dir with filenames `checkpoint-xxx.pt`.
+ It does not affect checkpoints with name `epoch-xxx.pt`.
+ """,
+ )
+
+ parser.add_argument(
+ "--average-period",
+ type=int,
+ default=100,
+ help="""Update the averaged model, namely `model_avg`, after processing
+ this number of batches. `model_avg` is a separate version of model,
+ in which each floating-point parameter is the average of all the
+ parameters from the start of training. Each time we take the average,
+ we do: `model_avg = model * (average_period / batch_idx_train) +
+ model_avg * ((batch_idx_train - average_period) / batch_idx_train)`.
+ """,
+ )
+
+ parser.add_argument(
+ "--use-fp16",
+ type=str2bool,
+ default=False,
+ help="Whether to use half precision training.",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def get_params() -> AttributeDict:
+ """Return a dict containing training parameters.
+
+ All training related parameters that are not passed from the commandline
+ are 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`:
+
+ - 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
+
+ - reset_interval: Reset statistics if batch_idx % reset_interval is 0
+
+ - valid_interval: Run validation if batch_idx % valid_interval is 0
+
+ - feature_dim: The model input dim. It has to match the one used
+ in computing features.
+
+ - subsampling_factor: The subsampling factor for the model.
+
+ - warm_step: The warm_step for Noam optimizer.
+ """
+ params = AttributeDict(
+ {
+ "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,
+ "reset_interval": 200,
+ "valid_interval": 1000,
+ # parameters for conformer
+ "feature_dim": 80,
+ "subsampling_factor": 4,
+ # parameters for ctc loss
+ "beam_size": 10,
+ "reduction": "none",
+ "use_double_scores": True,
+ # parameters for Noam
+ "model_warm_step": 3000, # arg given to model, not for lrate
+ "env_info": get_env_info(),
+ }
+ )
+
+ return params
+
+
+def load_checkpoint_if_available(
+ params: AttributeDict,
+ model: torch.nn.Module,
+ model_avg: torch.nn.Module = None,
+ optimizer: Optional[torch.optim.Optimizer] = None,
+ scheduler: Optional[LRSchedulerType] = None,
+) -> Optional[Dict[str, Any]]:
+ """Load checkpoint from file.
+
+ If params.start_batch is positive, it will load the checkpoint from
+ `params.exp_dir/checkpoint-{params.start_batch}.pt`. Otherwise, if
+ params.start_epoch is larger than 1, it will load the checkpoint from
+ `params.start_epoch - 1`.
+
+ Apart from loading state dict for `model` and `optimizer` 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.
+ model_avg:
+ The stored model averaged from the start of training.
+ optimizer:
+ The optimizer that we are using.
+ scheduler:
+ The scheduler that is used for training.
+ Returns:
+ Return a dict containing previously saved training info.
+ """
+ if params.start_batch > 0:
+ filename = params.exp_dir / f"checkpoint-{params.start_batch}.pt"
+ elif params.start_epoch > 1:
+ filename = params.exp_dir / f"epoch-{params.start_epoch-1}.pt"
+ else:
+ return None
+
+ assert filename.is_file(), f"{filename} does not exist!"
+
+ saved_params = load_checkpoint(
+ filename,
+ model=model,
+ model_avg=model_avg,
+ 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]
+
+ if params.start_batch > 0:
+ if "cur_epoch" in saved_params:
+ params["start_epoch"] = saved_params["cur_epoch"]
+
+ return saved_params
+
+
+def save_checkpoint(
+ params: AttributeDict,
+ model: Union[torch.nn.Module, DDP],
+ model_avg: Optional[torch.nn.Module] = None,
+ optimizer: Optional[torch.optim.Optimizer] = None,
+ scheduler: Optional[LRSchedulerType] = None,
+ sampler: Optional[CutSampler] = None,
+ scaler: Optional[GradScaler] = 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.
+ model_avg:
+ The stored model averaged from the start of training.
+ optimizer:
+ The optimizer used for training.
+ scheduler:
+ The learning rate scheduler used for training.
+ sampler:
+ The sampler for the training dataset.
+ scaler:
+ The scaler used for mix precision training.
+ """
+ if rank != 0:
+ return
+ filename = params.exp_dir / f"epoch-{params.cur_epoch}.pt"
+ save_checkpoint_impl(
+ filename=filename,
+ model=model,
+ model_avg=model_avg,
+ params=params,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=sampler,
+ scaler=scaler,
+ 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: Union[torch.nn.Module, DDP],
+ graph_compiler: BpeCtcTrainingGraphCompiler,
+ batch: dict,
+ is_training: bool,
+ warmup: float = 1.0,
+) -> Tuple[Tensor, MetricsTracker]:
+ """
+ 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.
+ warmup: a floating point value which increases throughout training;
+ values >= 1.0 are fully warmed up and have all modules present.
+ """
+ device = model.device if isinstance(model, DDP) else next(model.parameters()).device
+ feature = batch["inputs"]
+ # at entry, feature is (N, T, C)
+ assert feature.ndim == 3
+ feature = feature.to(device)
+
+ supervisions = batch["supervisions"]
+ feature_lens = supervisions["num_frames"].to(device)
+
+ with torch.set_grad_enabled(is_training):
+ nnet_output, encoder_memory, memory_mask = model(
+ feature, supervisions, warmup=warmup
+ )
+
+ supervision_segments, texts = encode_supervisions(
+ supervisions, subsampling_factor=params.subsampling_factor
+ )
+
+ token_ids = convert_texts_into_ids(texts, graph_compiler.sp)
+ 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):
+ mmodel = model.module if hasattr(model, "module") else model
+ # Note: We need to generate an unsorted version of token_ids
+ # `encode_supervisions()` called above sorts text, but
+ # encoder_memory and memory_mask are not sorted, so we
+ # use an unsorted version `supervisions["text"]` to regenerate
+ # the token_ids
+ #
+ # See https://github.com/k2-fsa/icefall/issues/97
+ # for more details
+ unsorted_token_ids = graph_compiler.texts_to_ids(supervisions["text"])
+ att_loss = mmodel.decoder_forward(
+ encoder_memory,
+ memory_mask,
+ token_ids=unsorted_token_ids,
+ sos_id=graph_compiler.sos_id,
+ eos_id=graph_compiler.eos_id,
+ warmup=warmup,
+ )
+ else:
+ att_loss = torch.tensor([0])
+
+ ctc_loss_is_finite = torch.isfinite(ctc_loss)
+ att_loss_is_finite = torch.isfinite(att_loss)
+ if torch.any(~ctc_loss_is_finite) or torch.any(~att_loss_is_finite):
+ logging.info(
+ "Not all losses are finite!\n"
+ f"ctc_loss: {ctc_loss}\n"
+ f"att_loss: {att_loss}"
+ )
+ display_and_save_batch(batch, params=params, sp=graph_compiler.sp)
+ ctc_loss = ctc_loss[ctc_loss_is_finite]
+ att_loss = att_loss[att_loss_is_finite]
+
+ # If the batch contains more than 10 utterances AND
+ # if either all ctc_loss or att_loss is inf or nan,
+ # we stop the training process by raising an exception
+ if torch.all(~ctc_loss_is_finite) or torch.all(~att_loss_is_finite):
+ raise ValueError(
+ "There are too many utterances in this batch "
+ "leading to inf or nan losses."
+ )
+
+ ctc_loss = ctc_loss.sum()
+ att_loss = att_loss.sum()
+
+ if params.att_rate > 0.0:
+ loss = (1.0 - params.att_rate) * ctc_loss + params.att_rate * att_loss
+ else:
+ loss = ctc_loss
+
+ assert loss.requires_grad == is_training
+
+ info = MetricsTracker()
+ # info["frames"] is an approximate number for two reasons:
+ # (1) The acutal subsampling factor is ((lens - 1) // 2 - 1) // 2
+ # (2) If some utterances in the batch lead to inf/nan loss, they
+ # are filtered out.
+ info["frames"] = (
+ torch.div(feature_lens, params.subsampling_factor, rounding_mode="floor")
+ .sum()
+ .item()
+ )
+
+ # `utt_duration` and `utt_pad_proportion` would be normalized by `utterances` # noqa
+ info["utterances"] = feature.size(0)
+ # averaged input duration in frames over utterances
+ info["utt_duration"] = feature_lens.sum().item()
+ # averaged padding proportion over utterances
+ info["utt_pad_proportion"] = (
+ ((feature.size(1) - feature_lens) / feature.size(1)).sum().item()
+ )
+
+ # Note: We use reduction=sum while computing the loss.
+ info["loss"] = loss.detach().cpu().item()
+ info["ctc_loss"] = ctc_loss.detach().cpu().item()
+ if params.att_rate > 0.0:
+ info["att_loss"] = att_loss.detach().cpu().item()
+
+ return loss, info
+
+
+def compute_validation_loss(
+ params: AttributeDict,
+ model: Union[torch.nn.Module, DDP],
+ graph_compiler: BpeCtcTrainingGraphCompiler,
+ valid_dl: torch.utils.data.DataLoader,
+ world_size: int = 1,
+) -> MetricsTracker:
+ """Run the validation process."""
+ model.eval()
+
+ tot_loss = MetricsTracker()
+
+ for batch in valid_dl:
+ loss, loss_info = compute_loss(
+ params=params,
+ model=model,
+ graph_compiler=graph_compiler,
+ batch=batch,
+ is_training=False,
+ )
+ assert loss.requires_grad is False
+ tot_loss = tot_loss + loss_info
+
+ if world_size > 1:
+ tot_loss.reduce(loss.device)
+
+ loss_value = tot_loss["loss"] / tot_loss["frames"]
+ if loss_value < params.best_valid_loss:
+ params.best_valid_epoch = params.cur_epoch
+ params.best_valid_loss = loss_value
+
+ return tot_loss
+
+
+def train_one_epoch(
+ params: AttributeDict,
+ model: Union[torch.nn.Module, DDP],
+ optimizer: torch.optim.Optimizer,
+ scheduler: LRSchedulerType,
+ graph_compiler: BpeCtcTrainingGraphCompiler,
+ train_dl: torch.utils.data.DataLoader,
+ valid_dl: torch.utils.data.DataLoader,
+ scaler: GradScaler,
+ model_avg: Optional[torch.nn.Module] = None,
+ tb_writer: Optional[SummaryWriter] = None,
+ world_size: int = 1,
+ rank: int = 0,
+) -> 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.
+ scheduler:
+ The learning rate scheduler, we call step() every step.
+ graph_compiler:
+ It is used to convert transcripts to FSAs.
+ train_dl:
+ Dataloader for the training dataset.
+ valid_dl:
+ Dataloader for the validation dataset.
+ scaler:
+ The scaler used for mix precision training.
+ model_avg:
+ The stored model averaged from the start of training.
+ tb_writer:
+ Writer to write log messages to tensorboard.
+ world_size:
+ Number of nodes in DDP training. If it is 1, DDP is disabled.
+ rank:
+ The rank of the node in DDP training. If no DDP is used, it should
+ be set to 0.
+ """
+ model.train()
+
+ tot_loss = MetricsTracker()
+
+ for batch_idx, batch in enumerate(train_dl):
+ params.batch_idx_train += 1
+ batch_size = len(batch["supervisions"]["text"])
+
+ try:
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, loss_info = compute_loss(
+ params=params,
+ model=model,
+ graph_compiler=graph_compiler,
+ batch=batch,
+ is_training=True,
+ warmup=(params.batch_idx_train / params.model_warm_step),
+ )
+ # summary stats
+ tot_loss = (tot_loss * (1 - 1 / params.reset_interval)) + loss_info
+
+ # NOTE: We use reduction==sum and loss is computed over utterances
+ # in the batch and there is no normalization to it so far.
+ scaler.scale(loss).backward()
+ scheduler.step_batch(params.batch_idx_train)
+ scaler.step(optimizer)
+ scaler.update()
+ optimizer.zero_grad()
+ except: # noqa
+ display_and_save_batch(batch, params=params, sp=graph_compiler.sp)
+ raise
+
+ if params.print_diagnostics and batch_idx == 5:
+ return
+
+ if (
+ rank == 0
+ and params.batch_idx_train > 0
+ and params.batch_idx_train % params.average_period == 0
+ ):
+ update_averaged_model(
+ params=params,
+ model_cur=model,
+ model_avg=model_avg,
+ )
+
+ if (
+ params.batch_idx_train > 0
+ and params.batch_idx_train % params.save_every_n == 0
+ ):
+ save_checkpoint_with_global_batch_idx(
+ out_dir=params.exp_dir,
+ global_batch_idx=params.batch_idx_train,
+ model=model,
+ model_avg=model_avg,
+ params=params,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=train_dl.sampler,
+ scaler=scaler,
+ rank=rank,
+ )
+ remove_checkpoints(
+ out_dir=params.exp_dir,
+ topk=params.keep_last_k,
+ rank=rank,
+ )
+
+ if batch_idx % params.log_interval == 0:
+ cur_lr = scheduler.get_last_lr()[0]
+ logging.info(
+ f"Epoch {params.cur_epoch}, "
+ f"batch {batch_idx}, loss[{loss_info}], "
+ f"tot_loss[{tot_loss}], batch size: {batch_size}, "
+ f"lr: {cur_lr:.2e}"
+ )
+
+ if tb_writer is not None:
+ tb_writer.add_scalar(
+ "train/learning_rate", cur_lr, params.batch_idx_train
+ )
+
+ loss_info.write_summary(
+ tb_writer, "train/current_", params.batch_idx_train
+ )
+ tot_loss.write_summary(tb_writer, "train/tot_", params.batch_idx_train)
+
+ if batch_idx > 0 and batch_idx % params.valid_interval == 0:
+ logging.info("Computing validation loss")
+ valid_info = 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}, validation: {valid_info}")
+ if tb_writer is not None:
+ valid_info.write_summary(
+ tb_writer, "train/valid_", params.batch_idx_train
+ )
+
+ loss_value = tot_loss["loss"] / tot_loss["frames"]
+ params.train_loss = loss_value
+ 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(params.seed)
+ 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)
+ logging.info(f"Device: {device}")
+
+ if "lang_bpe" not in str(params.lang_dir):
+ raise ValueError(
+ f"Unsupported type of lang dir (we expected it to have "
+ f"'lang_bpe' in its name): {params.lang_dir}"
+ )
+
+ 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,
+ num_classes=num_classes,
+ subsampling_factor=params.subsampling_factor,
+ d_model=params.dim_model,
+ nhead=params.nhead,
+ dim_feedforward=params.dim_feedforward,
+ num_encoder_layers=params.num_encoder_layers,
+ num_decoder_layers=params.num_decoder_layers,
+ )
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ assert params.save_every_n >= params.average_period
+ model_avg: Optional[torch.nn.Module] = None
+ if rank == 0:
+ # model_avg is only used with rank 0
+ model_avg = copy.deepcopy(model)
+
+ assert params.start_epoch > 0, params.start_epoch
+ checkpoints = load_checkpoint_if_available(
+ params=params, model=model, model_avg=model_avg
+ )
+
+ model.to(device)
+ if world_size > 1:
+ logging.info("Using DDP")
+ model = DDP(model, device_ids=[rank])
+
+ optimizer = optim.Eve(model.parameters(), lr=params.initial_lr)
+ scheduler = optim.Eden(optimizer, params.lr_batches, params.lr_epochs)
+
+ if checkpoints and checkpoints.get("optimizer") is not None:
+ logging.info("Loading optimizer state dict")
+ optimizer.load_state_dict(checkpoints["optimizer"])
+
+ if checkpoints and checkpoints.get("scheduler") is not None:
+ logging.info("Loading scheduler state dict")
+ scheduler.load_state_dict(checkpoints["scheduler"])
+
+ if params.print_diagnostics:
+ opts = diagnostics.TensorDiagnosticOptions(
+ 2**22
+ ) # allow 4 megabytes per sub-module
+ diagnostic = diagnostics.attach_diagnostics(model, opts)
+
+ tedlium = TedLiumAsrDataModule(args)
+
+ train_cuts = tedlium.train_cuts()
+
+ if params.start_batch > 0 and checkpoints and "sampler" in checkpoints:
+ # We only load the sampler's state dict when it loads a checkpoint
+ # saved in the middle of an epoch
+ sampler_state_dict = checkpoints["sampler"]
+ else:
+ sampler_state_dict = None
+
+ train_dl = tedlium.train_dataloaders(
+ train_cuts, sampler_state_dict=sampler_state_dict
+ )
+
+ valid_cuts = tedlium.dev_cuts()
+ valid_dl = tedlium.valid_dataloaders(valid_cuts)
+
+ if (
+ params.start_epoch <= 1
+ and params.start_batch <= 0
+ and not params.print_diagnostics
+ ):
+ scan_pessimistic_batches_for_oom(
+ model=model,
+ train_dl=train_dl,
+ optimizer=optimizer,
+ graph_compiler=graph_compiler,
+ params=params,
+ warmup=0.0 if params.start_epoch == 1 else 1.0,
+ )
+
+ scaler = GradScaler(enabled=params.use_fp16)
+ if checkpoints and "grad_scaler" in checkpoints:
+ logging.info("Loading grad scaler state dict")
+ scaler.load_state_dict(checkpoints["grad_scaler"])
+
+ for epoch in range(params.start_epoch, params.num_epochs + 1):
+ scheduler.step_epoch(epoch - 1)
+ fix_random_seed(params.seed + epoch - 1)
+ train_dl.sampler.set_epoch(epoch - 1)
+ train_dl.dataset.epoch = epoch - 1
+
+ if tb_writer is not None:
+ tb_writer.add_scalar("train/epoch", epoch, params.batch_idx_train)
+
+ params.cur_epoch = epoch
+
+ train_one_epoch(
+ params=params,
+ model=model,
+ model_avg=model_avg,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ graph_compiler=graph_compiler,
+ train_dl=train_dl,
+ valid_dl=valid_dl,
+ scaler=scaler,
+ tb_writer=tb_writer,
+ world_size=world_size,
+ rank=rank,
+ )
+
+ if params.print_diagnostics:
+ diagnostic.print_diagnostics()
+ break
+
+ save_checkpoint(
+ params=params,
+ model=model,
+ model_avg=model_avg,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=train_dl.sampler,
+ scaler=scaler,
+ rank=rank,
+ )
+
+ logging.info("Done!")
+
+ if world_size > 1:
+ torch.distributed.barrier()
+ cleanup_dist()
+
+
+def scan_pessimistic_batches_for_oom(
+ model: Union[torch.nn.Module, DDP],
+ train_dl: torch.utils.data.DataLoader,
+ optimizer: torch.optim.Optimizer,
+ graph_compiler: BpeCtcTrainingGraphCompiler,
+ params: AttributeDict,
+ warmup: float,
+):
+ from lhotse.dataset import find_pessimistic_batches
+
+ logging.info(
+ "Sanity check -- see if any of the batches in epoch 1 would cause OOM."
+ )
+ batches, crit_values = find_pessimistic_batches(train_dl.sampler)
+ for criterion, cuts in batches.items():
+ batch = train_dl.dataset[cuts]
+ try:
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, _ = compute_loss(
+ params=params,
+ model=model,
+ graph_compiler=graph_compiler,
+ batch=batch,
+ is_training=True,
+ warmup=warmup,
+ )
+ loss.backward()
+ optimizer.step()
+ optimizer.zero_grad()
+ except Exception as e:
+ if "CUDA out of memory" in str(e):
+ logging.error(
+ "Your GPU ran out of memory with the current "
+ "max_duration setting. We recommend decreasing "
+ "max_duration and trying again.\n"
+ f"Failing criterion: {criterion} "
+ f"(={crit_values[criterion]}) ..."
+ )
+ display_and_save_batch(batch, params=params, sp=graph_compiler.sp)
+ raise
+
+
+def main():
+ parser = get_parser()
+ TedLiumAsrDataModule.add_arguments(parser)
+ args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ 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)
+
+
+torch.set_num_threads(1)
+torch.set_num_interop_threads(1)
+
+# The flag below controls whether to allow TF32 on matmul. This flag defaults to False
+# in PyTorch 1.12 and later.
+torch.backends.cuda.matmul.allow_tf32 = True
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/tedlium3/ASR/conformer_ctc2/transformer.py b/egs/tedlium3/ASR/conformer_ctc2/transformer.py
new file mode 100644
index 000000000..9dbf32e48
--- /dev/null
+++ b/egs/tedlium3/ASR/conformer_ctc2/transformer.py
@@ -0,0 +1,1093 @@
+# Copyright 2021 University of Chinese Academy of Sciences (author: Han Zhu)
+# Copyright 2022 Xiaomi Corp. (author: Quandong Wang)
+#
+# 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.
+
+import copy
+import math
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from attention import MultiheadAttention
+from combiner import RandomCombine
+from label_smoothing import LabelSmoothingLoss
+from scaling import (
+ ActivationBalancer,
+ BasicNorm,
+ DoubleSwish,
+ ScaledEmbedding,
+ ScaledLinear,
+)
+from subsampling import Conv2dSubsampling
+from torch.nn.utils.rnn import pad_sequence
+
+# Note: TorchScript requires Dict/List/etc. to be fully typed.
+Supervisions = Dict[str, torch.Tensor]
+
+
+class Transformer(nn.Module):
+ 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,
+ layer_dropout: float = 0.075,
+ aux_layer_period: int = 3,
+ ) -> None:
+ """
+ Args:
+ num_features:
+ the input dimension of the model.
+ num_classes:
+ the output dimension of the model.
+ subsampling_factor:
+ number of output frames is num_in_frames // subsampling_factor;
+ currently, subsampling_factor MUST be 4.
+ d_model:
+ attention dimension.
+ nhead:
+ number of heads in multi-head attention;
+ must satisfy d_model // nhead == 0.
+ dim_feedforward:
+ the output dimension of the feedforward layers in encoder/decoder.
+ num_encoder_layers:
+ number of encoder layers.
+ num_decoder_layers:
+ number of decoder layers.
+ dropout:
+ dropout in encoder/decoder.
+ layer_dropout:
+ layer-dropout rate.
+ aux_layer_period:
+ determines the auxiliary encoder layers.
+ """
+ super().__init__()
+
+ 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 converts the input of shape (N, T, num_classes)
+ # to the shape (N, T//subsampling_factor, d_model).
+ # That is, it does two things simultaneously:
+ # (1) subsampling: T -> T//subsampling_factor
+ # (2) embedding: num_classes -> d_model
+ self.encoder_embed = Conv2dSubsampling(num_features, d_model)
+
+ self.encoder_pos = PositionalEncoding(d_model, dropout)
+
+ encoder_layer = TransformerEncoderLayer(
+ d_model=d_model,
+ nhead=nhead,
+ dim_feedforward=dim_feedforward,
+ dropout=dropout,
+ layer_dropout=layer_dropout,
+ )
+ # aux_layers from 1/3
+ self.encoder = TransformerEncoder(
+ encoder_layer=encoder_layer,
+ num_layers=num_encoder_layers,
+ aux_layers=list(
+ range(
+ num_encoder_layers // 3,
+ num_encoder_layers - 1,
+ aux_layer_period,
+ )
+ ),
+ )
+
+ # TODO(fangjun): remove dropout
+ self.encoder_output_layer = nn.Sequential(
+ nn.Dropout(p=dropout), ScaledLinear(d_model, num_classes, bias=True)
+ )
+
+ if num_decoder_layers > 0:
+ self.decoder_num_class = (
+ self.num_classes
+ ) # bpe model already has sos/eos symbol
+
+ self.decoder_embed = ScaledEmbedding(
+ num_embeddings=self.decoder_num_class, embedding_dim=d_model
+ )
+ self.decoder_pos = PositionalEncoding(d_model, dropout)
+
+ decoder_layer = TransformerDecoderLayer(
+ d_model=d_model,
+ nhead=nhead,
+ dim_feedforward=dim_feedforward,
+ dropout=dropout,
+ )
+
+ self.decoder = TransformerDecoder(
+ decoder_layer=decoder_layer,
+ num_layers=num_decoder_layers,
+ aux_layers=[],
+ )
+
+ self.decoder_output_layer = ScaledLinear(
+ d_model, self.decoder_num_class, bias=True
+ )
+
+ self.decoder_criterion = LabelSmoothingLoss(reduction="none")
+ else:
+ self.decoder_criterion = None
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ supervision: Optional[Supervisions] = None,
+ warmup: float = 1.0,
+ ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+ """
+ Args:
+ x:
+ The input tensor. Its shape is (N, S, C).
+ supervision:
+ Supervision in lhotse format.
+ See https://github.com/lhotse-speech/lhotse/blob/master/lhotse/dataset/speech_recognition.py#L32 # noqa
+ (CAUTION: It contains length information, i.e., start and number of
+ frames, before subsampling)
+ warmup:
+ a floating point value that gradually increases from 0 throughout
+ training; when it is >= 1.0 we are "fully warmed up". It is used
+ to turn modules on sequentially.
+
+ Returns:
+ Return a tuple containing 3 tensors:
+ - CTC output for ctc decoding. Its shape is (N, S, C)
+ - Encoder output with shape (S, N, C). It can be used as key and
+ value for the decoder.
+ - Encoder output padding mask. It can be used as
+ memory_key_padding_mask for the decoder. Its shape is (N, S).
+ It is None if `supervision` is None.
+ """
+
+ encoder_memory, memory_key_padding_mask = self.run_encoder(
+ x, supervision, warmup
+ )
+
+ x = self.ctc_output(encoder_memory)
+ return x, encoder_memory, memory_key_padding_mask
+
+ def run_encoder(
+ self,
+ x: torch.Tensor,
+ supervisions: Optional[Supervisions] = None,
+ warmup: float = 1.0,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ """Run the transformer encoder.
+
+ Args:
+ x:
+ The model input. Its shape is (N, S, C).
+ supervisions:
+ Supervision in lhotse format.
+ See https://github.com/lhotse-speech/lhotse/blob/master/lhotse/dataset/speech_recognition.py#L32 # noqa
+ CAUTION: It contains length information, i.e., start and number of
+ frames, before subsampling
+ It is read directly from the batch, without any sorting. It is used
+ to compute the encoder padding mask, which is used as memory key
+ padding mask for the decoder.
+ warmup:
+ a floating point value that gradually increases from 0 throughout
+ training; when it is >= 1.0 we are "fully warmed up". It is used
+ to turn modules on sequentially.
+
+ Returns:
+ Return a tuple with two tensors:
+ - The encoder output, with shape (S, N, C)
+ - encoder padding mask, with shape (N, S).
+ The mask is None if `supervisions` is None.
+ It is used as memory key padding mask in the decoder.
+ """
+ x = self.encoder_embed(x)
+ x = self.encoder_pos(x)
+ x = x.permute(1, 0, 2) # (N, S, C) -> (S, N, C)
+ mask = encoder_padding_mask(x.size(0), supervisions)
+ mask = mask.to(x.device) if mask is not None else None
+ x = self.encoder(x, src_key_padding_mask=mask, warmup=warmup) # (S, N, C)
+
+ return x, mask
+
+ def ctc_output(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Args:
+ x:
+ the output tensor from the transformer encoder;
+ its shape is (S, N, C)
+
+ Returns:
+ Return a tensor that can be used for CTC decoding.
+ Its shape is (N, S, C)
+ """
+ x = self.encoder_output_layer(x)
+ x = x.permute(1, 0, 2) # (S, N, C) -> (N, S, C)
+ x = nn.functional.log_softmax(x, dim=-1) # (N, S, C)
+ return x
+
+ @torch.jit.export
+ def decoder_forward(
+ self,
+ memory: torch.Tensor,
+ memory_key_padding_mask: torch.Tensor,
+ token_ids: List[List[int]],
+ sos_id: int,
+ eos_id: int,
+ warmup: float = 1.0,
+ ) -> torch.Tensor:
+ """
+ Args:
+ memory:
+ It's the output of the encoder of shape (S, N, C)
+ memory_key_padding_mask:
+ The padding mask from the encoder of shape (N, S).
+ token_ids:
+ A list-of-list IDs. Each sublist contains IDs for an utterance.
+ The IDs can be either phone IDs or word piece IDs.
+ sos_id:
+ sos token id
+ eos_id:
+ eos token id
+ warmup:
+ a floating point value that gradually increases from 0 throughout
+ training; when it is >= 1.0 we are "fully warmed up". It is used
+ to turn modules on sequentially.
+
+ Returns:
+ A scalar, the **sum** of label smoothing loss over utterances
+ in the batch without any normalization.
+ """
+ ys_in = add_sos(token_ids, sos_id=sos_id)
+ ys_in = [torch.tensor(y) for y in ys_in]
+ ys_in_pad = pad_sequence(ys_in, batch_first=True, padding_value=float(eos_id))
+
+ ys_out = add_eos(token_ids, eos_id=eos_id)
+ ys_out = [torch.tensor(y) for y in ys_out]
+ ys_out_pad = pad_sequence(ys_out, batch_first=True, padding_value=float(-1))
+
+ device = memory.device
+ ys_in_pad = ys_in_pad.to(device)
+ ys_out_pad = ys_out_pad.to(device)
+
+ tgt_mask = generate_square_subsequent_mask(ys_in_pad.shape[-1]).to(device)
+
+ tgt_key_padding_mask = decoder_padding_mask(ys_in_pad, ignore_id=eos_id)
+ # TODO: Use length information to create the decoder padding mask
+ # We set the first column to False since the first column in ys_in_pad
+ # contains sos_id, which is the same as eos_id in our current setting.
+ tgt_key_padding_mask[:, 0] = False
+
+ tgt = self.decoder_embed(ys_in_pad) # (N, T) -> (N, T, C)
+ tgt = self.decoder_pos(tgt)
+ tgt = tgt.permute(1, 0, 2) # (N, T, C) -> (T, N, C)
+ pred_pad = self.decoder(
+ tgt=tgt,
+ memory=memory,
+ tgt_mask=tgt_mask,
+ tgt_key_padding_mask=tgt_key_padding_mask,
+ memory_key_padding_mask=memory_key_padding_mask,
+ warmup=warmup,
+ ) # (T, N, C)
+ pred_pad = pred_pad.permute(1, 0, 2) # (T, N, C) -> (N, T, C)
+ pred_pad = self.decoder_output_layer(pred_pad) # (N, T, C)
+
+ decoder_loss = self.decoder_criterion(pred_pad, ys_out_pad)
+
+ return decoder_loss
+
+ @torch.jit.export
+ def decoder_nll(
+ self,
+ memory: torch.Tensor,
+ memory_key_padding_mask: torch.Tensor,
+ token_ids: List[torch.Tensor],
+ sos_id: int,
+ eos_id: int,
+ warmup: float = 1.0,
+ ) -> torch.Tensor:
+ """
+ Args:
+ memory:
+ It's the output of the encoder of shape (S, N, C).
+ memory_key_padding_mask:
+ The padding mask from the encoder of shape (N, S).
+ token_ids:
+ A list-of-list IDs (e.g., word piece IDs).
+ Each sublist represents an utterance.
+ sos_id:
+ The token ID for SOS.
+ eos_id:
+ The token ID for EOS.
+ warmup:
+ a floating point value that gradually increases from 0 throughout
+ training; when it is >= 1.0 we are "fully warmed up". It is used
+ to turn modules on sequentially.
+
+ Returns:
+ A 2-D tensor of shape (len(token_ids), max_token_length)
+ representing the cross entropy loss (i.e., negative log-likelihood).
+ """
+ # The common part between this function and decoder_forward could be
+ # extracted as a separate function.
+ if isinstance(token_ids[0], torch.Tensor):
+ # This branch is executed by torchscript in C++.
+ # See https://github.com/k2-fsa/k2/pull/870
+ # https://github.com/k2-fsa/k2/blob/3c1c18400060415b141ccea0115fd4bf0ad6234e/k2/torch/bin/attention_rescore.cu#L286
+ token_ids = [tolist(t) for t in token_ids]
+
+ ys_in = add_sos(token_ids, sos_id=sos_id)
+ ys_in = [torch.tensor(y) for y in ys_in]
+ ys_in_pad = pad_sequence(ys_in, batch_first=True, padding_value=float(eos_id))
+
+ ys_out = add_eos(token_ids, eos_id=eos_id)
+ ys_out = [torch.tensor(y) for y in ys_out]
+ ys_out_pad = pad_sequence(ys_out, batch_first=True, padding_value=float(-1))
+
+ device = memory.device
+ ys_in_pad = ys_in_pad.to(device, dtype=torch.int64)
+ ys_out_pad = ys_out_pad.to(device, dtype=torch.int64)
+
+ tgt_mask = generate_square_subsequent_mask(ys_in_pad.shape[-1]).to(device)
+
+ tgt_key_padding_mask = decoder_padding_mask(ys_in_pad, ignore_id=eos_id)
+ # TODO: Use length information to create the decoder padding mask
+ # We set the first column to False since the first column in ys_in_pad
+ # contains sos_id, which is the same as eos_id in our current setting.
+ tgt_key_padding_mask[:, 0] = False
+
+ tgt = self.decoder_embed(ys_in_pad) # (N, T) -> (N, T, C)
+ tgt = self.decoder_pos(tgt)
+ tgt = tgt.permute(1, 0, 2) # (N, T, С) -> (T, N, C)
+ pred_pad = self.decoder(
+ tgt=tgt,
+ memory=memory,
+ tgt_mask=tgt_mask,
+ tgt_key_padding_mask=tgt_key_padding_mask,
+ memory_key_padding_mask=memory_key_padding_mask,
+ warmup=warmup,
+ ) # (T, B, F)
+ pred_pad = pred_pad.permute(1, 0, 2) # (T, N, C) -> (N, T, C)
+ pred_pad = self.decoder_output_layer(pred_pad) # (N, T, C)
+ # 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.
+
+ Example:
+ >>> 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,
+ bypass_scale: float = 0.1,
+ layer_dropout: float = 0.075,
+ ) -> None:
+ """
+ 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).
+ bypass_scale:
+ a scale on the layer's output, used in bypass (resnet-type) skip-connection;
+ when the layer is bypassed the final output will be a
+ weighted sum of the layer's input and layer's output with weights
+ (1.0-bypass_scale) and bypass_scale correspondingly (default=0.1).
+ layer_dropout:
+ the probability to bypass the layer (default=0.075).
+ """
+
+ super().__init__()
+
+ if bypass_scale < 0.0 or bypass_scale > 1.0:
+ raise ValueError("bypass_scale should be between 0.0 and 1.0")
+
+ if layer_dropout < 0.0 or layer_dropout > 1.0:
+ raise ValueError("layer_dropout should be between 0.0 and 1.0")
+
+ self.bypass_scale = bypass_scale
+ self.layer_dropout = layer_dropout
+
+ self.self_attn = MultiheadAttention(d_model, nhead)
+ # Implementation of Feedforward model
+
+ self.feed_forward = nn.Sequential(
+ ScaledLinear(d_model, dim_feedforward),
+ ActivationBalancer(channel_dim=-1),
+ DoubleSwish(),
+ nn.Dropout(dropout),
+ ScaledLinear(dim_feedforward, d_model, initial_scale=0.25),
+ )
+
+ self.norm_final = BasicNorm(d_model)
+
+ # try to ensure the output is close to zero-mean (or at least, zero-median).
+ self.balancer = ActivationBalancer(
+ channel_dim=-1, min_positive=0.45, max_positive=0.55, max_abs=6.0
+ )
+
+ self.dropout = nn.Dropout(dropout)
+
+ def forward(
+ self,
+ src: torch.Tensor,
+ src_mask: Optional[torch.Tensor] = None,
+ src_key_padding_mask: Optional[torch.Tensor] = None,
+ warmup: float = 1.0,
+ ) -> torch.Tensor:
+ """
+ Pass the input through the encoder layer.
+
+ Args:
+ src:
+ the sequence to the encoder layer of shape (S, N, C) (required).
+ src_mask:
+ the mask for the src sequence of shape (S, S) (optional).
+ src_key_padding_mask:
+ the mask for the src keys per batch of shape (N, S) (optional)
+ warmup:
+ controls selective bypass of layers; if < 1.0, we will
+ bypass the layer more frequently (default=1.0).
+
+ Returns:
+ Output tensor of the shape (S, N, C), where
+ S is the source sequence length,
+ N is the batch size,
+ C is the feature number.
+
+ """
+ src_orig = src
+
+ warmup_scale = min(self.bypass_scale + warmup, 1.0)
+ # alpha = 1.0 means fully use this encoder layer, 0.0 would mean
+ # completely bypass it.
+ if self.training:
+ alpha = (
+ warmup_scale
+ if torch.rand(()).item() <= (1.0 - self.layer_dropout)
+ else self.bypass_scale
+ )
+ else:
+ alpha = 1.0
+
+ src_att = self.self_attn(
+ src,
+ src,
+ src,
+ attn_mask=src_mask,
+ key_padding_mask=src_key_padding_mask,
+ )[0]
+ src = src + self.dropout(src_att)
+
+ src = src + self.dropout(self.feed_forward(src))
+
+ src = self.norm_final(self.balancer(src))
+
+ if alpha != 1.0:
+ src = alpha * src + (1.0 - alpha) * src_orig
+
+ return src
+
+
+class TransformerDecoderLayer(nn.Module):
+ """Modified from torch.nn.TransformerDecoderLayer.
+
+ Example:
+ >>> 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,
+ bypass_scale: float = 0.1,
+ layer_dropout: float = 0.075,
+ ) -> None:
+
+ """
+ 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).
+ bypass_scale:
+ a scale on the layer's output, used in bypass (resnet-type) skip-connection;
+ when the layer is bypassed, the final output will be a
+ weighted sum of the layer's input and layer's output with weights
+ (1.0-bypass_scale) and bypass_scale correspondingly (default=0.1).
+ layer_dropout:
+ the probability to bypass the layer (default=0.075).
+ """
+
+ super().__init__()
+
+ if bypass_scale < 0.0 or bypass_scale > 1.0:
+ raise ValueError("bypass_scale should be between 0.0 and 1.0")
+
+ if layer_dropout < 0.0 or layer_dropout > 1.0:
+ raise ValueError("layer_dropout should be between 0.0 and 1.0")
+
+ self.bypass_scale = bypass_scale
+ self.layer_dropout = layer_dropout
+
+ self.self_attn = MultiheadAttention(d_model, nhead)
+ self.src_attn = MultiheadAttention(d_model, nhead)
+
+ # Implementation of Feedforward model
+ self.feed_forward = nn.Sequential(
+ ScaledLinear(d_model, dim_feedforward),
+ ActivationBalancer(channel_dim=-1),
+ DoubleSwish(),
+ nn.Dropout(dropout),
+ ScaledLinear(dim_feedforward, d_model, initial_scale=0.25),
+ )
+
+ self.norm_final = BasicNorm(d_model)
+
+ # try to ensure the output is close to zero-mean (or at least, zero-median).
+ self.balancer = ActivationBalancer(
+ channel_dim=-1, min_positive=0.45, max_positive=0.55, max_abs=6.0
+ )
+
+ self.dropout = nn.Dropout(dropout)
+
+ def forward(
+ self,
+ tgt: torch.Tensor,
+ memory: torch.Tensor,
+ tgt_mask: Optional[torch.Tensor] = None,
+ memory_mask: Optional[torch.Tensor] = None,
+ tgt_key_padding_mask: Optional[torch.Tensor] = None,
+ memory_key_padding_mask: Optional[torch.Tensor] = None,
+ warmup: float = 1.0,
+ ) -> torch.Tensor:
+ """Pass the inputs (and mask) through the decoder layer.
+
+ Args:
+ tgt:
+ the sequence to the decoder layer of shape (T, N, C) (required).
+ memory:
+ the sequence from the last layer of the encoder of shape (S, N, C) (required).
+ tgt_mask:
+ the mask for the tgt sequence of shape (T, T) (optional).
+ memory_mask:
+ the mask for the memory sequence of shape (T, S) (optional).
+ tgt_key_padding_mask:
+ the mask for the tgt keys per batch of shape (N, T) (optional).
+ memory_key_padding_mask:
+ the mask for the memory keys per batch of shape (N, S) (optional).
+ warmup: controls selective bypass of layers; if < 1.0, we will
+ bypass the layer more frequently (default=1.0).
+
+ Returns:
+ Output tensor of the shape (T, N, C), where
+ S is the source sequence length,
+ T is the target sequence length,
+ N is the batch size,
+ C is the feature number.
+
+ """
+ tgt_orig = tgt
+
+ warmup_scale = min(self.bypass_scale + warmup, 1.0)
+ # alpha = 1.0 means fully use this encoder layer, 0.0 would mean
+ # completely bypass it.
+ if self.training:
+ alpha = (
+ warmup_scale
+ if torch.rand(()).item() <= (1.0 - self.layer_dropout)
+ else self.bypass_scale
+ )
+ else:
+ alpha = 1.0
+
+ tgt_att = self.self_attn(
+ tgt,
+ tgt,
+ tgt,
+ attn_mask=tgt_mask,
+ key_padding_mask=tgt_key_padding_mask,
+ )[0]
+ tgt = tgt + self.dropout(tgt_att)
+
+ src_att = self.src_attn(
+ tgt,
+ memory,
+ memory,
+ attn_mask=memory_mask,
+ key_padding_mask=memory_key_padding_mask,
+ )[0]
+ tgt = tgt + self.dropout(src_att)
+
+ tgt = tgt + self.dropout(self.feed_forward(tgt))
+
+ tgt = self.norm_final(self.balancer(tgt))
+
+ if alpha != 1.0:
+ tgt = alpha * tgt + (1.0 - alpha) * tgt_orig
+
+ return tgt
+
+
+class TransformerEncoder(nn.Module):
+ """TransformerEncoder is a stack of N encoder layers
+
+ Examples:
+ >>> encoder_layer = TransformerEncoderLayer(d_model=512, nhead=8)
+ >>> transformer_encoder = TransformerEncoder(encoder_layer, num_layers=6)
+ >>> src = torch.rand(10, 32, 512)
+ >>> out = transformer_encoder(src)
+ """
+
+ def __init__(
+ self,
+ encoder_layer: nn.Module,
+ num_layers: int,
+ aux_layers: List[int],
+ ) -> None:
+ """
+ Args:
+ encoder_layer:
+ an instance of the TransformerEncoderLayer() class (required).
+ num_layers:
+ the number of sub-encoder-layers in the encoder (required).
+ aux_layers:
+ list of indexes of sub-encoder-layers outputs to be combined (required).
+ """
+
+ super().__init__()
+ self.layers = nn.ModuleList(
+ [copy.deepcopy(encoder_layer) for i in range(num_layers)]
+ )
+ self.num_layers = num_layers
+
+ assert len(set(aux_layers)) == len(aux_layers)
+
+ assert num_layers - 1 not in aux_layers
+ self.aux_layers = aux_layers + [num_layers - 1]
+
+ self.combiner = RandomCombine(
+ num_inputs=len(self.aux_layers),
+ final_weight=0.5,
+ pure_prob=0.333,
+ stddev=2.0,
+ )
+
+ def forward(
+ self,
+ src: torch.Tensor,
+ mask: Optional[torch.Tensor] = None,
+ src_key_padding_mask: Optional[torch.Tensor] = None,
+ warmup: float = 1.0,
+ ) -> torch.Tensor:
+ """Pass the input through the encoder layers in turn.
+
+ Args:
+ src:
+ the input to the encoder of shape (S, N, C) (required).
+ mask:
+ the mask for the src sequence of shape (S, S) (optional).
+ src_key_padding_mask:
+ the mask for the src keys per batch of shape (N, S) (optional).
+ warmup:
+ controls selective bypass of layer; if < 1.0, we will
+ bypass the layer more frequently (default=1.0).
+
+ Returns:
+ Output tensor of the shape (S, N, C), where
+ S is the source sequence length,
+ N is the batch size,
+ C is the feature number.
+
+ """
+ output = src
+
+ outputs = []
+ for i, mod in enumerate(self.layers):
+ output = mod(
+ output,
+ src_mask=mask,
+ src_key_padding_mask=src_key_padding_mask,
+ warmup=warmup,
+ )
+
+ if i in self.aux_layers:
+ outputs.append(output)
+
+ output = self.combiner(outputs)
+
+ return output
+
+
+class TransformerDecoder(nn.Module):
+ """TransformerDecoder is a stack of N decoder layers
+
+ Examples:
+ >>> decoder_layer = TransformerDecoderLayer(d_model=512, nhead=8)
+ >>> transformer_decoder = TransformerDecoder(decoder_layer, num_layers=6)
+ >>> memory = torch.rand(10, 32, 512)
+ >>> tgt = torch.rand(20, 32, 512)
+ >>> out = transformer_decoder(tgt, memory)
+ """
+
+ def __init__(
+ self,
+ decoder_layer: nn.Module,
+ num_layers: int,
+ aux_layers: List[int],
+ ) -> None:
+ """
+ Args:
+ decoder_layer:
+ an instance of the TransformerDecoderLayer() class (required).
+ num_layers:
+ the number of decoder layers in the decoder (required).
+ aux_layers:
+ list of indexes of decoder layer outputs to be combined (required).
+ """
+
+ super().__init__()
+ self.layers = nn.ModuleList(
+ [copy.deepcopy(decoder_layer) for i in range(num_layers)]
+ )
+ self.num_layers = num_layers
+
+ assert len(set(aux_layers)) == len(aux_layers)
+
+ assert num_layers - 1 not in aux_layers
+ self.aux_layers = aux_layers + [num_layers - 1]
+
+ self.combiner = RandomCombine(
+ num_inputs=len(self.aux_layers),
+ final_weight=0.5,
+ pure_prob=0.333,
+ stddev=2.0,
+ )
+
+ def forward(
+ self,
+ tgt: torch.Tensor,
+ memory: torch.Tensor,
+ tgt_mask: Optional[torch.Tensor] = None,
+ memory_mask: Optional[torch.Tensor] = None,
+ tgt_key_padding_mask: Optional[torch.Tensor] = None,
+ memory_key_padding_mask: Optional[torch.Tensor] = None,
+ warmup: float = 1.0,
+ ) -> torch.Tensor:
+ """Pass the input (and mask) through the decoder layers in turn.
+
+ Args:
+ tgt:
+ the sequence to the decoder of shape (T, N, C) (required).
+ memory:
+ the sequence from the last layer of the encoder of shape (S, N, C) (required).
+ tgt_mask:
+ the mask for the tgt sequence of shape (T, T) (optional).
+ memory_mask:
+ the mask for the memory sequence of shape (T, S) (optional).
+ tgt_key_padding_mask:
+ the mask for the tgt keys per batch of shape (N, T) (optional).
+ memory_key_padding_mask:
+ the mask for the memory keys per batch of shape (N, S) (optional).
+ warmup:
+ controls selective bypass of layer; if < 1.0, we will
+ bypass the layer more frequently (default=1.0).
+
+ Returns:
+ Output tensor of the shape (T, N, C), where
+ S is the source sequence length,
+ T is the target sequence length,
+ N is the batch size,
+ C is the feature number.
+
+ """
+ output = tgt
+
+ outputs = []
+ for i, mod in enumerate(self.layers):
+ output = mod(
+ output,
+ memory,
+ tgt_mask=tgt_mask,
+ memory_mask=memory_mask,
+ tgt_key_padding_mask=tgt_key_padding_mask,
+ memory_key_padding_mask=memory_key_padding_mask,
+ warmup=warmup,
+ )
+
+ if i in self.aux_layers:
+ outputs.append(output)
+
+ output = self.combiner(outputs)
+
+ return output
+
+
+class PositionalEncoding(nn.Module):
+ """This class implements the positional encoding
+ proposed in the following paper:
+
+ - Attention Is All You Need: https://arxiv.org/pdf/1706.03762.pdf
+
+ PE(pos, 2i) = sin(pos / (10000^(2i/d_modle))
+ PE(pos, 2i+1) = cos(pos / (10000^(2i/d_modle))
+
+ Note:
+
+ 1 / (10000^(2i/d_model)) = exp(-log(10000^(2i/d_model)))
+ = exp(-1* 2i / d_model * log(100000))
+ = exp(2i * -(log(10000) / d_model))
+ """
+
+ def __init__(self, d_model: int, dropout: float = 0.1) -> None:
+ """
+ Args:
+ d_model: Embedding dimension.
+ dropout: Dropout probability to be applied to the output of this module.
+ """
+ super().__init__()
+ self.d_model = d_model
+ self.xscale = math.sqrt(self.d_model)
+ self.dropout = nn.Dropout(p=dropout)
+ # not doing: self.pe = None because of errors thrown by torchscript
+ self.pe = torch.zeros(1, 0, self.d_model, dtype=torch.float32)
+
+ def extend_pe(self, x: torch.Tensor) -> None:
+ """Extend the time t in the positional encoding if required.
+ The shape of `self.pe` is (1, T1, d_model). The shape of the input x
+ is (N, T, d_model). If T > T1, then we change the shape of self.pe
+ to (N, T, d_model). Otherwise, nothing is done.
+
+ Args:
+ x:
+ It is a tensor of shape (N, T, C).
+ T is the target sequence length,
+ N is the batch size,
+ C is the feature number.
+ """
+ if self.pe is not None:
+ if self.pe.size(1) >= x.size(1):
+ self.pe = self.pe.to(dtype=x.dtype, device=x.device)
+ return
+ pe = torch.zeros(x.size(1), self.d_model, dtype=torch.float32)
+ 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)
+ # Now pe is of shape (1, T, d_model), where T is x.size(1)
+ self.pe = pe.to(device=x.device, dtype=x.dtype)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Add positional encoding.
+
+ Args:
+ x: Input of shape is (N, T, C)
+
+ Returns:
+ A tensor of the same shape (N, T, C),
+ T is the target sequence length,
+ N is the batch size,
+ C is the feature number.
+
+ """
+ self.extend_pe(x)
+ x = x + self.pe[:, : x.size(1), :]
+ return self.dropout(x)
+
+
+def encoder_padding_mask(
+ max_len: int, supervisions: Optional[Supervisions] = None
+) -> Optional[torch.Tensor]:
+ """Make mask tensor containing indexes of padded part.
+
+ TODO:
+ This function **assumes** that the model uses
+ a subsampling factor of 4. We should remove that
+ assumption later.
+
+ Args:
+ max_len:
+ Maximum length of input features.
+ CAUTION: It is the length after subsampling.
+ supervisions:
+ Supervision in lhotse format.
+ See https://github.com/lhotse-speech/lhotse/blob/master/lhotse/dataset/speech_recognition.py#L32 # noqa
+ (CAUTION: It contains length information, i.e., start and number of
+ frames, before subsampling)
+
+ Returns:
+ Mask tensor of dimension (batch_size, input_length),
+ True denotes 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: torch.Tensor, ignore_id: int = -1) -> torch.Tensor:
+ """Generate a length mask for input.
+
+ The masked position are filled with True,
+ Unmasked positions are filled with False.
+
+ Args:
+ ys_pad:
+ padded tensor of dimension (batch_size, input_length).
+ ignore_id:
+ the ignored number (the padding number) in ys_pad
+
+ Returns:
+ A bool tensor of the same shape as the input tensor.
+ """
+ ys_mask = ys_pad == ignore_id
+ return ys_mask
+
+
+def generate_square_subsequent_mask(sz: int) -> torch.Tensor:
+ """Generate a square mask for the sequence. The masked positions are
+ filled with float('-inf'). Unmasked positions are filled with float(0.0).
+ The mask can be used for masked self-attention.
+
+ For instance, if sz is 3, it returns::
+
+ tensor([[0., -inf, -inf],
+ [0., 0., -inf],
+ [0., 0., 0]])
+
+ Args:
+ sz: mask size
+
+ Returns:
+ A square mask tensor 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(token_ids: List[List[int]], sos_id: int) -> List[List[int]]:
+ """Prepend sos_id to each utterance.
+
+ Args:
+ token_ids:
+ A list-of-list of token IDs. Each sublist contains
+ token IDs (e.g., word piece IDs) of an utterance.
+ sos_id:
+ The ID of the SOS token.
+
+ Return:
+ Return a new list-of-list, where each sublist starts
+ with SOS ID.
+ """
+ return [[sos_id] + utt for utt in token_ids]
+
+
+def add_eos(token_ids: List[List[int]], eos_id: int) -> List[List[int]]:
+ """Append eos_id to each utterance.
+
+ Args:
+ token_ids:
+ A list-of-lists of token IDs. Each sublist contains
+ token IDs (e.g., word piece IDs) of an utterance.
+ eos_id:
+ The ID of the EOS token.
+
+ Return:
+ Return a new list-of-lists, where each sublist ends
+ with EOS ID.
+ """
+ return [utt + [eos_id] for utt in token_ids]
+
+
+def tolist(t: torch.Tensor) -> List[int]:
+ """Used by jit"""
+ return torch.jit.annotate(List[int], t.tolist())
diff --git a/egs/tedlium3/ASR/local/convert_transcript_words_to_bpe_ids.py b/egs/tedlium3/ASR/local/convert_transcript_words_to_bpe_ids.py
index 9dbcc9d9e..19ba8d24b 100644
--- a/egs/tedlium3/ASR/local/convert_transcript_words_to_bpe_ids.py
+++ b/egs/tedlium3/ASR/local/convert_transcript_words_to_bpe_ids.py
@@ -4,16 +4,18 @@
"""
Convert a transcript based on words to a list of BPE ids.
-For example, if we use 2 as the encoding id of :
+For example, if we use 2 as the encoding id of
+Note: it, inserts a space token before each
texts = ['this is a day']
-spm_ids = [[38, 33, 6, 2, 316]]
+spm_ids = [[38, 33, 6, 15, 2, 316]]
texts = [' this is a sunny day']
-spm_ids = [[2, 38, 33, 6, 118, 11, 11, 21, 316]]
+spm_ids = [[15, 2, 38, 33, 6, 118, 11, 11, 21, 316]]
texts = ['']
-spm_ids = [[2]]
+spm_ids = [[15, 2]]
+
"""
import argparse
@@ -38,29 +40,27 @@ def get_args():
def convert_texts_into_ids(
texts: List[str],
- unk_id: int,
sp: spm.SentencePieceProcessor,
) -> List[List[int]]:
"""
Args:
texts:
A string list of transcripts, such as ['Today is Monday', 'It's sunny'].
- unk_id:
- A number id for the token ''.
+ sp:
+ A sentencepiece BPE model.
Returns:
Return an integer list of bpe ids.
"""
y = []
for text in texts:
- y_ids = []
if "" in text:
- text_segments = text.split("")
- id_segments = sp.encode(text_segments, out_type=int)
+ id_segments = sp.encode(text.split(""), out_type=int)
+
+ y_ids = []
for i in range(len(id_segments)):
- if i != len(id_segments) - 1:
- y_ids.extend(id_segments[i] + [unk_id])
- else:
- y_ids.extend(id_segments[i])
+ y_ids += id_segments[i]
+ if i < len(id_segments) - 1:
+ y_ids += [sp.piece_to_id("▁"), sp.unk_id()]
else:
y_ids = sp.encode(text, out_type=int)
y.append(y_ids)
@@ -70,19 +70,13 @@ def convert_texts_into_ids(
def main():
args = get_args()
- texts = args.texts
- bpe_model = args.bpe_model
sp = spm.SentencePieceProcessor()
- sp.load(bpe_model)
- unk_id = sp.piece_to_id("")
+ sp.load(args.bpe_model)
- y = convert_texts_into_ids(
- texts=texts,
- unk_id=unk_id,
- sp=sp,
- )
- logging.info(f"The input texts: {texts}")
+ y = convert_texts_into_ids(texts=args.texts, sp=sp)
+
+ logging.info(f"The input texts: {args.texts}")
logging.info(f"The encoding ids: {y}")
diff --git a/egs/tedlium3/ASR/local/convert_transcript_words_to_tokens.py b/egs/tedlium3/ASR/local/convert_transcript_words_to_tokens.py
deleted file mode 120000
index 2ce13fd69..000000000
--- a/egs/tedlium3/ASR/local/convert_transcript_words_to_tokens.py
+++ /dev/null
@@ -1 +0,0 @@
-../../../librispeech/ASR/local/convert_transcript_words_to_tokens.py
\ No newline at end of file
diff --git a/egs/tedlium3/ASR/local/generate_unique_lexicon.py b/egs/tedlium3/ASR/local/generate_unique_lexicon.py
deleted file mode 120000
index c0aea1403..000000000
--- a/egs/tedlium3/ASR/local/generate_unique_lexicon.py
+++ /dev/null
@@ -1 +0,0 @@
-../../../librispeech/ASR/local/generate_unique_lexicon.py
\ No newline at end of file
diff --git a/egs/tedlium3/ASR/local/prepare_lang.py b/egs/tedlium3/ASR/local/prepare_lang.py
deleted file mode 120000
index 747f2ab39..000000000
--- a/egs/tedlium3/ASR/local/prepare_lang.py
+++ /dev/null
@@ -1 +0,0 @@
-../../../librispeech/ASR/local/prepare_lang.py
\ No newline at end of file
diff --git a/egs/tedlium3/ASR/local/prepare_lexicon.py b/egs/tedlium3/ASR/local/prepare_lexicon.py
deleted file mode 100755
index b9160b6d4..000000000
--- a/egs/tedlium3/ASR/local/prepare_lexicon.py
+++ /dev/null
@@ -1,94 +0,0 @@
-#!/usr/bin/env python3
-# Copyright 2022 Xiaomi Corp. (authors: Mingshuang Luo)
-#
-# 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.
-
-
-"""
-This script takes as input supervisions json dir "data/manifests"
-consisting of supervisions_train.json and does the following:
-
-1. Generate lexicon_words.txt.
-
-"""
-import argparse
-import logging
-from pathlib import Path
-
-import lhotse
-
-
-def get_args():
- parser = argparse.ArgumentParser()
- parser.add_argument(
- "--manifests-dir",
- type=str,
- help="""Input directory.
- """,
- )
- parser.add_argument(
- "--lang-dir",
- type=str,
- help="""Output directory.
- """,
- )
-
- return parser.parse_args()
-
-
-def prepare_lexicon(manifests_dir: str, lang_dir: str):
- """
- Args:
- manifests_dir:
- The manifests directory, e.g., data/manifests.
- lang_dir:
- The language directory, e.g., data/lang_phone.
-
- Return:
- The lexicon_words.txt file.
- """
- words = set()
-
- lexicon = Path(lang_dir) / "lexicon_words.txt"
- sups = lhotse.load_manifest(f"{manifests_dir}/tedlium_supervisions_train.jsonl.gz")
- for s in sups:
- # list the words units and filter the empty item
- words_list = list(filter(None, s.text.split()))
-
- for word in words_list:
- if word not in words and word != "":
- words.add(word)
-
- with open(lexicon, "w") as f:
- for word in sorted(words):
- f.write(word + " " + word)
- f.write("\n")
-
-
-def main():
- args = get_args()
- manifests_dir = Path(args.manifests_dir)
- lang_dir = Path(args.lang_dir)
-
- logging.info("Generating lexicon_words.txt")
- prepare_lexicon(manifests_dir, lang_dir)
-
-
-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/tedlium3/ASR/local/prepare_transcripts.py b/egs/tedlium3/ASR/local/prepare_transcripts.py
index 7ea4e89a4..d4ccdd1e3 100755
--- a/egs/tedlium3/ASR/local/prepare_transcripts.py
+++ b/egs/tedlium3/ASR/local/prepare_transcripts.py
@@ -1,5 +1,6 @@
#!/usr/bin/env python3
-# Copyright 2021 Xiaomi Corp. (authors: Mingshuang Luo)
+# Copyright 2021 Xiaomi Corp. (author: Mingshuang Luo)
+# Copyright 2022 Behavox LLC. (author: Daniil Kulko)
#
# See ../../../../LICENSE for clarification regarding multiple authors
#
@@ -17,68 +18,67 @@
"""
-This script takes as input supervisions json dir "data/manifests"
-consisting of supervisions_train.json and does the following:
-
-1. Generate train.text.
+This script takes input text file and removes all words
+that iclude any character out of English alphabet.
"""
import argparse
import logging
+import re
from pathlib import Path
-import lhotse
-
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument(
- "--manifests-dir",
+ "--input-text-path",
type=str,
- help="""Input directory.
- """,
+ help="Input text file path.",
)
parser.add_argument(
- "--lang-dir",
+ "--output-text-path",
type=str,
- help="""Output directory.
- """,
+ help="Output text file path.",
)
return parser.parse_args()
-def prepare_transcripts(manifests_dir: str, lang_dir: str):
+def prepare_transcripts(input_text_path: Path, output_text_path: Path) -> None:
"""
Args:
- manifests_dir:
- The manifests directory, e.g., data/manifests.
- lang_dir:
- The language directory, e.g., data/lang_phone.
+ input_text_path:
+ The input data text file path, e.g., data/lang/train_orig.txt.
+ output_text_path:
+ The output data text file path, e.g., data/lang/train.txt.
Return:
- The train.text in lang_dir.
+ Saved text file in output_text_path.
"""
- texts = []
- train_text = Path(lang_dir) / "train.text"
- sups = lhotse.load_manifest(f"{manifests_dir}/tedlium_supervisions_train.jsonl.gz")
- for s in sups:
- texts.append(s.text)
+ foreign_chr_check = re.compile(r"[^a-z']")
- with open(train_text, "w") as f:
- for text in texts:
- f.write(text)
- f.write("\n")
+ logging.info(f"Loading {input_text_path.name}")
+ with open(input_text_path, "r", encoding="utf8") as f:
+ texts = {t.rstrip("\n") for t in f}
+
+ texts = {
+ " ".join([w for w in t.split() if foreign_chr_check.search(w) is None])
+ for t in texts
+ }
+
+ with open(output_text_path, "w+", encoding="utf8") as f:
+ for t in texts:
+ f.write(f"{t}\n")
-def main():
+def main() -> None:
args = get_args()
- manifests_dir = Path(args.manifests_dir)
- lang_dir = Path(args.lang_dir)
+ input_text_path = Path(args.input_text_path)
+ output_text_path = Path(args.output_text_path)
- logging.info("Generating train.text")
- prepare_transcripts(manifests_dir, lang_dir)
+ logging.info(f"Generating {output_text_path.name}")
+ prepare_transcripts(input_text_path, output_text_path)
if __name__ == "__main__":
diff --git a/egs/tedlium3/ASR/local/prepare_words.py b/egs/tedlium3/ASR/local/prepare_words.py
new file mode 100755
index 000000000..a37d0f08f
--- /dev/null
+++ b/egs/tedlium3/ASR/local/prepare_words.py
@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+# Copyright 2022 Behavox LLC. (authors: Daniil Kulko)
+#
+# 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.
+
+
+"""
+This script takes as input supervisions json dir "data/manifests"
+consisting of tedlium_supervisions_train.json and does the following:
+
+1. Generate words.txt.
+
+"""
+import argparse
+import logging
+import re
+from pathlib import Path
+
+
+def get_args():
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--lang-dir",
+ type=str,
+ help="Output directory.",
+ )
+
+ return parser.parse_args()
+
+
+def prepare_words(lang_dir: str) -> None:
+ """
+ Args:
+ lang_dir:
+ The language directory, e.g., data/lang.
+
+ Return:
+ The words.txt file.
+ """
+
+ words_orig_path = Path(lang_dir) / "words_orig.txt"
+ words_path = Path(lang_dir) / "words.txt"
+
+ foreign_chr_check = re.compile(r"[^a-z']")
+
+ logging.info(f"Loading {words_orig_path.name}")
+ with open(words_orig_path, "r", encoding="utf8") as f:
+ words = {w for w_compl in f for w in w_compl.strip("-\n").split("_")}
+ words = {w for w in words if foreign_chr_check.search(w) is None and w != ""}
+ words.add("")
+ words = ["", "!SIL"] + sorted(words) + ["#0", "", ""]
+
+ with open(words_path, "w+", encoding="utf8") as f:
+ for idx, word in enumerate(words):
+ f.write(f"{word} {idx}\n")
+
+
+def main() -> None:
+ args = get_args()
+ lang_dir = Path(args.lang_dir)
+
+ logging.info("Generating words.txt")
+ prepare_words(lang_dir)
+
+
+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/tedlium3/ASR/local/test_prepare_lang.py b/egs/tedlium3/ASR/local/test_prepare_lang.py
deleted file mode 120000
index f0f864998..000000000
--- a/egs/tedlium3/ASR/local/test_prepare_lang.py
+++ /dev/null
@@ -1 +0,0 @@
-../../../librispeech/ASR/local/test_prepare_lang.py
\ No newline at end of file
diff --git a/egs/tedlium3/ASR/prepare.sh b/egs/tedlium3/ASR/prepare.sh
index 272cf7aed..3d90436ff 100755
--- a/egs/tedlium3/ASR/prepare.sh
+++ b/egs/tedlium3/ASR/prepare.sh
@@ -5,7 +5,6 @@ export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
set -eou pipefail
-nj=15
stage=0
stop_stage=100
@@ -63,6 +62,13 @@ if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then
mv $dl_dir/TEDLIUM_release-3 $dl_dir/tedlium3
fi
+ # Download big and small 4 gram lanuage models
+ if [ ! -d $dl_dir/lm ]; then
+ wget --continue http://kaldi-asr.org/models/5/4gram_small.arpa.gz -P $dl_dir/lm
+ wget --continue http://kaldi-asr.org/models/5/4gram_big.arpa.gz -P $dl_dir/lm
+ gzip -d $dl_dir/lm/4gram_small.arpa.gz $dl_dir/lm/4gram_big.arpa.gz
+ fi
+
# If you have pre-downloaded it to /path/to/musan,
# you can create a symlink
#
@@ -100,7 +106,14 @@ if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
if [ ! -e data/fbank/.tedlium3.done ]; then
mkdir -p data/fbank
+
python3 ./local/compute_fbank_tedlium.py
+
+ gunzip -c data/fbank/tedlium_cuts_train.jsonl.gz | shuf | \
+ gzip -c > data/fbank/tedlium_cuts_train-shuf.jsonl.gz
+ mv data/fbank/tedlium_cuts_train-shuf.jsonl.gz \
+ data/fbank/tedlium_cuts_train.jsonl.gz
+
touch data/fbank/.tedlium3.done
fi
fi
@@ -115,28 +128,24 @@ if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
fi
if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then
- log "Stage 5: Prepare phone based lang"
- lang_dir=data/lang_phone
+ log "Stage 5: Prepare BPE train data and set of words"
+ lang_dir=data/lang
mkdir -p $lang_dir
- if [ ! -f $lang_dir/train.text ]; then
+ if [ ! -f $lang_dir/train.txt ]; then
+ gunzip -c $dl_dir/tedlium3/LM/*.en.gz | sed 's: <\/s>::g' > $lang_dir/train_orig.txt
+
./local/prepare_transcripts.py \
- --lang-dir $lang_dir \
- --manifests-dir data/manifests
+ --input-text-path $lang_dir/train_orig.txt \
+ --output-text-path $lang_dir/train.txt
fi
- if [ ! -f $lang_dir/lexicon_words.txt ]; then
- ./local/prepare_lexicon.py \
- --lang-dir $lang_dir \
- --manifests-dir data/manifests
- fi
+ if [ ! -f $lang_dir/words.txt ]; then
- (echo '!SIL SIL'; echo ' '; ) |
- cat - $lang_dir/lexicon_words.txt |
- sort | uniq > $lang_dir/lexicon.txt
+ awk '{print $1}' $dl_dir/tedlium3/TEDLIUM.152k.dic |
+ sed 's:([0-9])::g' | sort | uniq > $lang_dir/words_orig.txt
- if [ ! -f $lang_dir/L_disambig.pt ]; then
- ./local/prepare_lang.py --lang-dir $lang_dir
+ ./local/prepare_words.py --lang-dir $lang_dir
fi
fi
@@ -148,25 +157,56 @@ if [ $stage -le 6 ] && [ $stop_stage -ge 6 ]; then
mkdir -p $lang_dir
# We reuse words.txt from phone based lexicon
# so that the two can share G.pt later.
- cp data/lang_phone/words.txt $lang_dir
-
- if [ ! -f $lang_dir/transcript_words.txt ]; then
- log "Generate data for BPE training"
- cat data/lang_phone/train.text |
- cut -d " " -f 2- > $lang_dir/transcript_words.txt
- # remove the for transcript_words.txt
- sed -i 's/ //g' $lang_dir/transcript_words.txt
- sed -i 's/ //g' $lang_dir/transcript_words.txt
- sed -i 's///g' $lang_dir/transcript_words.txt
- fi
+ cp data/lang/words.txt $lang_dir
./local/train_bpe_model.py \
--lang-dir $lang_dir \
--vocab-size $vocab_size \
- --transcript $lang_dir/transcript_words.txt
+ --transcript data/lang/train.txt
if [ ! -f $lang_dir/L_disambig.pt ]; then
- ./local/prepare_lang_bpe.py --lang-dir $lang_dir
+ ./local/prepare_lang_bpe.py --lang-dir $lang_dir --oov ""
+ fi
+ done
+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
+
+ mkdir -p data/lm
+ if [ ! -f data/lm/G_4_gram_small.fst.txt ]; then
+ # It is used in building HLG
+ python3 -m kaldilm \
+ --read-symbol-table="data/lang/words.txt" \
+ --disambig-symbol='#0' \
+ --max-order=4 \
+ --max-arpa-warnings=-1 \
+ $dl_dir/lm/4gram_small.arpa > data/lm/G_4_gram_small.fst.txt
+ fi
+
+ if [ ! -f data/lm/G_4_gram_big.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 \
+ --max-arpa-warnings=-1 \
+ $dl_dir/lm/4gram_big.arpa > data/lm/G_4_gram_big.fst.txt
+ fi
+fi
+
+if [ $stage -le 8 ] && [ $stop_stage -ge 8 ]; then
+ log "Stage 8: Compile HLG"
+
+ for vocab_size in ${vocab_sizes[@]}; do
+ lang_dir=data/lang_bpe_${vocab_size}
+
+ if [ ! -f $lang_dir/HLG.pt ]; then
+ ./local/compile_hlg.py \
+ --lang-dir $lang_dir \
+ --lm G_4_gram_small
fi
done
fi
diff --git a/egs/wenetspeech/ASR/pruned_transducer_stateless2/train.py b/egs/wenetspeech/ASR/pruned_transducer_stateless2/train.py
index 43fa0d01b..48b347b64 100644
--- a/egs/wenetspeech/ASR/pruned_transducer_stateless2/train.py
+++ b/egs/wenetspeech/ASR/pruned_transducer_stateless2/train.py
@@ -861,15 +861,41 @@ def run(rank, world_size, args):
valid_cuts = wenetspeech.valid_cuts()
def remove_short_and_long_utt(c: Cut):
- # Keep only utterances with duration between 1 second and 15.0 seconds
+ # Keep only utterances with duration between 1 second and 10 seconds
#
- # Caution: There is a reason to select 15.0 here. Please see
+ # Caution: There is a reason to select 10.0 here. Please see
# ../local/display_manifest_statistics.py
#
# You should use ../local/display_manifest_statistics.py to get
# an utterance duration distribution for your dataset to select
# the threshold
- return 1.0 <= c.duration <= 15.0
+ if c.duration < 1.0 or c.duration > 10.0:
+ logging.warning(
+ f"Exclude cut with ID {c.id} from training. Duration: {c.duration}"
+ )
+ return False
+
+ # In pruned RNN-T, we require that T >= S
+ # where T is the number of feature frames after subsampling
+ # and S is the number of tokens in the utterance
+
+ # In ./conformer.py, the conv module uses the following expression
+ # for subsampling
+ T = ((c.num_frames - 1) // 2 - 1) // 2
+ tokens = c.supervisions[0].text.replace(" ", "")
+
+ if T < len(tokens):
+ logging.warning(
+ f"Exclude cut with ID {c.id} from training. "
+ f"Number of frames (before subsampling): {c.num_frames}. "
+ f"Number of frames (after subsampling): {T}. "
+ f"Text: {c.supervisions[0].text}. "
+ f"Tokens: {tokens}. "
+ f"Number of tokens: {len(tokens)}"
+ )
+ return False
+
+ return True
train_cuts = train_cuts.filter(remove_short_and_long_utt)
diff --git a/egs/wenetspeech/ASR/pruned_transducer_stateless5/conformer.py b/egs/wenetspeech/ASR/pruned_transducer_stateless5/conformer.py
index 9bb55d07a..23a877b2f 100644
--- a/egs/wenetspeech/ASR/pruned_transducer_stateless5/conformer.py
+++ b/egs/wenetspeech/ASR/pruned_transducer_stateless5/conformer.py
@@ -966,20 +966,32 @@ class RelPositionMultiheadAttention(nn.Module):
(batch_size, num_heads, time1, n) = x.shape
time2 = time1 + left_context
- assert (
- n == left_context + 2 * time1 - 1
- ), f"{n} == {left_context} + 2 * {time1} - 1"
+ if not torch.jit.is_tracing():
+ assert (
+ n == left_context + 2 * time1 - 1
+ ), f"{n} == {left_context} + 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, time2),
- (batch_stride, head_stride, time1_stride - n_stride, n_stride),
- storage_offset=n_stride * (time1 - 1),
- )
+ if torch.jit.is_tracing():
+ rows = torch.arange(start=time1 - 1, end=-1, step=-1)
+ cols = torch.arange(time2)
+ rows = rows.repeat(batch_size * num_heads).unsqueeze(-1)
+ indexes = rows + cols
+
+ x = x.reshape(-1, n)
+ x = torch.gather(x, dim=1, index=indexes)
+ x = x.reshape(batch_size, num_heads, time1, time2)
+ return x
+ else:
+ # 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, time2),
+ (batch_stride, head_stride, time1_stride - n_stride, n_stride),
+ storage_offset=n_stride * (time1 - 1),
+ )
def multi_head_attention_forward(
self,
diff --git a/egs/wenetspeech/ASR/pruned_transducer_stateless5/train.py b/egs/wenetspeech/ASR/pruned_transducer_stateless5/train.py
index 440b65f32..34a72be8f 100755
--- a/egs/wenetspeech/ASR/pruned_transducer_stateless5/train.py
+++ b/egs/wenetspeech/ASR/pruned_transducer_stateless5/train.py
@@ -1006,15 +1006,41 @@ def run(rank, world_size, args):
valid_cuts = wenetspeech.valid_cuts()
def remove_short_and_long_utt(c: Cut):
- # Keep only utterances with duration between 1 second and 15.0 seconds
+ # Keep only utterances with duration between 1 second and 10 seconds
#
- # Caution: There is a reason to select 15.0 here. Please see
+ # Caution: There is a reason to select 10.0 here. Please see
# ../local/display_manifest_statistics.py
#
# You should use ../local/display_manifest_statistics.py to get
# an utterance duration distribution for your dataset to select
# the threshold
- return 1.0 <= c.duration <= 15.0
+ if c.duration < 1.0 or c.duration > 10.0:
+ logging.warning(
+ f"Exclude cut with ID {c.id} from training. Duration: {c.duration}"
+ )
+ return False
+
+ # In pruned RNN-T, we require that T >= S
+ # where T is the number of feature frames after subsampling
+ # and S is the number of tokens in the utterance
+
+ # In ./conformer.py, the conv module uses the following expression
+ # for subsampling
+ T = ((c.num_frames - 1) // 2 - 1) // 2
+ tokens = c.supervisions[0].text.replace(" ", "")
+
+ if T < len(tokens):
+ logging.warning(
+ f"Exclude cut with ID {c.id} from training. "
+ f"Number of frames (before subsampling): {c.num_frames}. "
+ f"Number of frames (after subsampling): {T}. "
+ f"Text: {c.supervisions[0].text}. "
+ f"Tokens: {tokens}. "
+ f"Number of tokens: {len(tokens)}"
+ )
+ return False
+
+ return True
train_cuts = train_cuts.filter(remove_short_and_long_utt)
diff --git a/icefall/__init__.py b/icefall/__init__.py
index 27ad74213..82d21706c 100644
--- a/icefall/__init__.py
+++ b/icefall/__init__.py
@@ -68,3 +68,5 @@ from .utils import (
)
from .ngram_lm import NgramLm, NgramLmStateCost
+
+from .lm_wrapper import LmScorer
diff --git a/icefall/decode.py b/icefall/decode.py
index e4c614c4e..23f9fb9b3 100644
--- a/icefall/decode.py
+++ b/icefall/decode.py
@@ -466,9 +466,7 @@ def one_best_decoding(
Return:
An FsaVec containing linear paths.
"""
-
if lm_scale_list is not None:
-
ans = dict()
saved_am_scores = lattice.scores - lattice.lm_scores
for lm_scale in lm_scale_list:
@@ -717,6 +715,107 @@ def rescore_with_n_best_list(
return ans
+def nbest_rescore_with_LM(
+ lattice: k2.Fsa,
+ LM: k2.Fsa,
+ num_paths: int,
+ lm_scale_list: List[float],
+ nbest_scale: float = 1.0,
+ use_double_scores: bool = True,
+) -> Dict[str, k2.Fsa]:
+ """Rescore an n-best list with an n-gram LM.
+ The path with the maximum score is used as the decoding output.
+
+ Args:
+ lattice:
+ An FsaVec with axes [utt][state][arc]. It must have the following
+ attributes: ``aux_labels`` and ``lm_scores``. They are both token
+ IDs.
+ LM:
+ An FsaVec containing only a single FSA. It is one of follows:
+ - LG, L is lexicon and G is word-level n-gram LM.
+ - G, token-level n-gram LM.
+ num_paths:
+ Size of nbest list.
+ lm_scale_list:
+ A list of floats representing LM score scales.
+ nbest_scale:
+ Scale to be applied to ``lattice.score`` when sampling paths
+ using ``k2.random_paths``.
+ use_double_scores:
+ True to use double precision during computation. False to use
+ single precision.
+ Returns:
+ A dict of FsaVec, whose key is an lm_scale and the value is the
+ best decoding path for each utterance in the lattice.
+ """
+ device = lattice.device
+
+ assert len(lattice.shape) == 3
+ assert hasattr(lattice, "aux_labels")
+ assert hasattr(lattice, "lm_scores")
+
+ assert LM.shape == (1, None, None)
+ assert LM.device == device
+
+ nbest = Nbest.from_lattice(
+ lattice=lattice,
+ num_paths=num_paths,
+ use_double_scores=use_double_scores,
+ nbest_scale=nbest_scale,
+ )
+ # nbest.fsa.scores contains 0s
+
+ nbest = nbest.intersect(lattice)
+
+ # Now nbest.fsa has its scores set
+ assert hasattr(nbest.fsa, "lm_scores")
+
+ # am scores + bi-gram scores
+ hp_scores = nbest.tot_scores()
+
+ # Now start to intersect nbest with LG or G
+ inv_fsa = k2.invert(nbest.fsa)
+ if hasattr(LM, "aux_labels"):
+ # LM is LG here
+ # delete token IDs as it is not needed
+ del inv_fsa.aux_labels
+ inv_fsa.scores.zero_()
+ inv_fsa_with_epsilon_loops = k2.linear_fsa_with_self_loops(inv_fsa)
+ path_to_utt_map = nbest.shape.row_ids(1)
+
+ LM = k2.arc_sort(LM)
+ path_lattice = k2.intersect_device(
+ LM,
+ inv_fsa_with_epsilon_loops,
+ b_to_a_map=torch.zeros_like(path_to_utt_map),
+ sorted_match_a=True,
+ )
+
+ # Its labels are token IDs.
+ # If LM is G, its aux_labels are tokens IDs;
+ # If LM is LG, its aux_labels are words IDs.
+ path_lattice = k2.top_sort(k2.connect(path_lattice))
+ one_best = k2.shortest_path(path_lattice, use_double_scores=use_double_scores)
+
+ lm_scores = one_best.get_tot_scores(
+ use_double_scores=use_double_scores,
+ log_semiring=True, # Note: we always use True
+ )
+ # If LM is LG, we might get empty paths
+ lm_scores[lm_scores == float("-inf")] = -1e9
+
+ ans = dict()
+ for lm_scale in lm_scale_list:
+ tot_scores = hp_scores.values / lm_scale + lm_scores
+ tot_scores = k2.RaggedTensor(nbest.shape, tot_scores)
+ max_indexes = tot_scores.argmax()
+ best_path = k2.index_fsa(nbest.fsa, max_indexes)
+ key = f"lm_scale_{lm_scale}"
+ ans[key] = best_path
+ return ans
+
+
def rescore_with_whole_lattice(
lattice: k2.Fsa,
G_with_epsilon_loops: k2.Fsa,
diff --git a/icefall/dist.py b/icefall/dist.py
index 9df1c5bd1..922f31a2f 100644
--- a/icefall/dist.py
+++ b/icefall/dist.py
@@ -21,12 +21,16 @@ import torch
from torch import distributed as dist
-def setup_dist(rank, world_size, master_port=None, use_ddp_launch=False):
+def setup_dist(
+ rank, world_size, master_port=None, use_ddp_launch=False, master_addr=None
+):
"""
rank and world_size are used only if use_ddp_launch is False.
"""
if "MASTER_ADDR" not in os.environ:
- os.environ["MASTER_ADDR"] = "localhost"
+ os.environ["MASTER_ADDR"] = (
+ "localhost" if master_addr is None else str(master_addr)
+ )
if "MASTER_PORT" not in os.environ:
os.environ["MASTER_PORT"] = "12354" if master_port is None else str(master_port)
diff --git a/icefall/lm_wrapper.py b/icefall/lm_wrapper.py
new file mode 100644
index 000000000..0468befd0
--- /dev/null
+++ b/icefall/lm_wrapper.py
@@ -0,0 +1,254 @@
+# Copyright (c) 2022 Xiaomi Corporation (authors: Xiaoyu Yang)
+#
+# 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.
+
+import argparse
+import logging
+
+import torch
+
+from icefall.checkpoint import average_checkpoints, load_checkpoint
+from icefall.rnn_lm.model import RnnLmModel
+from icefall.transformer_lm.model import TransformerLM
+from icefall.utils import AttributeDict, str2bool
+
+
+class LmScorer(torch.nn.Module):
+ """This is a wrapper for NN LMs
+ The language models supported include:
+ RNN,
+ Transformer
+ """
+
+ def __init__(
+ self,
+ lm_type: str,
+ params: AttributeDict,
+ device,
+ lm_scale: float = 0.3,
+ ):
+ super(LmScorer, self).__init__()
+ assert lm_type in ["rnn", "transformer"], f"{lm_type} is not supported"
+ self.lm_type = lm_type
+ self.lm = self.get_lm(lm_type, device, params)
+ self.lm_scale = lm_scale
+ self.params = params
+
+ @classmethod
+ def add_arguments(cls, parser):
+ # LM general arguments
+ parser.add_argument(
+ "--vocab-size",
+ type=int,
+ default=500,
+ )
+
+ parser.add_argument(
+ "--lm-epoch",
+ type=int,
+ default=7,
+ help="""Which epoch to be used
+ """,
+ )
+
+ parser.add_argument(
+ "--lm-avg",
+ type=int,
+ default=1,
+ help="""Number of checkpoints to be averaged
+ """,
+ )
+
+ parser.add_argument("--lm-exp-dir", type=str, help="Path to LM experiments")
+
+ # Now RNNLM related arguments
+ parser.add_argument(
+ "--rnn-lm-embedding-dim",
+ type=int,
+ default=2048,
+ help="Embedding dim of the model",
+ )
+
+ parser.add_argument(
+ "--rnn-lm-hidden-dim",
+ type=int,
+ default=2048,
+ help="Hidden dim of the model",
+ )
+
+ parser.add_argument(
+ "--rnn-lm-num-layers",
+ type=int,
+ default=3,
+ help="Number of RNN layers the model",
+ )
+
+ parser.add_argument(
+ "--rnn-lm-tie-weights",
+ type=str2bool,
+ default=True,
+ help="""True to share the weights between the input embedding layer and the
+ last output linear layer
+ """,
+ )
+
+ # Now transformers
+ parser.add_argument(
+ "--transformer-lm-exp-dir", type=str, help="Directory of transformer LM exp"
+ )
+
+ parser.add_argument(
+ "--transformer-lm-dim-feedforward",
+ type=int,
+ default=2048,
+ help="Dimension of FFW module in transformer",
+ )
+
+ parser.add_argument(
+ "--transformer-lm-encoder-dim",
+ type=int,
+ default=768,
+ help="Encoder dimension of transformer",
+ )
+
+ parser.add_argument(
+ "--transformer-lm-embedding-dim",
+ type=int,
+ default=768,
+ help="Input embedding dimension of transformer",
+ )
+
+ parser.add_argument(
+ "--transformer-lm-nhead",
+ type=int,
+ default=8,
+ help="Number of attention heads in transformer",
+ )
+
+ parser.add_argument(
+ "--transformer-lm-num-layers",
+ type=int,
+ default=16,
+ help="Number of encoder layers in transformer",
+ )
+
+ parser.add_argument(
+ "--transformer-lm-tie-weights",
+ type=str2bool,
+ default=True,
+ help="If tie weights in transformer LM",
+ )
+
+ def get_lm(self, lm_type: str, device, params: AttributeDict) -> torch.nn.Module:
+ """Return the neural network LM
+
+ Args:
+ lm_type (str): Type name of NN LM
+ """
+ if lm_type == "rnn":
+ model = RnnLmModel(
+ vocab_size=params.vocab_size,
+ embedding_dim=params.rnn_lm_embedding_dim,
+ hidden_dim=params.rnn_lm_hidden_dim,
+ num_layers=params.rnn_lm_num_layers,
+ tie_weights=params.rnn_lm_tie_weights,
+ )
+
+ if params.lm_avg == 1:
+ load_checkpoint(
+ f"{params.lm_exp_dir}/epoch-{params.lm_epoch}.pt", model
+ )
+ model.to(device)
+ else:
+ start = params.lm_epoch - params.lm_avg + 1
+ filenames = []
+ for i in range(start, params.lm_epoch + 1):
+ if start >= 0:
+ filenames.append(f"{params.lm_exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+
+ elif lm_type == "transformer":
+ model = TransformerLM(
+ vocab_size=params.vocab_size,
+ d_model=params.transformer_lm_encoder_dim,
+ embedding_dim=params.transformer_lm_embedding_dim,
+ dim_feedforward=params.transformer_lm_dim_feedforward,
+ nhead=params.transformer_lm_nhead,
+ num_layers=params.transformer_lm_num_layers,
+ tie_weights=params.transformer_lm_tie_weights,
+ params=params,
+ )
+
+ if params.lm_avg == 1:
+ load_checkpoint(
+ f"{params.lm_exp_dir}/epoch-{params.lm_epoch}.pt", model
+ )
+ model.to(device)
+ else:
+ start = params.lm_epoch - params.lm_avg + 1
+ filenames = []
+ for i in range(start, params.lm_epoch + 1):
+ if start >= 0:
+ filenames.append(f"{params.lm_exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ raise NotImplementedError()
+
+ return model
+
+ def score_token(self, x: torch.Tensor, x_lens: torch.Tensor, state=None):
+ """Score the input and return the prediction
+ This requires the lm to have the method `score_token`
+ Args:
+ x (torch.Tensor): Input tokens
+ x_lens (torch.Tensor): Length of the input tokens
+ state (optional): LM states
+
+ """
+ return self.lm.score_token(x, x_lens, state)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ LmScorer.add_arguments(parser)
+ args = parser.parse_args()
+
+ params = AttributeDict()
+ params.update(vars(args))
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ Scorer = LmScorer(params=params, device=device)
+ Scorer.eval()
+
+ x = (
+ torch.tensor([[1, 4, 19, 256, 77], [1, 4, 19, 256, 77]])
+ .to(device)
+ .to(torch.int64)
+ )
+ x_lens = torch.tensor([5, 5]).to(device)
+
+ state = None
+
+ score, state = Scorer.score(x, x_lens)
+ print(score.shape)
+ print(score[0])
+ print(score[1])
diff --git a/icefall/mmi.py b/icefall/mmi.py
index 16ed6e032..b7777b434 100644
--- a/icefall/mmi.py
+++ b/icefall/mmi.py
@@ -112,8 +112,12 @@ def _compute_mmi_loss_exact_non_optimized(
num_graphs, den_graphs = graph_compiler.compile(texts, replicate_den=True)
# TODO: pass output_beam as function argument
- num_lats = k2.intersect_dense(num_graphs, dense_fsa_vec, output_beam=beam_size)
- den_lats = k2.intersect_dense(den_graphs, dense_fsa_vec, output_beam=beam_size)
+ num_lats = k2.intersect_dense(
+ num_graphs, dense_fsa_vec, output_beam=beam_size, max_arcs=2147483600
+ )
+ den_lats = k2.intersect_dense(
+ den_graphs, dense_fsa_vec, output_beam=beam_size, max_arcs=2147483600
+ )
num_tot_scores = num_lats.get_tot_scores(log_semiring=True, use_double_scores=True)
@@ -144,7 +148,7 @@ def _compute_mmi_loss_pruned(
"""
num_graphs, den_graphs = graph_compiler.compile(texts, replicate_den=False)
- num_lats = k2.intersect_dense(num_graphs, dense_fsa_vec, output_beam=10.0)
+ num_lats = k2.intersect_dense(num_graphs, dense_fsa_vec, output_beam=8.0)
# the values for search_beam/output_beam/min_active_states/max_active_states
# are not tuned. You may want to tune them.
diff --git a/icefall/rnn_lm/model.py b/icefall/rnn_lm/model.py
index 3598a4857..08eb753b5 100644
--- a/icefall/rnn_lm/model.py
+++ b/icefall/rnn_lm/model.py
@@ -153,9 +153,24 @@ class RnnLmModel(torch.nn.Module):
def clean_cache(self):
self.cache = {}
- def score_token(self, tokens: torch.Tensor, state=None):
+ def score_token(self, x: torch.Tensor, x_lens: torch.Tensor, state=None):
+ """Score a batch of tokens
+
+ Args:
+ x (torch.Tensor):
+ A batch of tokens
+ x_lens (torch.Tensor):
+ The length of tokens in the batch before padding
+ state (_type_, optional):
+ Either None or a tuple of two torch.Tensor. Each tensor has
+ the shape of (hidden_dim)
+
+
+ Returns:
+ _type_: _description_
+ """
device = next(self.parameters()).device
- batch_size = tokens.size(0)
+ batch_size = x.size(0)
if state:
h, c = state
else:
@@ -166,7 +181,7 @@ class RnnLmModel(torch.nn.Module):
device
)
- embedding = self.input_embedding(tokens)
+ embedding = self.input_embedding(x)
rnn_out, states = self.rnn(embedding, (h, c))
logits = self.output_linear(rnn_out)
diff --git a/icefall/rnn_lm/train.py b/icefall/rnn_lm/train.py
index 803da99d6..f43e66cd2 100755
--- a/icefall/rnn_lm/train.py
+++ b/icefall/rnn_lm/train.py
@@ -531,6 +531,9 @@ def run(rank, world_size, args):
tie_weights=params.tie_weights,
)
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
checkpoints = load_checkpoint_if_available(params=params, model=model)
model.to(device)
diff --git a/icefall/transformer_lm/attention.py b/icefall/transformer_lm/attention.py
new file mode 100644
index 000000000..5ce83b15e
--- /dev/null
+++ b/icefall/transformer_lm/attention.py
@@ -0,0 +1,510 @@
+# Copyright (c) 2021 University of Chinese Academy of Sciences (author: Han Zhu)
+#
+# 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.
+
+import warnings
+from typing import List, Optional, Tuple
+
+import torch
+from torch import Tensor, nn
+
+from icefall.transformer_lm.scaling import (
+ ActivationBalancer,
+ BasicNorm,
+ DoubleSwish,
+ ScaledConv1d,
+ ScaledConv2d,
+ ScaledLinear,
+)
+from icefall.utils import is_jit_tracing
+
+
+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,
+ ) -> 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 = ScaledLinear(embed_dim, 3 * embed_dim, bias=True)
+ self.out_proj = ScaledLinear(
+ embed_dim, embed_dim, bias=True, initial_scale=0.25
+ )
+
+ # linear transformation for positional encoding.
+ self.linear_pos = ScaledLinear(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.pos_bias_u_scale = nn.Parameter(torch.zeros(()).detach())
+ self.pos_bias_v_scale = nn.Parameter(torch.zeros(()).detach())
+ self._reset_parameters()
+
+ def _pos_bias_u(self):
+ return self.pos_bias_u * self.pos_bias_u_scale.exp()
+
+ def _pos_bias_v(self):
+ return self.pos_bias_v * self.pos_bias_v_scale.exp()
+
+ def _reset_parameters(self) -> None:
+ nn.init.normal_(self.pos_bias_u, std=0.01)
+ nn.init.normal_(self.pos_bias_v, std=0.01)
+
+ def forward(
+ self,
+ query: Tensor,
+ key: Tensor,
+ value: Tensor,
+ pos_emb: Tensor,
+ key_padding_mask: Optional[Tensor] = None,
+ need_weights: bool = False,
+ attn_mask: Optional[Tensor] = None,
+ left_context: int = 0,
+ ) -> 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.
+ left_context (int): left context (in frames) used during streaming decoding.
+ this is used only in real streaming decoding, in other circumstances,
+ it MUST be 0.
+
+ 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.get_weight(),
+ self.in_proj.get_bias(),
+ self.dropout,
+ self.out_proj.get_weight(),
+ self.out_proj.get_bias(),
+ training=self.training,
+ key_padding_mask=key_padding_mask,
+ need_weights=need_weights,
+ attn_mask=attn_mask,
+ left_context=left_context,
+ )
+
+ def rel_shift(self, x: Tensor, left_context: int = 0) -> Tensor:
+ """Compute relative positional encoding.
+
+ Args:
+ x: Input tensor (batch, head, time1, 2*time1-1+left_context).
+ time1 means the length of query vector.
+ left_context (int): left context (in frames) used during streaming decoding.
+ this is used only in real streaming decoding, in other circumstances,
+ it MUST be 0.
+
+ 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
+
+ time2 = time1 + left_context
+ if not is_jit_tracing():
+ assert (
+ n == left_context + 2 * time1 - 1
+ ), f"{n} == {left_context} + 2 * {time1} - 1"
+
+ if is_jit_tracing():
+ rows = torch.arange(start=time1 - 1, end=-1, step=-1)
+ cols = torch.arange(time2)
+ rows = rows.repeat(batch_size * num_heads).unsqueeze(-1)
+ indexes = rows + cols
+
+ x = x.reshape(-1, n)
+ x = torch.gather(x, dim=1, index=indexes)
+ x = x.reshape(batch_size, num_heads, time1, time2)
+ return x
+ else:
+ # 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, time2),
+ (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 = False,
+ attn_mask: Optional[Tensor] = None,
+ left_context: int = 0,
+ ) -> 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.
+ left_context (int): left context (in frames) used during streaming decoding.
+ this is used only in real streaming decoding, in other circumstances,
+ it MUST be 0.
+
+ 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()
+ if not is_jit_tracing():
+ 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
+ if not is_jit_tracing():
+ 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 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 * scaling).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 and not is_jit_tracing():
+ 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)
+ if not is_jit_tracing():
+ 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)
+ # (batch, 2*time1, head, d_k) --> (batch, head, d_k, 2*time -1)
+ p = p.permute(0, 2, 3, 1)
+
+ 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) # (batch, head, time1, 2*time1-1)
+ matrix_bd = self.rel_shift(matrix_bd, left_context)
+
+ attn_output_weights = matrix_ac + matrix_bd # (batch, head, time1, time2)
+
+ attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len, -1)
+
+ if not is_jit_tracing():
+ 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)
+
+ # If we are using dynamic_chunk_training and setting a limited
+ # num_left_chunks, the attention may only see the padding values which
+ # will also be masked out by `key_padding_mask`, at this circumstances,
+ # the whole column of `attn_output_weights` will be `-inf`
+ # (i.e. be `nan` after softmax), so, we fill `0.0` at the masking
+ # positions to avoid invalid loss value below.
+ if (
+ attn_mask is not None
+ and attn_mask.dtype == torch.bool
+ and key_padding_mask is not None
+ ):
+ if attn_mask.size(0) != 1:
+ attn_mask = attn_mask.view(bsz, num_heads, tgt_len, src_len)
+ combined_mask = attn_mask | key_padding_mask.unsqueeze(1).unsqueeze(2)
+ else:
+ # attn_mask.shape == (1, tgt_len, src_len)
+ combined_mask = attn_mask.unsqueeze(0) | key_padding_mask.unsqueeze(
+ 1
+ ).unsqueeze(2)
+
+ attn_output_weights = attn_output_weights.view(
+ bsz, num_heads, tgt_len, src_len
+ )
+ attn_output_weights = attn_output_weights.masked_fill(combined_mask, 0.0)
+ attn_output_weights = attn_output_weights.view(
+ bsz * num_heads, tgt_len, src_len
+ )
+
+ attn_output_weights = nn.functional.dropout(
+ attn_output_weights, p=dropout_p, training=training
+ )
+
+ attn_output = torch.bmm(attn_output_weights, v)
+
+ if not is_jit_tracing():
+ 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
diff --git a/icefall/transformer_lm/compute_perplexity.py b/icefall/transformer_lm/compute_perplexity.py
new file mode 100644
index 000000000..72d7c477b
--- /dev/null
+++ b/icefall/transformer_lm/compute_perplexity.py
@@ -0,0 +1,195 @@
+#!/usr/bin/env python3
+# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang
+# Xiaoyu Yang)
+#
+# 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.
+
+import argparse
+import logging
+import math
+from pathlib import Path
+
+import torch
+from dataset import get_dataloader
+from train import get_params
+
+from icefall.checkpoint import average_checkpoints, load_checkpoint
+from icefall.transformer_lm.model import TransformerLM
+from icefall.utils import AttributeDict, setup_logger, str2bool
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=7,
+ 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'. ",
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="transformer_lm/exp_full_libri_16layer_maxlen200_8gpu",
+ )
+
+ parser.add_argument(
+ "--lm-data",
+ type=str,
+ help="Path to the LM test data for computing perplexity",
+ default="transformer_lm/libri_lm_training_bpe500/sorted_lm_data-test.pt",
+ )
+
+ parser.add_argument(
+ "--vocab-size",
+ type=int,
+ default=500,
+ help="Vocabulary size of the model",
+ )
+
+ parser.add_argument(
+ "--num-layers",
+ type=int,
+ default=16,
+ help="Number of RNN layers the model",
+ )
+
+ parser.add_argument(
+ "--tie-weights",
+ type=str2bool,
+ default=False,
+ help="""True to share the weights between the input embedding layer and the
+ last output linear layer
+ """,
+ )
+
+ parser.add_argument(
+ "--batch-size",
+ type=int,
+ default=50,
+ help="Number of RNN layers the model",
+ )
+
+ parser.add_argument(
+ "--max-sent-len",
+ type=int,
+ default=100,
+ help="Number of RNN layers the model",
+ )
+
+ return parser
+
+
+def main():
+ parser = get_parser()
+ args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
+ args.lm_data = Path(args.lm_data)
+
+ params = get_params()
+ params.update(vars(args))
+
+ setup_logger(f"{params.exp_dir}/log-ppl/")
+ logging.info("Computing perplexity started")
+ logging.info(params)
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"Device: {device}")
+
+ logging.info("About to create model")
+ model = TransformerLM(
+ vocab_size=params.vocab_size,
+ d_model=params.encoder_dim,
+ embedding_dim=params.embedding_dim,
+ dim_feedforward=params.dim_feedforward,
+ nhead=params.nhead,
+ num_layers=params.num_layers,
+ tie_weights=params.tie_weights,
+ params=params,
+ )
+
+ if params.avg == 1:
+ load_checkpoint(f"{params.exp_dir}/epoch-{params.epoch}.pt", model)
+ model.to(device)
+ 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.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+
+ model.eval()
+ num_param = sum([p.numel() for p in model.parameters()])
+ num_param_requires_grad = sum(
+ [p.numel() for p in model.parameters() if p.requires_grad]
+ )
+
+ logging.info(f"Number of model parameters: {num_param}")
+ logging.info(
+ f"Number of model parameters (requires_grad): "
+ f"{num_param_requires_grad} "
+ f"({num_param_requires_grad/num_param_requires_grad*100}%)"
+ )
+
+ logging.info(f"Loading LM test data from {params.lm_data}")
+ test_dl = get_dataloader(
+ filename=params.lm_data,
+ is_distributed=False,
+ params=params,
+ )
+
+ tot_loss = 0.0
+ num_tokens = 0
+ num_sentences = 0
+ for batch_idx, batch in enumerate(test_dl):
+ x, y, sentence_lengths = batch
+ x = x.to(device)
+ y = y.to(device)
+ sentence_lengths = sentence_lengths.to(device)
+
+ nll = model(x, y, sentence_lengths)
+ loss = nll.sum().cpu().item()
+
+ tot_loss += loss
+ num_tokens += sentence_lengths.sum().cpu().item()
+ num_sentences += x.size(0)
+
+ ppl = math.exp(tot_loss / num_tokens)
+ logging.info(
+ f"total nll: {tot_loss}, num tokens: {num_tokens}, "
+ f"num sentences: {num_sentences}, ppl: {ppl:.3f}"
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/icefall/transformer_lm/dataset.py b/icefall/transformer_lm/dataset.py
new file mode 120000
index 000000000..5792a6cf0
--- /dev/null
+++ b/icefall/transformer_lm/dataset.py
@@ -0,0 +1 @@
+../rnn_lm/dataset.py
\ No newline at end of file
diff --git a/icefall/transformer_lm/encoder.py b/icefall/transformer_lm/encoder.py
new file mode 100644
index 000000000..4357b83d7
--- /dev/null
+++ b/icefall/transformer_lm/encoder.py
@@ -0,0 +1,329 @@
+# Copyright (c) 2021 Xiaomi Corporation (authors: Xiaoyu Yang)
+#
+# 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.
+
+import copy
+import math
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from icefall.transformer_lm.attention import RelPositionMultiheadAttention
+from icefall.transformer_lm.scaling import (
+ ActivationBalancer,
+ BasicNorm,
+ DoubleSwish,
+ ScaledConv1d,
+ ScaledConv2d,
+ ScaledLinear,
+)
+from icefall.utils import is_jit_tracing, make_pad_mask
+
+
+class Transformer(torch.nn.Module):
+ """_summary_
+
+ Args:
+ input_dim (int): Input feature dimension
+ d_mode (int): The dimension of the transformer
+ dim_feedforward (int ): The dimension of the ffw module
+ nhead (int): The number of attention heads
+ dropout_rate (float): dropout rate
+ att_dropout (float): dropout rate in attention module
+ """
+
+ def __init__(
+ self,
+ input_dim: int,
+ d_model: int,
+ dim_feedforward: int,
+ nhead: int = 4,
+ num_layers: int = 6,
+ dropout_rate: float = 0.1,
+ att_dropout: float = 0.0,
+ ):
+ super().__init__()
+
+ self.encoder_layers = num_layers
+ self.d_model = d_model
+
+ self.embed = ScaledLinear(input_dim, d_model)
+ self.norm_before = BasicNorm(d_model, learn_eps=False)
+
+ self.encoder_pos = RelPositionalEncoding(d_model, dropout_rate)
+
+ encoder_layer = TransformerEncoderLayer(
+ d_model=d_model,
+ dim_feedforward=dim_feedforward,
+ nhead=nhead,
+ dropout_rate=dropout_rate,
+ )
+
+ self.encoder = TransformerEncoder(encoder_layer, num_layers)
+
+ def _create_attention_mask(self, x_lens: torch.Tensor):
+ # create a 2D attention mask to mask out
+ # the upper right half of the attention matrix
+ max_len = max(x_lens)
+ ones = torch.ones(max_len, max_len, device=x_lens.device, dtype=torch.bool)
+ return torch.triu(ones, diagonal=1)
+
+ def forward(
+ self, x: torch.Tensor, x_lens: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Transformer forward
+
+ Args:
+ x (torch.Tensor): Input tensor (B,T,input_dim)
+ x_lens (torch.Tensor): The length of input tensors before padding (B,)
+
+ Returns:
+ Return a tuple of 2 tensors:
+ - x: output feature of the transformer (B,T,d_model)
+ - x_lens: output feature lens of the transformer
+ """
+
+ attention_mask = self._create_attention_mask(x_lens)
+ src_key_padding_mask = make_pad_mask(x_lens)
+
+ x = self.norm_before(self.embed(x))
+
+ x, pos_emb = self.encoder_pos(x)
+ x = x.permute(1, 0, 2)
+
+ x = self.encoder(
+ x,
+ pos_emb,
+ mask=attention_mask, # pass the attention mast
+ src_key_padding_mask=src_key_padding_mask,
+ ) # (T, N, C)
+
+ x = x.permute(1, 0, 2) # (T, N, C) ->(N, T, C)
+ return x, x_lens
+
+
+class TransformerEncoder(torch.nn.Module):
+ def __init__(self, encoder_layer: torch.nn.Module, num_layers: int) -> None:
+ """TransformerEncoder is a stack of N encoder layers
+
+ Args:
+ encoder_layer (torch.nn.Module): an instance of the TransformerEncoderLayer()
+ num_layers (int): Number of layers to be stacked
+ """
+ super().__init__()
+ self.layers = nn.ModuleList(
+ [copy.deepcopy(encoder_layer) for i in range(num_layers)]
+ )
+ self.num_layers = num_layers
+
+ def forward(
+ self,
+ src: torch.Tensor,
+ pos_emb: torch.Tensor,
+ src_key_padding_mask: Optional[torch.Tensor] = None,
+ mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """_summary_
+
+ 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).
+
+ Returns:
+ output: transformer encoded features
+ """
+ output = src
+
+ for layer_index, mod in enumerate(self.layers):
+ output = mod(
+ output,
+ pos_emb,
+ src_key_padding_mask=src_key_padding_mask,
+ src_mask=mask,
+ )
+
+ return output
+
+
+class TransformerEncoderLayer(torch.nn.Module):
+ def __init__(
+ self,
+ d_model: int,
+ dim_feedforward: int,
+ nhead: int,
+ dropout_rate: float,
+ ):
+ """TransformerEncoderLayer is made up of self-attn and feedforward module
+
+ Args:
+ d_model (int): The model size
+ dim_feedforward (int): Dimension of ffw module
+ nhead (int): Number of heads
+ dropout_rate (float): Dropout rate
+ """
+ super().__init__()
+
+ self.d_model = d_model
+
+ self.self_attn = RelPositionMultiheadAttention(d_model, nhead, dropout=0.0)
+ self.feed_forward = nn.Sequential(
+ ScaledLinear(d_model, dim_feedforward),
+ ActivationBalancer(channel_dim=-1),
+ DoubleSwish(),
+ nn.Dropout(dropout_rate),
+ ScaledLinear(dim_feedforward, d_model, initial_scale=0.25),
+ )
+
+ self.norm_final = BasicNorm(d_model)
+
+ self.balancer = ActivationBalancer(
+ channel_dim=-1, min_positive=0.45, max_positive=0.55, max_abs=6.0
+ )
+
+ self.dropout = nn.Dropout(dropout_rate)
+
+ def forward(
+ self,
+ src: torch.Tensor,
+ pos_emb: torch.Tensor,
+ src_key_padding_mask: Optional[torch.Tensor] = None,
+ src_mask: Optional[torch.Tensor] = None,
+ cache=None,
+ ):
+ """
+ Pass the input through the encoder layer.
+
+ Args:
+ src: the sequence to the encoder layer (required).
+ pos_emb: Positional embedding tensor (required).
+ src_key_padding_mask: the mask for the src keys per batch (optional).
+ src_mask: the mask for the src sequence (optional).
+ """
+ src_orig = 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 = src + self.dropout(src_att)
+
+ # feed forward module
+ src = src + self.dropout(self.feed_forward(src))
+
+ src = self.norm_final(self.balancer(src))
+
+ return src
+
+
+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__()
+ if is_jit_tracing():
+ # 10k frames correspond to ~100k ms, e.g., 100 seconds, i.e.,
+ # It assumes that the maximum input won't have more than
+ # 10k frames.
+ #
+ # TODO(fangjun): Use torch.jit.script() for this module
+ max_len = 10000
+
+ self.d_model = 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: torch.Tensor, left_context: int = 0) -> None:
+ """Reset the positional encodings."""
+ x_size_1 = x.size(1) + left_context
+ 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 vector 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[torch.Tensor, torch.Tensor]:
+ """Add positional encoding.
+
+ Args:
+ x (torch.Tensor): Input tensor (batch, time, `*`).
+ left_context (int): left context (in frames) used during streaming decoding.
+ this is used only in real streaming decoding, in other circumstances,
+ it MUST be 0.
+
+ Returns:
+ torch.Tensor: Encoded tensor (batch, time, `*`).
+ torch.Tensor: Encoded tensor (batch, 2*time-1, `*`).
+
+ """
+ self.extend_pe(x, left_context)
+ x_size_1 = x.size(1) + left_context
+ pos_emb = self.pe[
+ :,
+ self.pe.size(1) // 2
+ - x_size_1
+ + 1 : self.pe.size(1) // 2 # noqa E203
+ + x.size(1),
+ ]
+ return self.dropout(x), self.dropout(pos_emb)
diff --git a/icefall/transformer_lm/export.py b/icefall/transformer_lm/export.py
new file mode 100644
index 000000000..c08982e37
--- /dev/null
+++ b/icefall/transformer_lm/export.py
@@ -0,0 +1,186 @@
+#!/usr/bin/env python3
+# Copyright (c) 2022 Xiaomi Corporation (authors: Xiaoyu Yang)
+#
+# 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.
+
+# This script converts several saved checkpoints
+# to a single one using model averaging.
+
+import argparse
+import logging
+from pathlib import Path
+
+import torch
+from model import TransformerLM
+
+from icefall.checkpoint import load_checkpoint
+from icefall.utils import AttributeDict, load_averaged_model, str2bool
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=11,
+ 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'. ",
+ )
+
+ parser.add_argument(
+ "--vocab-size",
+ type=int,
+ default=500,
+ help="Vocabulary size of the model",
+ )
+
+ parser.add_argument(
+ "--embedding-dim",
+ type=int,
+ default=768,
+ help="Embedding dim of the model",
+ )
+
+ parser.add_argument(
+ "--encoder-dim",
+ type=int,
+ default=768,
+ help="Encoder dim of the model",
+ )
+
+ parser.add_argument(
+ "--dim_feedforward",
+ type=int,
+ default=2048,
+ help="Hidden dim of the model",
+ )
+
+ parser.add_argument(
+ "--nhead",
+ type=int,
+ default=8,
+ help="Number of attention heads",
+ )
+
+ parser.add_argument(
+ "--num-layers",
+ type=int,
+ default=16,
+ help="Number of Transformer layers",
+ )
+
+ parser.add_argument(
+ "--tie-weights",
+ type=str2bool,
+ default=True,
+ help="""True to share the weights between the input embedding layer and the
+ last output linear layer
+ """,
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="rnn_lm/exp",
+ help="""It specifies the directory where all training related
+ files, e.g., checkpoints, log, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--jit",
+ type=str2bool,
+ default=True,
+ help="""True to save a model after applying torch.jit.script.
+ """,
+ )
+
+ return parser
+
+
+def main():
+ args = get_parser().parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ params = AttributeDict({})
+ params.update(vars(args))
+
+ logging.info(params)
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"device: {device}")
+
+ logging.info("About to create model")
+ model = TransformerLM(
+ vocab_size=params.vocab_size,
+ d_model=params.encoder_dim,
+ embedding_dim=params.embedding_dim,
+ dim_feedforward=params.dim_feedforward,
+ nhead=params.nhead,
+ num_layers=params.num_layers,
+ tie_weights=params.tie_weights,
+ params=params,
+ )
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ model.to(device)
+
+ if params.avg == 1:
+ load_checkpoint(f"{params.exp_dir}/epoch-{params.epoch}.pt", model)
+ else:
+ model = load_averaged_model(
+ params.exp_dir, model, params.epoch, params.avg, device
+ )
+
+ model.to("cpu")
+ model.eval()
+
+ if params.jit:
+ logging.info("Using torch.jit.script")
+ model = torch.jit.script(model)
+ filename = params.exp_dir / "cpu_jit.pt"
+ model.save(str(filename))
+ logging.info(f"Saved to {filename}")
+ else:
+ logging.info("Not using torch.jit.script")
+ # Save it using a format so that it can be loaded
+ # by :func:`load_checkpoint`
+ filename = params.exp_dir / "pretrained.pt"
+ torch.save({"model": model.state_dict()}, str(filename))
+ logging.info(f"Saved to {filename}")
+
+
+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/icefall/transformer_lm/model.py b/icefall/transformer_lm/model.py
new file mode 100644
index 000000000..79dda3168
--- /dev/null
+++ b/icefall/transformer_lm/model.py
@@ -0,0 +1,115 @@
+# Copyright (c) 2022 Xiaomi Corporation (authors: Xiaoyu Yang)
+#
+# 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.
+
+import logging
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+from icefall.transformer_lm.encoder import Transformer
+from icefall.utils import AttributeDict, add_eos, add_sos, make_pad_mask
+
+
+class TransformerLM(torch.nn.Module):
+ def __init__(
+ self,
+ vocab_size: int,
+ embedding_dim: int,
+ d_model: int,
+ dim_feedforward: int,
+ nhead: int = 8,
+ num_layers: int = 16,
+ tie_weights: bool = True,
+ dropout: float = 0.1,
+ emb_dropout_rate: float = 0.0,
+ params: AttributeDict = None,
+ ):
+ super().__init__()
+
+ self.vocab_size = vocab_size
+ self.params = params
+
+ self.input_embedding = torch.nn.Embedding(
+ num_embeddings=vocab_size,
+ embedding_dim=embedding_dim,
+ )
+
+ self.encoder = Transformer(
+ input_dim=embedding_dim,
+ d_model=d_model,
+ dim_feedforward=dim_feedforward,
+ nhead=nhead,
+ num_layers=num_layers,
+ dropout_rate=dropout,
+ )
+
+ self.output_linear = torch.nn.Linear(
+ in_features=d_model, out_features=vocab_size
+ )
+ if tie_weights:
+ logging.info("Tying weights")
+ assert d_model == embedding_dim, (d_model, embedding_dim)
+ self.output_linear.weight = self.input_embedding.weight
+ else:
+ logging.info("Not tying weights")
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ y: torch.Tensor,
+ x_lens: torch.Tensor,
+ return_logits: bool = False,
+ ):
+ """Forward transformer language model
+
+ Args:
+ x (torch.Tensor): Input tokens (B,L)
+ y (torch.Tensor): Output tokens (with EOS appended) (B,L)
+ x_lens (torch.Tensor): Length of input tokens before padding (B,)
+ return_logits (bool, optional): Return logits instead of NLL
+
+ """
+
+ x = self.input_embedding(x)
+
+ x, x_lens = self.encoder(x, x_lens)
+
+ logits = self.output_linear(x)
+
+ if return_logits:
+ return logits
+
+ nll_loss = F.cross_entropy(
+ logits.reshape(-1, self.vocab_size), y.reshape(-1), reduction="none"
+ )
+
+ mask = make_pad_mask(x_lens).reshape(-1)
+ nll_loss.masked_fill_(mask, 0)
+
+ return nll_loss
+
+ def score_token(self, x: torch.Tensor, x_lens: torch.Tensor, state=None):
+
+ bs = x.size(0)
+
+ state = None
+ logits = self.forward(x, x, x_lens, return_logits=True)
+ index = torch.arange(bs)
+
+ last_logits = logits[index, x_lens - 1, :]
+
+ return last_logits.log_softmax(-1), state
diff --git a/icefall/transformer_lm/scaling.py b/icefall/transformer_lm/scaling.py
new file mode 120000
index 000000000..0876c0704
--- /dev/null
+++ b/icefall/transformer_lm/scaling.py
@@ -0,0 +1 @@
+../../egs/librispeech/ASR/pruned_transducer_stateless2/scaling.py
\ No newline at end of file
diff --git a/icefall/transformer_lm/train.py b/icefall/transformer_lm/train.py
new file mode 100644
index 000000000..c36abfcdf
--- /dev/null
+++ b/icefall/transformer_lm/train.py
@@ -0,0 +1,609 @@
+#!/usr/bin/env python3
+# Copyright 2021 Xiaomi Corp. (authors: Xiaoyu Yang)
+#
+# 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.
+
+
+"""
+Usage:
+ ./transformer_lm/train.py \
+ --start-epoch 0 \
+ --world-size 2 \
+ --num-epochs 1 \
+ --use-fp16 0 \
+ --num-layers 12 \
+ --batch-size 400
+
+"""
+
+import argparse
+import logging
+import math
+from pathlib import Path
+from shutil import copyfile
+from typing import Optional, Tuple
+
+import torch
+import torch.multiprocessing as mp
+import torch.nn as nn
+import torch.optim as optim
+from dataset import get_dataloader
+from lhotse.utils import fix_random_seed
+from model import TransformerLM
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.nn.utils import clip_grad_norm_
+from torch.utils.tensorboard import SummaryWriter
+
+from icefall.checkpoint import load_checkpoint
+from icefall.checkpoint import save_checkpoint as save_checkpoint_impl
+from icefall.dist import cleanup_dist, setup_dist
+from icefall.env import get_env_info
+from icefall.utils import AttributeDict, MetricsTracker, 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.",
+ )
+
+ parser.add_argument(
+ "--num-epochs",
+ type=int,
+ default=30,
+ help="Number of epochs to train.",
+ )
+
+ parser.add_argument(
+ "--start-epoch",
+ type=int,
+ default=0,
+ help="""Resume training from from this epoch.
+ If it is positive, it will load checkpoint from
+ exp_dir/epoch-{start_epoch-1}.pt
+ """,
+ )
+
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="transformer_lm/exp",
+ help="""The experiment dir.
+ It specifies the directory where all training related
+ files, e.g., checkpoints, logs, etc, are saved
+ """,
+ )
+
+ parser.add_argument(
+ "--use-fp16",
+ type=str2bool,
+ default=True,
+ help="Whether to use half precision training.",
+ )
+
+ parser.add_argument(
+ "--batch-size",
+ type=int,
+ default=400,
+ )
+
+ parser.add_argument(
+ "--lm-data",
+ type=str,
+ default="data/lm_training_bpe_500/sorted_lm_data.pt",
+ help="LM training data",
+ )
+
+ parser.add_argument(
+ "--lm-data-valid",
+ type=str,
+ default="data/lm_training_bpe_500/sorted_lm_data-valid.pt",
+ help="LM validation data",
+ )
+
+ parser.add_argument(
+ "--vocab-size",
+ type=int,
+ default=500,
+ help="Vocabulary size of the model",
+ )
+
+ parser.add_argument(
+ "--num-layers",
+ type=int,
+ default=12,
+ help="Number of Transformer layers in the model",
+ )
+
+ parser.add_argument(
+ "--tie-weights",
+ type=str2bool,
+ default=True,
+ help="""True to share the weights between the input embedding layer and the
+ last output linear layer
+ """,
+ )
+
+ parser.add_argument(
+ "--seed",
+ type=int,
+ default=42,
+ help="The seed for random generators intended for reproducibility",
+ )
+
+ return parser
+
+
+def get_params() -> AttributeDict:
+ """Return a dict containing training parameters."""
+
+ params = AttributeDict(
+ {
+ "max_sent_len": 200,
+ "sos_id": 1,
+ "eos_id": 1,
+ "blank_id": 0,
+ "lr": 1e-3,
+ "weight_decay": 1e-6,
+ "best_train_loss": float("inf"),
+ "best_valid_loss": float("inf"),
+ "best_train_epoch": -1,
+ "best_valid_epoch": -1,
+ "batch_idx_train": 0,
+ "log_interval": 200,
+ "reset_interval": 2000,
+ "valid_interval": 1000,
+ "nhead": 8,
+ "embedding_dim": 768,
+ "encoder_dim": 768,
+ "dim_feedforward": 2048,
+ "dropout": 0.1,
+ "env_info": get_env_info(),
+ }
+ )
+ 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"
+ logging.info(f"Loading checkpoint: {filename}")
+ 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(
+ model: nn.Module,
+ x: torch.Tensor,
+ y: torch.Tensor,
+ sentence_lengths: torch.Tensor,
+ is_training: bool,
+) -> Tuple[torch.Tensor, MetricsTracker]:
+ """Compute the negative log-likelihood loss given a model and its input.
+ Args:
+ model:
+ The NN model,
+ x:
+ A 2-D tensor. Each row contains BPE token IDs for a sentence. Also,
+ each row starts with SOS ID.
+ y:
+ A 2-D tensor. Each row is a shifted version of the corresponding row
+ in `x` but ends with an EOS ID (before padding).
+ sentence_lengths:
+ A 1-D tensor containing number of tokens of each sentence
+ before padding.
+ is_training:
+ True for training. False for validation.
+ """
+ with torch.set_grad_enabled(is_training):
+ device = model.device
+ x = x.to(device)
+ y = y.to(device)
+ sentence_lengths = sentence_lengths.to(device)
+
+ nll = model(x, y, sentence_lengths)
+ loss = nll.sum()
+
+ num_tokens = sentence_lengths.sum().item()
+
+ loss_info = MetricsTracker()
+ # Note: Due to how MetricsTracker() is designed,
+ # we use "frames" instead of "num_tokens" as a key here
+ loss_info["frames"] = num_tokens
+ loss_info["loss"] = loss.detach().item()
+ return loss, loss_info
+
+
+def compute_validation_loss(
+ params: AttributeDict,
+ model: nn.Module,
+ valid_dl: torch.utils.data.DataLoader,
+ world_size: int = 1,
+) -> MetricsTracker:
+ """Run the validation process. The validation loss
+ is saved in `params.valid_loss`.
+ """
+ model.eval()
+
+ tot_loss = MetricsTracker()
+
+ for batch_idx, batch in enumerate(valid_dl):
+ x, y, sentence_lengths = batch
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, loss_info = compute_loss(
+ model=model,
+ x=x,
+ y=y,
+ sentence_lengths=sentence_lengths,
+ is_training=False,
+ )
+
+ assert loss.requires_grad is False
+ tot_loss = tot_loss + loss_info
+
+ if world_size > 1:
+ tot_loss.reduce(loss.device)
+
+ loss_value = tot_loss["loss"] / tot_loss["frames"]
+ if loss_value < params.best_valid_loss:
+ params.best_valid_epoch = params.cur_epoch
+ params.best_valid_loss = loss_value
+
+ return tot_loss
+
+
+def train_one_epoch(
+ params: AttributeDict,
+ model: nn.Module,
+ optimizer: torch.optim.Optimizer,
+ 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 sentences 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.
+ 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 = MetricsTracker()
+
+ for batch_idx, batch in enumerate(train_dl):
+ params.batch_idx_train += 1
+ x, y, sentence_lengths = batch
+ batch_size = x.size(0)
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, loss_info = compute_loss(
+ model=model,
+ x=x,
+ y=y,
+ sentence_lengths=sentence_lengths,
+ is_training=True,
+ )
+
+ # summary stats
+ tot_loss = (tot_loss * (1 - 1 / params.reset_interval)) + loss_info
+
+ optimizer.zero_grad()
+ loss.backward()
+ clip_grad_norm_(model.parameters(), 5.0, 2.0)
+ optimizer.step()
+
+ if batch_idx % params.log_interval == 0:
+ # Note: "frames" here means "num_tokens"
+ this_batch_ppl = math.exp(loss_info["loss"] / loss_info["frames"])
+ tot_ppl = math.exp(tot_loss["loss"] / tot_loss["frames"])
+
+ logging.info(
+ f"Epoch {params.cur_epoch}, "
+ f"batch {batch_idx}, loss[{loss_info}, ppl: {this_batch_ppl}] "
+ f"tot_loss[{tot_loss}, ppl: {tot_ppl}], "
+ f"batch size: {batch_size}"
+ )
+
+ if tb_writer is not None:
+ loss_info.write_summary(
+ tb_writer, "train/current_", params.batch_idx_train
+ )
+ tot_loss.write_summary(tb_writer, "train/tot_", params.batch_idx_train)
+
+ tb_writer.add_scalar(
+ "train/current_ppl", this_batch_ppl, params.batch_idx_train
+ )
+
+ tb_writer.add_scalar("train/tot_ppl", tot_ppl, params.batch_idx_train)
+
+ if batch_idx > 0 and batch_idx % params.valid_interval == 0:
+ logging.info("Computing validation loss")
+
+ valid_info = compute_validation_loss(
+ params=params,
+ model=model,
+ valid_dl=valid_dl,
+ world_size=world_size,
+ )
+ model.train()
+
+ valid_ppl = math.exp(valid_info["loss"] / valid_info["frames"])
+ logging.info(
+ f"Epoch {params.cur_epoch}, validation: {valid_info}, "
+ f"ppl: {valid_ppl}"
+ )
+
+ if tb_writer is not None:
+ valid_info.write_summary(
+ tb_writer, "train/valid_", params.batch_idx_train
+ )
+
+ tb_writer.add_scalar(
+ "train/valid_ppl", valid_ppl, params.batch_idx_train
+ )
+
+ loss_value = tot_loss["loss"] / tot_loss["frames"]
+ params.train_loss = loss_value
+ 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))
+ is_distributed = world_size > 1
+
+ fix_random_seed(params.seed)
+ if is_distributed:
+ 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
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", rank)
+
+ logging.info(f"Device: {device}")
+
+ logging.info("About to create model")
+ model = TransformerLM(
+ vocab_size=params.vocab_size,
+ d_model=params.encoder_dim,
+ embedding_dim=params.embedding_dim,
+ dim_feedforward=params.dim_feedforward,
+ nhead=params.nhead,
+ num_layers=params.num_layers,
+ tie_weights=params.tie_weights,
+ params=params,
+ )
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ checkpoints = load_checkpoint_if_available(params=params, model=model)
+
+ model.to(device)
+ if is_distributed:
+ model = DDP(model, device_ids=[rank])
+
+ model.device = device
+
+ optimizer = optim.Adam(
+ model.parameters(),
+ lr=params.lr,
+ weight_decay=params.weight_decay,
+ )
+ if checkpoints:
+ logging.info("Load optimizer state_dict from checkpoint")
+ optimizer.load_state_dict(checkpoints["optimizer"])
+
+ logging.info(f"Loading LM training data from {params.lm_data}")
+ train_dl = get_dataloader(
+ filename=params.lm_data,
+ is_distributed=is_distributed,
+ params=params,
+ )
+
+ logging.info(f"Loading LM validation data from {params.lm_data_valid}")
+ valid_dl = get_dataloader(
+ filename=params.lm_data_valid,
+ is_distributed=is_distributed,
+ params=params,
+ )
+
+ # Note: No learning rate scheduler is used here
+ for epoch in range(params.start_epoch, params.num_epochs):
+ if is_distributed:
+ train_dl.sampler.set_epoch(epoch)
+
+ params.cur_epoch = epoch
+
+ train_one_epoch(
+ params=params,
+ model=model,
+ optimizer=optimizer,
+ 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 is_distributed:
+ torch.distributed.barrier()
+ cleanup_dist()
+
+
+def main():
+ parser = get_parser()
+ args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ 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)
+
+
+torch.set_num_threads(1)
+torch.set_num_interop_threads(1)
+
+if __name__ == "__main__":
+ main()
diff --git a/test/test_lexicon.py b/test/test_lexicon.py
index 69867efc7..b1beab3f6 100755
--- a/test/test_lexicon.py
+++ b/test/test_lexicon.py
@@ -112,7 +112,7 @@ def uniq_lexicon_test():
# But there is no word "ca" in the lexicon, so our
# implementation returns the id of "