=
+; # 言語音と独立に講演者の笑いが生じている場合
+<笑> =
+; # 言語音と独立に講演者の咳が生じている場合
+<咳> =
+; # 言語音と独立に講演者の息が生じている場合
+<息> =
+; # 講演者の泣き声
+<泣> =
+; # 聴衆（司会者なども含む）の発話
+<フロア発話> =
+; # 聴衆の笑い
+<フロア笑> =
+; # 聴衆の拍手
+<拍手> =
+; # 講演者が発表中に用いたデモンストレーションの音声
+<デモ> =
+; # 学会講演に発表時間を知らせるためにならすベルの音
+<ベル> =
+; # 転記単位全体が再度読み直された場合
+<朗読間違い> =
+; # 上記以外の音で特に目立った音
+<雑音> =
+; # 0.2秒以上のポーズ
+ =
+; # Redacted information, for R
+; # It is \x00D7 multiplication sign, not your normal 'x'
+× = ×
+
+[FIELDS]
+; # Time information for segment
+time = 3
+; # Word surface
+surface = 5
+; # Word surface root form without CSJ tags
+notag = 9
+; # Part Of Speech
+pos1 = 11
+; # Conjugated Form
+cForm = 12
+; # Conjugation Type
+cType1 = 13
+; # Subcategory of POS
+pos2 = 14
+; # Euphonic Change / Subcategory of Conjugation Type
+cType2 = 15
+; # Other information
+other = 16
+; # Pronunciation for lexicon
+pron = 10
+; # Speaker ID
+spk_id = 2
+
+[KATAKANA2ROMAJI]
+ア = 'a
+イ = 'i
+ウ = 'u
+エ = 'e
+オ = 'o
+カ = ka
+キ = ki
+ク = ku
+ケ = ke
+コ = ko
+ガ = ga
+ギ = gi
+グ = gu
+ゲ = ge
+ゴ = go
+サ = sa
+シ = si
+ス = su
+セ = se
+ソ = so
+ザ = za
+ジ = zi
+ズ = zu
+ゼ = ze
+ゾ = zo
+タ = ta
+チ = ti
+ツ = tu
+テ = te
+ト = to
+ダ = da
+ヂ = di
+ヅ = du
+デ = de
+ド = do
+ナ = na
+ニ = ni
+ヌ = nu
+ネ = ne
+ノ = no
+ハ = ha
+ヒ = hi
+フ = hu
+ヘ = he
+ホ = ho
+バ = ba
+ビ = bi
+ブ = bu
+ベ = be
+ボ = bo
+パ = pa
+ピ = pi
+プ = pu
+ペ = pe
+ポ = po
+マ = ma
+ミ = mi
+ム = mu
+メ = me
+モ = mo
+ヤ = ya
+ユ = yu
+ヨ = yo
+ラ = ra
+リ = ri
+ル = ru
+レ = re
+ロ = ro
+ワ = wa
+ヰ = we
+ヱ = wi
+ヲ = wo
+ン = ŋ
+ッ = q
+ー = -
+キャ = kǐa
+キュ = kǐu
+キョ = kǐo
+ギャ = gǐa
+ギュ = gǐu
+ギョ = gǐo
+シャ = sǐa
+シュ = sǐu
+ショ = sǐo
+ジャ = zǐa
+ジュ = zǐu
+ジョ = zǐo
+チャ = tǐa
+チュ = tǐu
+チョ = tǐo
+ヂャ = dǐa
+ヂュ = dǐu
+ヂョ = dǐo
+ニャ = nǐa
+ニュ = nǐu
+ニョ = nǐo
+ヒャ = hǐa
+ヒュ = hǐu
+ヒョ = hǐo
+ビャ = bǐa
+ビュ = bǐu
+ビョ = bǐo
+ピャ = pǐa
+ピュ = pǐu
+ピョ = pǐo
+ミャ = mǐa
+ミュ = mǐu
+ミョ = mǐo
+リャ = rǐa
+リュ = rǐu
+リョ = rǐo
+ァ = a
+ィ = i
+ゥ = u
+ェ = e
+ォ = o
+ヮ = ʍ
+ヴ = vu
+ャ = ǐa
+ュ = ǐu
+ョ = ǐo
+
diff --git a/egs/csj/ASR/local/conf/fluent.ini b/egs/csj/ASR/local/conf/fluent.ini
new file mode 100644
index 000000000..5d22f9eb8
--- /dev/null
+++ b/egs/csj/ASR/local/conf/fluent.ini
@@ -0,0 +1,321 @@
+; # This section is ignored if this file is not supplied as the first config file to
+; # lhotse prepare csj
+[SEGMENTS]
+; # Allowed period of nonverbal noise. If exceeded, a new segment is created.
+gap = 0.5
+; # Maximum length of segment (s).
+maxlen = 10
+; # Minimum length of segment (s). Segments shorter than `minlen` will be dropped silently.
+minlen = 0.02
+; # Use this symbol to represent a period of allowed nonverbal noise, i.e. `gap`.
+; # Pass an empty string to avoid adding any symbol. It was "" in kaldi.
+; # If you intend to use a multicharacter string for gap_sym, remember to register the
+; # multicharacter string as part of userdef-string in prepare_lang_char.py.
+gap_sym =
+
+[CONSTANTS]
+; # Name of this mode
+MODE = fluent
+; # Suffixes to use after the word surface (no longer used)
+MORPH = pos1 cForm cType2 pos2
+; # Used to differentiate between A tag and A_num tag
+JPN_NUM = ゼロ ０ 零 一 二 三 四 五 六 七 八 九 十 百 千 ．
+; # Dummy character to delineate multiline words
+PLUS = ＋
+
+[DECISIONS]
+; # TAG+'^'とは、タグが一つの転記単位に独立していない場合
+; # The PLUS (fullwidth) sign '＋' marks line boundaries for multiline entries
+
+; # フィラー、感情表出系感動詞
+; # 0 to remain, 1 to delete
+; # Example: '(F ぎょっ)'
+F = 1
+; # Example: '(L (F ン))', '比べ(F えー)る'
+F^ = 1
+; # 言い直し、いいよどみなどによる語断片
+; # 0 to remain, 1 to delete
+; # Example: '(D だ)(D だいが) 大学の学部の会議'
+D = 1
+; # Example: '(L (D ドゥ)＋(D ヒ))'
+D^ = 1
+; # 助詞、助動詞、接辞の言い直し
+; # 0 to remain, 1 to delete
+; # Example: '西洋 (D2 的)(F えー)(D ふ) 風というか'
+D2 = 1
+; # Example: '(X (D2 ノ))'
+D2^ = 1
+; # 聞き取りや語彙の判断に自信がない場合
+; # 0 to remain, 1 to delete
+; # Example: (? 字数) の
+; # If no option: empty string is returned regardless of output
+; # Example: '(?) で'
+? = 0
+; # Example: '(D (? すー))＋そう＋です＋よ＋ね'
+?^ = 0
+; # タグ?で、値は複数の候補が想定される場合
+; # 0 for main guess with matching morph info, 1 for second guess
+; # Example: '(? 次数, 実数)', '(? これ,ここで)＋(? 説明＋し＋た＋方＋が＋いい＋か＋な)'
+?, = 0
+; # Example: '(W (? テユクー);(? ケッキョク,テユウコトデ))', '(W マシ;(? マシ＋タ,マス))'
+?,^ = 0
+; # 音や言葉に関するメタ的な引用
+; # 0 to remain, 1 to delete
+; # Example: '助詞の (M は) は (M は) と書くが発音は (M わ)'
+M = 0
+; # Example: '(L (M ヒ)＋(M ヒ))', '(L (M (? ヒ＋ヒ)))'
+M^ = 0
+; # 外国語や古語、方言など
+; # 0 to remain, 1 to delete
+; # Example: '(O ザッツファイン)'
+O = 0
+; # Example: '(笑 (O エクスキューズ＋ミー))', '(笑 メダッ＋テ＋(O ナンボ))'
+O^ = 0
+; # 講演者の名前、差別語、誹謗中傷など
+; # 0 to remain, 1 to delete
+; # Example: '国語研の (R ××) です'
+R = 0
+R^ = 0
+; # 非朗読対象発話（朗読における言い間違い等）
+; # 0 to remain, 1 to delete
+; # Example: '(X 実際は) 実際には'
+X = 0
+; # Example: '(L (X (D2 ニ)))'
+X^ = 0
+; # アルファベットや算用数字、記号の表記
+; # 0 to use Japanese form, 1 to use alphabet form
+; # Example: '(A シーディーアール;ＣＤ－Ｒ)'
+A = 1
+; # Example: 'スモール(A エヌ;Ｎ)', 'ラージ(A キュー;Ｑ)', '(A ティーエフ;ＴＦ)＋(A アイディーエフ;ＩＤＦ)' (Strung together by pron: '(W (? ティーワイド);ティーエフ＋アイディーエフ)')
+A^ = 1
+; # タグAで、単語は算用数字の場合
+; # 0 to use Japanese form, 1 to use Arabic numerals
+; # Example: (A 二千;２０００)
+A_num = eval:self.notag
+A_num^ = eval:self.notag
+; # 何らかの原因で漢字表記できなくなった場合
+; # 0 to use broken form, 1 to use orthodox form
+; # Example: '(K たち (F えー) ばな;橘)'
+K = 1
+; # Example: '合(K か(?)く;格)', '宮(K ま(?)え;前)'
+K^ = 1
+; # 転訛、発音の怠けなど、一時的な発音エラー
+; # 0 to use wrong form, 1 to use orthodox form
+; # Example: '(W ギーツ;ギジュツ)'
+W = 1
+; # Example: '(F (W エド;エト))', 'イベント(W リレーティッド;リレーテッド)'
+W^ = 1
+; # 語の読みに関する知識レベルのいい間違い
+; # 0 to use wrong form, 1 to use orthodox form
+; # Example: '(B シブタイ;ジュータイ)'
+B = 0
+; # Example: 'データー(B カズ;スー)'
+B^ = 0
+; # 笑いながら発話
+; # 0 to remain, 1 to delete
+; # Example: '(笑 ナニガ)', '(笑 (F エー)＋ソー＋イッ＋タ＋ヨー＋ナ)'
+笑 = 0
+; # Example: 'コク(笑 サイ＋(D オン))',
+笑^ = 0
+; # 泣きながら発話
+; # 0 to remain, 1 to delete
+; # Example: '(泣 ドンナニ)'
+泣 = 0
+泣^ = 0
+; # 咳をしながら発話
+; # 0 to remain, 1 to delete
+; # Example: 'シャ(咳 リン) ノ'
+咳 = 0
+; # Example: 'イッ(咳 パン)', 'ワズ(咳 カ)'
+咳^ = 0
+; # ささやき声や独り言などの小さな声
+; # 0 to remain, 1 to delete
+; # Example: '(L アレコレナンダッケ)', '(L (W コデ;(? コレ,ココデ))＋(? セツメー＋シ＋タ＋ホー＋ガ＋イー＋カ＋ナ))'
+L = 0
+; # Example: 'デ(L ス)', 'ッ(L テ＋コ)ト'
+L^ = 0
+
+[REPLACEMENTS]
+; # ボーカルフライなどで母音が同定できない場合
+ =
+; # 「うん/うーん/ふーん」の音の特定が困難な場合
+ =
+; # 非語彙的な母音の引き延ばし
+ =
+; # 非語彙的な子音の引き延ばし
+ =
+; # 言語音と独立に講演者の笑いが生じている場合
+<笑> =
+; # 言語音と独立に講演者の咳が生じている場合
+<咳> =
+; # 言語音と独立に講演者の息が生じている場合
+<息> =
+; # 講演者の泣き声
+<泣> =
+; # 聴衆（司会者なども含む）の発話
+<フロア発話> =
+; # 聴衆の笑い
+<フロア笑> =
+; # 聴衆の拍手
+<拍手> =
+; # 講演者が発表中に用いたデモンストレーションの音声
+<デモ> =
+; # 学会講演に発表時間を知らせるためにならすベルの音
+<ベル> =
+; # 転記単位全体が再度読み直された場合
+<朗読間違い> =
+; # 上記以外の音で特に目立った音
+<雑音> =
+; # 0.2秒以上のポーズ
+ =
+; # Redacted information, for R
+; # It is \x00D7 multiplication sign, not your normal 'x'
+× = ×
+
+[FIELDS]
+; # Time information for segment
+time = 3
+; # Word surface
+surface = 5
+; # Word surface root form without CSJ tags
+notag = 9
+; # Part Of Speech
+pos1 = 11
+; # Conjugated Form
+cForm = 12
+; # Conjugation Type
+cType1 = 13
+; # Subcategory of POS
+pos2 = 14
+; # Euphonic Change / Subcategory of Conjugation Type
+cType2 = 15
+; # Other information
+other = 16
+; # Pronunciation for lexicon
+pron = 10
+; # Speaker ID
+spk_id = 2
+
+[KATAKANA2ROMAJI]
+ア = 'a
+イ = 'i
+ウ = 'u
+エ = 'e
+オ = 'o
+カ = ka
+キ = ki
+ク = ku
+ケ = ke
+コ = ko
+ガ = ga
+ギ = gi
+グ = gu
+ゲ = ge
+ゴ = go
+サ = sa
+シ = si
+ス = su
+セ = se
+ソ = so
+ザ = za
+ジ = zi
+ズ = zu
+ゼ = ze
+ゾ = zo
+タ = ta
+チ = ti
+ツ = tu
+テ = te
+ト = to
+ダ = da
+ヂ = di
+ヅ = du
+デ = de
+ド = do
+ナ = na
+ニ = ni
+ヌ = nu
+ネ = ne
+ノ = no
+ハ = ha
+ヒ = hi
+フ = hu
+ヘ = he
+ホ = ho
+バ = ba
+ビ = bi
+ブ = bu
+ベ = be
+ボ = bo
+パ = pa
+ピ = pi
+プ = pu
+ペ = pe
+ポ = po
+マ = ma
+ミ = mi
+ム = mu
+メ = me
+モ = mo
+ヤ = ya
+ユ = yu
+ヨ = yo
+ラ = ra
+リ = ri
+ル = ru
+レ = re
+ロ = ro
+ワ = wa
+ヰ = we
+ヱ = wi
+ヲ = wo
+ン = ŋ
+ッ = q
+ー = -
+キャ = kǐa
+キュ = kǐu
+キョ = kǐo
+ギャ = gǐa
+ギュ = gǐu
+ギョ = gǐo
+シャ = sǐa
+シュ = sǐu
+ショ = sǐo
+ジャ = zǐa
+ジュ = zǐu
+ジョ = zǐo
+チャ = tǐa
+チュ = tǐu
+チョ = tǐo
+ヂャ = dǐa
+ヂュ = dǐu
+ヂョ = dǐo
+ニャ = nǐa
+ニュ = nǐu
+ニョ = nǐo
+ヒャ = hǐa
+ヒュ = hǐu
+ヒョ = hǐo
+ビャ = bǐa
+ビュ = bǐu
+ビョ = bǐo
+ピャ = pǐa
+ピュ = pǐu
+ピョ = pǐo
+ミャ = mǐa
+ミュ = mǐu
+ミョ = mǐo
+リャ = rǐa
+リュ = rǐu
+リョ = rǐo
+ァ = a
+ィ = i
+ゥ = u
+ェ = e
+ォ = o
+ヮ = ʍ
+ヴ = vu
+ャ = ǐa
+ュ = ǐu
+ョ = ǐo
+
diff --git a/egs/csj/ASR/local/conf/number.ini b/egs/csj/ASR/local/conf/number.ini
new file mode 100644
index 000000000..2613c3409
--- /dev/null
+++ b/egs/csj/ASR/local/conf/number.ini
@@ -0,0 +1,321 @@
+; # This section is ignored if this file is not supplied as the first config file to
+; # lhotse prepare csj
+[SEGMENTS]
+; # Allowed period of nonverbal noise. If exceeded, a new segment is created.
+gap = 0.5
+; # Maximum length of segment (s).
+maxlen = 10
+; # Minimum length of segment (s). Segments shorter than `minlen` will be dropped silently.
+minlen = 0.02
+; # Use this symbol to represent a period of allowed nonverbal noise, i.e. `gap`.
+; # Pass an empty string to avoid adding any symbol. It was "" in kaldi.
+; # If you intend to use a multicharacter string for gap_sym, remember to register the
+; # multicharacter string as part of userdef-string in prepare_lang_char.py.
+gap_sym =
+
+[CONSTANTS]
+; # Name of this mode
+MODE = number
+; # Suffixes to use after the word surface (no longer used)
+MORPH = pos1 cForm cType2 pos2
+; # Used to differentiate between A tag and A_num tag
+JPN_NUM = ゼロ ０ 零 一 二 三 四 五 六 七 八 九 十 百 千 ．
+; # Dummy character to delineate multiline words
+PLUS = ＋
+
+[DECISIONS]
+; # TAG+'^'とは、タグが一つの転記単位に独立していない場合
+; # The PLUS (fullwidth) sign '＋' marks line boundaries for multiline entries
+
+; # フィラー、感情表出系感動詞
+; # 0 to remain, 1 to delete
+; # Example: '(F ぎょっ)'
+F = 1
+; # Example: '(L (F ン))', '比べ(F えー)る'
+F^ = 1
+; # 言い直し、いいよどみなどによる語断片
+; # 0 to remain, 1 to delete
+; # Example: '(D だ)(D だいが) 大学の学部の会議'
+D = 1
+; # Example: '(L (D ドゥ)＋(D ヒ))'
+D^ = 1
+; # 助詞、助動詞、接辞の言い直し
+; # 0 to remain, 1 to delete
+; # Example: '西洋 (D2 的)(F えー)(D ふ) 風というか'
+D2 = 1
+; # Example: '(X (D2 ノ))'
+D2^ = 1
+; # 聞き取りや語彙の判断に自信がない場合
+; # 0 to remain, 1 to delete
+; # Example: (? 字数) の
+; # If no option: empty string is returned regardless of output
+; # Example: '(?) で'
+? = 0
+; # Example: '(D (? すー))＋そう＋です＋よ＋ね'
+?^ = 0
+; # タグ?で、値は複数の候補が想定される場合
+; # 0 for main guess with matching morph info, 1 for second guess
+; # Example: '(? 次数, 実数)', '(? これ,ここで)＋(? 説明＋し＋た＋方＋が＋いい＋か＋な)'
+?, = 0
+; # Example: '(W (? テユクー);(? ケッキョク,テユウコトデ))', '(W マシ;(? マシ＋タ,マス))'
+?,^ = 0
+; # 音や言葉に関するメタ的な引用
+; # 0 to remain, 1 to delete
+; # Example: '助詞の (M は) は (M は) と書くが発音は (M わ)'
+M = 0
+; # Example: '(L (M ヒ)＋(M ヒ))', '(L (M (? ヒ＋ヒ)))'
+M^ = 0
+; # 外国語や古語、方言など
+; # 0 to remain, 1 to delete
+; # Example: '(O ザッツファイン)'
+O = 0
+; # Example: '(笑 (O エクスキューズ＋ミー))', '(笑 メダッ＋テ＋(O ナンボ))'
+O^ = 0
+; # 講演者の名前、差別語、誹謗中傷など
+; # 0 to remain, 1 to delete
+; # Example: '国語研の (R ××) です'
+R = 0
+R^ = 0
+; # 非朗読対象発話（朗読における言い間違い等）
+; # 0 to remain, 1 to delete
+; # Example: '(X 実際は) 実際には'
+X = 0
+; # Example: '(L (X (D2 ニ)))'
+X^ = 0
+; # アルファベットや算用数字、記号の表記
+; # 0 to use Japanese form, 1 to use alphabet form
+; # Example: '(A シーディーアール;ＣＤ－Ｒ)'
+A = 1
+; # Example: 'スモール(A エヌ;Ｎ)', 'ラージ(A キュー;Ｑ)', '(A ティーエフ;ＴＦ)＋(A アイディーエフ;ＩＤＦ)' (Strung together by pron: '(W (? ティーワイド);ティーエフ＋アイディーエフ)')
+A^ = 1
+; # タグAで、単語は算用数字の場合
+; # 0 to use Japanese form, 1 to use Arabic numerals
+; # Example: (A 二千;２０００)
+A_num = 1
+A_num^ = 1
+; # 何らかの原因で漢字表記できなくなった場合
+; # 0 to use broken form, 1 to use orthodox form
+; # Example: '(K たち (F えー) ばな;橘)'
+K = 1
+; # Example: '合(K か(?)く;格)', '宮(K ま(?)え;前)'
+K^ = 1
+; # 転訛、発音の怠けなど、一時的な発音エラー
+; # 0 to use wrong form, 1 to use orthodox form
+; # Example: '(W ギーツ;ギジュツ)'
+W = 1
+; # Example: '(F (W エド;エト))', 'イベント(W リレーティッド;リレーテッド)'
+W^ = 1
+; # 語の読みに関する知識レベルのいい間違い
+; # 0 to use wrong form, 1 to use orthodox form
+; # Example: '(B シブタイ;ジュータイ)'
+B = 0
+; # Example: 'データー(B カズ;スー)'
+B^ = 0
+; # 笑いながら発話
+; # 0 to remain, 1 to delete
+; # Example: '(笑 ナニガ)', '(笑 (F エー)＋ソー＋イッ＋タ＋ヨー＋ナ)'
+笑 = 0
+; # Example: 'コク(笑 サイ＋(D オン))',
+笑^ = 0
+; # 泣きながら発話
+; # 0 to remain, 1 to delete
+; # Example: '(泣 ドンナニ)'
+泣 = 0
+泣^ = 0
+; # 咳をしながら発話
+; # 0 to remain, 1 to delete
+; # Example: 'シャ(咳 リン) ノ'
+咳 = 0
+; # Example: 'イッ(咳 パン)', 'ワズ(咳 カ)'
+咳^ = 0
+; # ささやき声や独り言などの小さな声
+; # 0 to remain, 1 to delete
+; # Example: '(L アレコレナンダッケ)', '(L (W コデ;(? コレ,ココデ))＋(? セツメー＋シ＋タ＋ホー＋ガ＋イー＋カ＋ナ))'
+L = 0
+; # Example: 'デ(L ス)', 'ッ(L テ＋コ)ト'
+L^ = 0
+
+[REPLACEMENTS]
+; # ボーカルフライなどで母音が同定できない場合
+ =
+; # 「うん/うーん/ふーん」の音の特定が困難な場合
+ =
+; # 非語彙的な母音の引き延ばし
+ =
+; # 非語彙的な子音の引き延ばし
+ =
+; # 言語音と独立に講演者の笑いが生じている場合
+<笑> =
+; # 言語音と独立に講演者の咳が生じている場合
+<咳> =
+; # 言語音と独立に講演者の息が生じている場合
+<息> =
+; # 講演者の泣き声
+<泣> =
+; # 聴衆（司会者なども含む）の発話
+<フロア発話> =
+; # 聴衆の笑い
+<フロア笑> =
+; # 聴衆の拍手
+<拍手> =
+; # 講演者が発表中に用いたデモンストレーションの音声
+<デモ> =
+; # 学会講演に発表時間を知らせるためにならすベルの音
+<ベル> =
+; # 転記単位全体が再度読み直された場合
+<朗読間違い> =
+; # 上記以外の音で特に目立った音
+<雑音> =
+; # 0.2秒以上のポーズ
+ =
+; # Redacted information, for R
+; # It is \x00D7 multiplication sign, not your normal 'x'
+× = ×
+
+[FIELDS]
+; # Time information for segment
+time = 3
+; # Word surface
+surface = 5
+; # Word surface root form without CSJ tags
+notag = 9
+; # Part Of Speech
+pos1 = 11
+; # Conjugated Form
+cForm = 12
+; # Conjugation Type
+cType1 = 13
+; # Subcategory of POS
+pos2 = 14
+; # Euphonic Change / Subcategory of Conjugation Type
+cType2 = 15
+; # Other information
+other = 16
+; # Pronunciation for lexicon
+pron = 10
+; # Speaker ID
+spk_id = 2
+
+[KATAKANA2ROMAJI]
+ア = 'a
+イ = 'i
+ウ = 'u
+エ = 'e
+オ = 'o
+カ = ka
+キ = ki
+ク = ku
+ケ = ke
+コ = ko
+ガ = ga
+ギ = gi
+グ = gu
+ゲ = ge
+ゴ = go
+サ = sa
+シ = si
+ス = su
+セ = se
+ソ = so
+ザ = za
+ジ = zi
+ズ = zu
+ゼ = ze
+ゾ = zo
+タ = ta
+チ = ti
+ツ = tu
+テ = te
+ト = to
+ダ = da
+ヂ = di
+ヅ = du
+デ = de
+ド = do
+ナ = na
+ニ = ni
+ヌ = nu
+ネ = ne
+ノ = no
+ハ = ha
+ヒ = hi
+フ = hu
+ヘ = he
+ホ = ho
+バ = ba
+ビ = bi
+ブ = bu
+ベ = be
+ボ = bo
+パ = pa
+ピ = pi
+プ = pu
+ペ = pe
+ポ = po
+マ = ma
+ミ = mi
+ム = mu
+メ = me
+モ = mo
+ヤ = ya
+ユ = yu
+ヨ = yo
+ラ = ra
+リ = ri
+ル = ru
+レ = re
+ロ = ro
+ワ = wa
+ヰ = we
+ヱ = wi
+ヲ = wo
+ン = ŋ
+ッ = q
+ー = -
+キャ = kǐa
+キュ = kǐu
+キョ = kǐo
+ギャ = gǐa
+ギュ = gǐu
+ギョ = gǐo
+シャ = sǐa
+シュ = sǐu
+ショ = sǐo
+ジャ = zǐa
+ジュ = zǐu
+ジョ = zǐo
+チャ = tǐa
+チュ = tǐu
+チョ = tǐo
+ヂャ = dǐa
+ヂュ = dǐu
+ヂョ = dǐo
+ニャ = nǐa
+ニュ = nǐu
+ニョ = nǐo
+ヒャ = hǐa
+ヒュ = hǐu
+ヒョ = hǐo
+ビャ = bǐa
+ビュ = bǐu
+ビョ = bǐo
+ピャ = pǐa
+ピュ = pǐu
+ピョ = pǐo
+ミャ = mǐa
+ミュ = mǐu
+ミョ = mǐo
+リャ = rǐa
+リュ = rǐu
+リョ = rǐo
+ァ = a
+ィ = i
+ゥ = u
+ェ = e
+ォ = o
+ヮ = ʍ
+ヴ = vu
+ャ = ǐa
+ュ = ǐu
+ョ = ǐo
+
diff --git a/egs/csj/ASR/local/conf/symbol.ini b/egs/csj/ASR/local/conf/symbol.ini
new file mode 100644
index 000000000..8ba451dd5
--- /dev/null
+++ b/egs/csj/ASR/local/conf/symbol.ini
@@ -0,0 +1,322 @@
+; # This section is ignored if this file is not supplied as the first config file to
+; # lhotse prepare csj
+[SEGMENTS]
+; # Allowed period of nonverbal noise. If exceeded, a new segment is created.
+gap = 0.5
+; # Maximum length of segment (s).
+maxlen = 10
+; # Minimum length of segment (s). Segments shorter than `minlen` will be dropped silently.
+minlen = 0.02
+; # Use this symbol to represent a period of allowed nonverbal noise, i.e. `gap`.
+; # Pass an empty string to avoid adding any symbol. It was "" in kaldi.
+; # If you intend to use a multicharacter string for gap_sym, remember to register the
+; # multicharacter string as part of userdef-string in prepare_lang_char.py.
+gap_sym =
+
+[CONSTANTS]
+; # Name of this mode
+; # See https://www.isca-speech.org/archive/pdfs/interspeech_2022/horii22_interspeech.pdf
+MODE = symbol
+; # Suffixes to use after the word surface (no longer used)
+MORPH = pos1 cForm cType2 pos2
+; # Used to differentiate between A tag and A_num tag
+JPN_NUM = ゼロ ０ 零 一 二 三 四 五 六 七 八 九 十 百 千 ．
+; # Dummy character to delineate multiline words
+PLUS = ＋
+
+[DECISIONS]
+; # TAG+'^'とは、タグが一つの転記単位に独立していない場合
+; # The PLUS (fullwidth) sign '＋' marks line boundaries for multiline entries
+
+; # フィラー、感情表出系感動詞
+; # 0 to remain, 1 to delete
+; # Example: '(F ぎょっ)'
+F = ＃
+; # Example: '(L (F ン))', '比べ(F えー)る'
+F^ = ＃
+; # 言い直し、いいよどみなどによる語断片
+; # 0 to remain, 1 to delete
+; # Example: '(D だ)(D だいが) 大学の学部の会議'
+D = ＠
+; # Example: '(L (D ドゥ)＋(D ヒ))'
+D^ = ＠
+; # 助詞、助動詞、接辞の言い直し
+; # 0 to remain, 1 to delete
+; # Example: '西洋 (D2 的)(F えー)(D ふ) 風というか'
+D2 = ＠
+; # Example: '(X (D2 ノ))'
+D2^ = ＠
+; # 聞き取りや語彙の判断に自信がない場合
+; # 0 to remain, 1 to delete
+; # Example: (? 字数) の
+; # If no option: empty string is returned regardless of output
+; # Example: '(?) で'
+? = 0
+; # Example: '(D (? すー))＋そう＋です＋よ＋ね'
+?^ = 0
+; # タグ?で、値は複数の候補が想定される場合
+; # 0 for main guess with matching morph info, 1 for second guess
+; # Example: '(? 次数, 実数)', '(? これ,ここで)＋(? 説明＋し＋た＋方＋が＋いい＋か＋な)'
+?, = 0
+; # Example: '(W (? テユクー);(? ケッキョク,テユウコトデ))', '(W マシ;(? マシ＋タ,マス))'
+?,^ = 0
+; # 音や言葉に関するメタ的な引用
+; # 0 to remain, 1 to delete
+; # Example: '助詞の (M は) は (M は) と書くが発音は (M わ)'
+M = 0
+; # Example: '(L (M ヒ)＋(M ヒ))', '(L (M (? ヒ＋ヒ)))'
+M^ = 0
+; # 外国語や古語、方言など
+; # 0 to remain, 1 to delete
+; # Example: '(O ザッツファイン)'
+O = 0
+; # Example: '(笑 (O エクスキューズ＋ミー))', '(笑 メダッ＋テ＋(O ナンボ))'
+O^ = 0
+; # 講演者の名前、差別語、誹謗中傷など
+; # 0 to remain, 1 to delete
+; # Example: '国語研の (R ××) です'
+R = 0
+R^ = 0
+; # 非朗読対象発話（朗読における言い間違い等）
+; # 0 to remain, 1 to delete
+; # Example: '(X 実際は) 実際には'
+X = 0
+; # Example: '(L (X (D2 ニ)))'
+X^ = 0
+; # アルファベットや算用数字、記号の表記
+; # 0 to use Japanese form, 1 to use alphabet form
+; # Example: '(A シーディーアール;ＣＤ－Ｒ)'
+A = 1
+; # Example: 'スモール(A エヌ;Ｎ)', 'ラージ(A キュー;Ｑ)', '(A ティーエフ;ＴＦ)＋(A アイディーエフ;ＩＤＦ)' (Strung together by pron: '(W (? ティーワイド);ティーエフ＋アイディーエフ)')
+A^ = 1
+; # タグAで、単語は算用数字の場合
+; # 0 to use Japanese form, 1 to use Arabic numerals
+; # Example: (A 二千;２０００)
+A_num = eval:self.notag
+A_num^ = eval:self.notag
+; # 何らかの原因で漢字表記できなくなった場合
+; # 0 to use broken form, 1 to use orthodox form
+; # Example: '(K たち (F えー) ばな;橘)'
+K = 1
+; # Example: '合(K か(?)く;格)', '宮(K ま(?)え;前)'
+K^ = 1
+; # 転訛、発音の怠けなど、一時的な発音エラー
+; # 0 to use wrong form, 1 to use orthodox form
+; # Example: '(W ギーツ;ギジュツ)'
+W = 1
+; # Example: '(F (W エド;エト))', 'イベント(W リレーティッド;リレーテッド)'
+W^ = 1
+; # 語の読みに関する知識レベルのいい間違い
+; # 0 to use wrong form, 1 to use orthodox form
+; # Example: '(B シブタイ;ジュータイ)'
+B = 0
+; # Example: 'データー(B カズ;スー)'
+B^ = 0
+; # 笑いながら発話
+; # 0 to remain, 1 to delete
+; # Example: '(笑 ナニガ)', '(笑 (F エー)＋ソー＋イッ＋タ＋ヨー＋ナ)'
+笑 = 0
+; # Example: 'コク(笑 サイ＋(D オン))',
+笑^ = 0
+; # 泣きながら発話
+; # 0 to remain, 1 to delete
+; # Example: '(泣 ドンナニ)'
+泣 = 0
+泣^ = 0
+; # 咳をしながら発話
+; # 0 to remain, 1 to delete
+; # Example: 'シャ(咳 リン) ノ'
+咳 = 0
+; # Example: 'イッ(咳 パン)', 'ワズ(咳 カ)'
+咳^ = 0
+; # ささやき声や独り言などの小さな声
+; # 0 to remain, 1 to delete
+; # Example: '(L アレコレナンダッケ)', '(L (W コデ;(? コレ,ココデ))＋(? セツメー＋シ＋タ＋ホー＋ガ＋イー＋カ＋ナ))'
+L = 0
+; # Example: 'デ(L ス)', 'ッ(L テ＋コ)ト'
+L^ = 0
+
+[REPLACEMENTS]
+; # ボーカルフライなどで母音が同定できない場合
+ =
+; # 「うん/うーん/ふーん」の音の特定が困難な場合
+ =
+; # 非語彙的な母音の引き延ばし
+ =
+; # 非語彙的な子音の引き延ばし
+ =
+; # 言語音と独立に講演者の笑いが生じている場合
+<笑> =
+; # 言語音と独立に講演者の咳が生じている場合
+<咳> =
+; # 言語音と独立に講演者の息が生じている場合
+<息> =
+; # 講演者の泣き声
+<泣> =
+; # 聴衆（司会者なども含む）の発話
+<フロア発話> =
+; # 聴衆の笑い
+<フロア笑> =
+; # 聴衆の拍手
+<拍手> =
+; # 講演者が発表中に用いたデモンストレーションの音声
+<デモ> =
+; # 学会講演に発表時間を知らせるためにならすベルの音
+<ベル> =
+; # 転記単位全体が再度読み直された場合
+<朗読間違い> =
+; # 上記以外の音で特に目立った音
+<雑音> =
+; # 0.2秒以上のポーズ
+ =
+; # Redacted information, for R
+; # It is \x00D7 multiplication sign, not your normal 'x'
+× = ×
+
+[FIELDS]
+; # Time information for segment
+time = 3
+; # Word surface
+surface = 5
+; # Word surface root form without CSJ tags
+notag = 9
+; # Part Of Speech
+pos1 = 11
+; # Conjugated Form
+cForm = 12
+; # Conjugation Type
+cType1 = 13
+; # Subcategory of POS
+pos2 = 14
+; # Euphonic Change / Subcategory of Conjugation Type
+cType2 = 15
+; # Other information
+other = 16
+; # Pronunciation for lexicon
+pron = 10
+; # Speaker ID
+spk_id = 2
+
+[KATAKANA2ROMAJI]
+ア = 'a
+イ = 'i
+ウ = 'u
+エ = 'e
+オ = 'o
+カ = ka
+キ = ki
+ク = ku
+ケ = ke
+コ = ko
+ガ = ga
+ギ = gi
+グ = gu
+ゲ = ge
+ゴ = go
+サ = sa
+シ = si
+ス = su
+セ = se
+ソ = so
+ザ = za
+ジ = zi
+ズ = zu
+ゼ = ze
+ゾ = zo
+タ = ta
+チ = ti
+ツ = tu
+テ = te
+ト = to
+ダ = da
+ヂ = di
+ヅ = du
+デ = de
+ド = do
+ナ = na
+ニ = ni
+ヌ = nu
+ネ = ne
+ノ = no
+ハ = ha
+ヒ = hi
+フ = hu
+ヘ = he
+ホ = ho
+バ = ba
+ビ = bi
+ブ = bu
+ベ = be
+ボ = bo
+パ = pa
+ピ = pi
+プ = pu
+ペ = pe
+ポ = po
+マ = ma
+ミ = mi
+ム = mu
+メ = me
+モ = mo
+ヤ = ya
+ユ = yu
+ヨ = yo
+ラ = ra
+リ = ri
+ル = ru
+レ = re
+ロ = ro
+ワ = wa
+ヰ = we
+ヱ = wi
+ヲ = wo
+ン = ŋ
+ッ = q
+ー = -
+キャ = kǐa
+キュ = kǐu
+キョ = kǐo
+ギャ = gǐa
+ギュ = gǐu
+ギョ = gǐo
+シャ = sǐa
+シュ = sǐu
+ショ = sǐo
+ジャ = zǐa
+ジュ = zǐu
+ジョ = zǐo
+チャ = tǐa
+チュ = tǐu
+チョ = tǐo
+ヂャ = dǐa
+ヂュ = dǐu
+ヂョ = dǐo
+ニャ = nǐa
+ニュ = nǐu
+ニョ = nǐo
+ヒャ = hǐa
+ヒュ = hǐu
+ヒョ = hǐo
+ビャ = bǐa
+ビュ = bǐu
+ビョ = bǐo
+ピャ = pǐa
+ピュ = pǐu
+ピョ = pǐo
+ミャ = mǐa
+ミュ = mǐu
+ミョ = mǐo
+リャ = rǐa
+リュ = rǐu
+リョ = rǐo
+ァ = a
+ィ = i
+ゥ = u
+ェ = e
+ォ = o
+ヮ = ʍ
+ヴ = vu
+ャ = ǐa
+ュ = ǐu
+ョ = ǐo
+
diff --git a/egs/csj/ASR/local/display_manifest_statistics.py b/egs/csj/ASR/local/display_manifest_statistics.py
new file mode 100644
index 000000000..c9de21073
--- /dev/null
+++ b/egs/csj/ASR/local/display_manifest_statistics.py
@@ -0,0 +1,182 @@
+#!/usr/bin/env python3
+# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang)
+# 2022 The University of Electro-Communications (author: Teo Wen Shen) # noqa
+#
+# 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
+from pathlib import Path
+
+from lhotse import CutSet, load_manifest
+
+ARGPARSE_DESCRIPTION = """
+This file displays duration statistics of utterances in a manifest.
+You can use the displayed value to choose minimum/maximum duration
+to remove short and long utterances during the training.
+
+See the function `remove_short_and_long_utt()` in
+pruned_transducer_stateless5/train.py for usage.
+"""
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ description=ARGPARSE_DESCRIPTION,
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+ )
+
+ parser.add_argument(
+ "--manifest-dir", type=Path, help="Path to cutset manifests"
+ )
+
+ return parser.parse_args()
+
+
+def main():
+ args = get_parser()
+
+ for path in args.manifest_dir.glob("csj_cuts_*.jsonl.gz"):
+
+ cuts: CutSet = load_manifest(path)
+
+ print("\n---------------------------------\n")
+ print(path.name + ":")
+ cuts.describe()
+
+
+if __name__ == "__main__":
+ main()
+
+"""
+## eval1
+Cuts count: 1272
+Total duration (hh:mm:ss): 01:50:07
+Speech duration (hh:mm:ss): 01:50:07 (100.0%)
+Duration statistics (seconds):
+mean 5.2
+std 3.9
+min 0.2
+25% 1.9
+50% 4.0
+75% 8.1
+99% 14.3
+99.5% 14.7
+99.9% 16.0
+max 16.9
+Recordings available: 1272
+Features available: 1272
+Supervisions available: 1272
+SUPERVISION custom fields:
+- fluent (in 1272 cuts)
+- disfluent (in 1272 cuts)
+- number (in 1272 cuts)
+- symbol (in 1272 cuts)
+
+## eval2
+Cuts count: 1292
+Total duration (hh:mm:ss): 01:56:50
+Speech duration (hh:mm:ss): 01:56:50 (100.0%)
+Duration statistics (seconds):
+mean 5.4
+std 3.9
+min 0.1
+25% 2.1
+50% 4.6
+75% 8.6
+99% 14.1
+99.5% 15.2
+99.9% 16.1
+max 16.9
+Recordings available: 1292
+Features available: 1292
+Supervisions available: 1292
+SUPERVISION custom fields:
+- fluent (in 1292 cuts)
+- number (in 1292 cuts)
+- symbol (in 1292 cuts)
+- disfluent (in 1292 cuts)
+
+## eval3
+Cuts count: 1385
+Total duration (hh:mm:ss): 01:19:21
+Speech duration (hh:mm:ss): 01:19:21 (100.0%)
+Duration statistics (seconds):
+mean 3.4
+std 3.0
+min 0.2
+25% 1.2
+50% 2.5
+75% 4.6
+99% 12.7
+99.5% 13.7
+99.9% 15.0
+max 15.9
+Recordings available: 1385
+Features available: 1385
+Supervisions available: 1385
+SUPERVISION custom fields:
+- number (in 1385 cuts)
+- symbol (in 1385 cuts)
+- fluent (in 1385 cuts)
+- disfluent (in 1385 cuts)
+
+## valid
+Cuts count: 4000
+Total duration (hh:mm:ss): 05:08:09
+Speech duration (hh:mm:ss): 05:08:09 (100.0%)
+Duration statistics (seconds):
+mean 4.6
+std 3.8
+min 0.1
+25% 1.5
+50% 3.4
+75% 7.0
+99% 13.8
+99.5% 14.8
+99.9% 16.0
+max 17.3
+Recordings available: 4000
+Features available: 4000
+Supervisions available: 4000
+SUPERVISION custom fields:
+- fluent (in 4000 cuts)
+- symbol (in 4000 cuts)
+- disfluent (in 4000 cuts)
+- number (in 4000 cuts)
+
+## train
+Cuts count: 1291134
+Total duration (hh:mm:ss): 1596:37:27
+Speech duration (hh:mm:ss): 1596:37:27 (100.0%)
+Duration statistics (seconds):
+mean 4.5
+std 3.6
+min 0.0
+25% 1.6
+50% 3.3
+75% 6.4
+99% 14.0
+99.5% 14.8
+99.9% 16.6
+max 27.8
+Recordings available: 1291134
+Features available: 1291134
+Supervisions available: 1291134
+SUPERVISION custom fields:
+- disfluent (in 1291134 cuts)
+- fluent (in 1291134 cuts)
+- symbol (in 1291134 cuts)
+- number (in 1291134 cuts)
+"""
diff --git a/egs/csj/ASR/local/prepare_lang_char.py b/egs/csj/ASR/local/prepare_lang_char.py
new file mode 100644
index 000000000..e4d996871
--- /dev/null
+++ b/egs/csj/ASR/local/prepare_lang_char.py
@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+# Copyright 2022 The University of Electro-Communications (Author: Teo Wen Shen) # noqa
+#
+# 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 pathlib import Path
+
+from lhotse import CutSet
+
+ARGPARSE_DESCRIPTION = """
+This script gathers all training transcripts of the specified {trans_mode} type
+and produces a token_list that would be output set of the ASR system.
+
+It splits transcripts by whitespace into lists, then, for each word in the
+list, if the word does not appear in the list of user-defined multicharacter
+strings, it further splits that word into individual characters to be counted
+into the output token set.
+
+It outputs 4 files into the lang directory:
+- trans_mode: the name of transcript mode. If trans_mode was not specified,
+ this will be an empty file.
+- userdef_string: a list of user defined strings that should not be split
+ further into individual characters. By default, it contains "", "",
+ ""
+- words_len: the total number of tokens in the output set.
+- words.txt: a list of tokens in the output set. The length matches words_len.
+
+"""
+
+
+def get_args():
+ parser = argparse.ArgumentParser(
+ description=ARGPARSE_DESCRIPTION,
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+ )
+
+ parser.add_argument(
+ "--train-cut", type=Path, required=True, help="Path to the train cut"
+ )
+
+ parser.add_argument(
+ "--trans-mode",
+ type=str,
+ default=None,
+ help=(
+ "Name of the transcript mode to use. "
+ "If lang-dir is not set, this will also name the lang-dir"
+ ),
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=Path,
+ default=None,
+ help=(
+ "Name of lang dir. "
+ "If not set, this will default to lang_char_{trans-mode}"
+ ),
+ )
+
+ parser.add_argument(
+ "--userdef-string",
+ type=Path,
+ default=None,
+ help="Multicharacter strings that do not need to be split",
+ )
+
+ return parser.parse_args()
+
+
+def main():
+ args = get_args()
+
+ logging.basicConfig(
+ format=(
+ "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] " "%(message)s"
+ ),
+ level=logging.INFO,
+ )
+
+ if not args.lang_dir:
+ p = "lang_char"
+ if args.trans_mode:
+ p += f"_{args.trans_mode}"
+ args.lang_dir = Path(p)
+
+ if args.userdef_string:
+ args.userdef_string = set(args.userdef_string.read_text().split())
+ else:
+ args.userdef_string = set()
+
+ sysdef_string = ["", "", ""]
+ args.userdef_string.update(sysdef_string)
+
+ train_set: CutSet = CutSet.from_file(args.train_cut)
+
+ words = set()
+ logging.info(
+ f"Creating vocabulary from {args.train_cut.name}"
+ f" at {args.trans_mode} mode."
+ )
+ for cut in train_set:
+ try:
+ text: str = (
+ cut.supervisions[0].custom[args.trans_mode]
+ if args.trans_mode
+ else cut.supervisions[0].text
+ )
+ except KeyError:
+ raise KeyError(
+ f"Could not find {args.trans_mode} in "
+ f"{cut.supervisions[0].custom}"
+ )
+ for t in text.split():
+ if t in args.userdef_string:
+ words.add(t)
+ else:
+ words.update(c for c in list(t))
+
+ words -= set(sysdef_string)
+ words = sorted(words)
+ words = [""] + words + ["", ""]
+
+ args.lang_dir.mkdir(parents=True, exist_ok=True)
+ (args.lang_dir / "words.txt").write_text(
+ "\n".join(f"{word}\t{i}" for i, word in enumerate(words))
+ )
+
+ (args.lang_dir / "words_len").write_text(f"{len(words)}")
+
+ (args.lang_dir / "userdef_string").write_text(
+ "\n".join(args.userdef_string)
+ )
+
+ (args.lang_dir / "trans_mode").write_text(args.trans_mode)
+ logging.info("Done.")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/csj/ASR/local/validate_manifest.py b/egs/csj/ASR/local/validate_manifest.py
new file mode 100644
index 000000000..0c4c6c1ea
--- /dev/null
+++ b/egs/csj/ASR/local/validate_manifest.py
@@ -0,0 +1,98 @@
+#!/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 checks the following assumptions of the generated manifest:
+
+- Single supervision per cut
+- Supervision time bounds are within cut time bounds
+
+We will add more checks later if needed.
+
+Usage example:
+
+ python3 ./local/validate_manifest.py \
+ ./data/fbank/librispeech_cuts_train-clean-100.jsonl.gz
+
+"""
+
+import argparse
+import logging
+from pathlib import Path
+
+from lhotse import CutSet, load_manifest
+from lhotse.cut import Cut
+
+
+def get_args():
+ parser = argparse.ArgumentParser()
+
+ parser.add_argument(
+ "--manifest",
+ type=Path,
+ help="Path to the manifest file",
+ )
+
+ return parser.parse_args()
+
+
+def validate_one_supervision_per_cut(c: Cut):
+ if len(c.supervisions) != 1:
+ raise ValueError(f"{c.id} has {len(c.supervisions)} supervisions")
+
+
+def validate_supervision_and_cut_time_bounds(c: Cut):
+ s = c.supervisions[0]
+
+ # Removed because when the cuts were trimmed from supervisions,
+ # the start time of the supervision can be lesser than cut start time.
+ # https://github.com/lhotse-speech/lhotse/issues/813
+ # if s.start < c.start:
+ # raise ValueError(
+ # f"{c.id}: Supervision start time {s.start} is less "
+ # f"than cut start time {c.start}"
+ # )
+
+ if s.end > c.end:
+ raise ValueError(
+ f"{c.id}: Supervision end time {s.end} is larger "
+ f"than cut end time {c.end}"
+ )
+
+
+def main():
+ args = get_args()
+
+ manifest = Path(args.manifest)
+ logging.info(f"Validating {manifest}")
+
+ assert manifest.is_file(), f"{manifest} does not exist"
+ cut_set = load_manifest(manifest)
+ assert isinstance(cut_set, CutSet)
+
+ for c in cut_set:
+ validate_one_supervision_per_cut(c)
+ validate_supervision_and_cut_time_bounds(c)
+
+
+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/csj/ASR/prepare.sh b/egs/csj/ASR/prepare.sh
new file mode 100755
index 000000000..269c1ec9a
--- /dev/null
+++ b/egs/csj/ASR/prepare.sh
@@ -0,0 +1,130 @@
+#!/usr/bin/env bash
+# We assume the following directories are downloaded.
+#
+# - $csj_dir
+# CSJ is assumed to be the USB-type directory, which should contain the following subdirectories:-
+# - DATA (not used in this script)
+# - DOC (not used in this script)
+# - MODEL (not used in this script)
+# - MORPH
+# - LDB (not used in this script)
+# - SUWDIC (not used in this script)
+# - SDB
+# - core
+# - ...
+# - noncore
+# - ...
+# - PLABEL (not used in this script)
+# - SUMMARY (not used in this script)
+# - TOOL (not used in this script)
+# - WAV
+# - core
+# - ...
+# - noncore
+# - ...
+# - XML (not used in this script)
+#
+# - $musan_dir
+# This directory contains the following directories downloaded from
+# http://www.openslr.org/17/
+# - music
+# - noise
+# - speech
+#
+# By default, this script produces the original transcript like kaldi and espnet. Optionally, you
+# can generate other transcript formats by supplying your own config files. A few examples of these
+# config files can be found in local/conf.
+
+set -eou pipefail
+
+nj=8
+stage=-1
+stop_stage=100
+
+csj_dir=/mnt/minami_data_server/t2131178/corpus/CSJ
+musan_dir=/mnt/minami_data_server/t2131178/corpus/musan/musan
+trans_dir=$csj_dir/retranscript
+csj_fbank_dir=/mnt/host/csj_data/fbank
+musan_fbank_dir=$musan_dir/fbank
+csj_manifest_dir=data/manifests
+musan_manifest_dir=$musan_dir/manifests
+
+. shared/parse_options.sh || exit 1
+
+mkdir -p data
+
+log() {
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
+ log "Stage 1: Prepare CSJ manifest"
+ # If you want to generate more transcript modes, append the path to those config files at c.
+ # Example: lhotse prepare csj $csj_dir $trans_dir $csj_manifest_dir -c local/conf/disfluent.ini
+ # NOTE: In case multiple config files are supplied, the second config file and onwards will inherit
+ # the segment boundaries of the first config file.
+ if [ ! -e $csj_manifest_dir/.librispeech.done ]; then
+ lhotse prepare csj $csj_dir $trans_dir $csj_manifest_dir -j 4
+ touch $csj_manifest_dir/.librispeech.done
+ fi
+fi
+
+if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
+ log "Stage 2: Prepare musan manifest"
+ mkdir -p $musan_manifest_dir
+ if [ ! -e $musan_manifest_dir/.musan.done ]; then
+ lhotse prepare musan $musan_dir $musan_manifest_dir
+ touch $musan_manifest_dir/.musan.done
+ fi
+fi
+
+if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
+ log "Stage 3: Compute CSJ fbank"
+ if [ ! -e $csj_fbank_dir/.csj-validated.done ]; then
+ python local/compute_fbank_csj.py --manifest-dir $csj_manifest_dir \
+ --fbank-dir $csj_fbank_dir
+ parts=(
+ train
+ valid
+ eval1
+ eval2
+ eval3
+ )
+ for part in ${parts[@]}; do
+ python local/validate_manifest.py --manifest $csj_manifest_dir/csj_cuts_$part.jsonl.gz
+ done
+ touch $csj_fbank_dir/.csj-validated.done
+ fi
+fi
+
+if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
+ log "Stage 4: Prepare CSJ lang"
+ modes=disfluent
+
+ # If you want prepare the lang directory for other transcript modes, just append
+ # the names of those modes behind. An example is shown as below:-
+ # modes="$modes fluent symbol number"
+
+ for mode in ${modes[@]}; do
+ python local/prepare_lang_char.py --trans-mode $mode \
+ --train-cut $csj_manifest_dir/csj_cuts_train.jsonl.gz \
+ --lang-dir lang_char_$mode
+ done
+fi
+
+if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then
+ log "Stage 5: Compute fbank for musan"
+ mkdir -p $musan_fbank_dir
+
+ if [ ! -e $musan_fbank_dir/.musan.done ]; then
+ python local/compute_fbank_musan.py --manifest-dir $musan_manifest_dir --fbank-dir $musan_fbank_dir
+ touch $musan_fbank_dir/.musan.done
+ fi
+fi
+
+if [ $stage -le 6 ] && [ $stop_stage -ge 6 ]; then
+ log "Stage 6: Show manifest statistics"
+ python local/display_manifest_statistics.py --manifest-dir $csj_manifest_dir > $csj_manifest_dir/manifest_statistics.txt
+ cat $csj_manifest_dir/manifest_statistics.txt
+fi
\ No newline at end of file
diff --git a/egs/csj/ASR/shared b/egs/csj/ASR/shared
new file mode 120000
index 000000000..4c5e91438
--- /dev/null
+++ b/egs/csj/ASR/shared
@@ -0,0 +1 @@
+../../../icefall/shared/
\ No newline at end of file
diff --git a/egs/gigaspeech/ASR/conformer_ctc/asr_datamodule.py b/egs/gigaspeech/ASR/conformer_ctc/asr_datamodule.py
index ab958fa68..d78e26240 100644
--- a/egs/gigaspeech/ASR/conformer_ctc/asr_datamodule.py
+++ b/egs/gigaspeech/ASR/conformer_ctc/asr_datamodule.py
@@ -20,9 +20,8 @@ import logging
from functools import lru_cache
from pathlib import Path
-from lhotse import CutSet, Fbank, FbankConfig, load_manifest
+from lhotse import CutSet, Fbank, FbankConfig, load_manifest, load_manifest_lazy
from lhotse.dataset import (
- BucketingSampler,
CutConcatenate,
CutMix,
DynamicBucketingSampler,
@@ -191,7 +190,7 @@ class GigaSpeechAsrDataModule:
def train_dataloaders(self, cuts_train: CutSet) -> DataLoader:
logging.info("About to get Musan cuts")
cuts_musan = load_manifest(
- self.args.manifest_dir / "cuts_musan.json.gz"
+ self.args.manifest_dir / "musan_cuts.jsonl.gz"
)
transforms = []
@@ -315,7 +314,7 @@ class GigaSpeechAsrDataModule:
cut_transforms=transforms,
return_cuts=self.args.return_cuts,
)
- valid_sampler = BucketingSampler(
+ valid_sampler = DynamicBucketingSampler(
cuts_valid,
max_duration=self.args.max_duration,
shuffle=False,
@@ -339,8 +338,10 @@ class GigaSpeechAsrDataModule:
else PrecomputedFeatures(),
return_cuts=self.args.return_cuts,
)
- sampler = BucketingSampler(
- cuts, max_duration=self.args.max_duration, shuffle=False
+ sampler = DynamicBucketingSampler(
+ cuts,
+ max_duration=self.args.max_duration,
+ shuffle=False,
)
logging.debug("About to create test dataloader")
test_dl = DataLoader(
@@ -361,7 +362,9 @@ class GigaSpeechAsrDataModule:
@lru_cache()
def dev_cuts(self) -> CutSet:
logging.info("About to get dev cuts")
- cuts_valid = load_manifest(self.args.manifest_dir / "cuts_DEV.jsonl.gz")
+ cuts_valid = load_manifest_lazy(
+ self.args.manifest_dir / "cuts_DEV.jsonl.gz"
+ )
if self.args.small_dev:
return cuts_valid.subset(first=1000)
else:
@@ -370,4 +373,4 @@ class GigaSpeechAsrDataModule:
@lru_cache()
def test_cuts(self) -> CutSet:
logging.info("About to get test cuts")
- return load_manifest(self.args.manifest_dir / "cuts_TEST.jsonl.gz")
+ return load_manifest_lazy(self.args.manifest_dir / "cuts_TEST.jsonl.gz")
diff --git a/egs/gigaspeech/ASR/conformer_ctc/conformer.py b/egs/gigaspeech/ASR/conformer_ctc/conformer.py
index 871712a46..6fac07f93 100644
--- a/egs/gigaspeech/ASR/conformer_ctc/conformer.py
+++ b/egs/gigaspeech/ASR/conformer_ctc/conformer.py
@@ -253,7 +253,9 @@ class ConformerEncoderLayer(nn.Module):
residual = src
if self.normalize_before:
src = self.norm_conv(src)
- src = residual + self.dropout(self.conv_module(src))
+ src = residual + self.dropout(
+ self.conv_module(src, src_key_padding_mask=src_key_padding_mask)
+ )
if not self.normalize_before:
src = self.norm_conv(src)
@@ -369,7 +371,7 @@ class RelPositionalEncoding(torch.nn.Module):
):
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
+ # 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 Tensor:
+ 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).
Returns:
Tensor: Output tensor (#time, batch, channels).
@@ -908,6 +915,8 @@ class ConvolutionModule(nn.Module):
x = nn.functional.glu(x, dim=1) # (batch, channels, time)
# 1D Depthwise Conv
+ if src_key_padding_mask is not None:
+ x.masked_fill_(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
x = self.depthwise_conv(x)
if self.use_batchnorm:
x = self.norm(x)
diff --git a/egs/gigaspeech/ASR/conformer_ctc/decode.py b/egs/gigaspeech/ASR/conformer_ctc/decode.py
index 6ab9852b4..51406667e 100755
--- a/egs/gigaspeech/ASR/conformer_ctc/decode.py
+++ b/egs/gigaspeech/ASR/conformer_ctc/decode.py
@@ -173,13 +173,13 @@ def get_params() -> AttributeDict:
def post_processing(
- results: List[Tuple[List[str], List[str]]],
-) -> List[Tuple[List[str], List[str]]]:
+ results: List[Tuple[str, List[str], List[str]]],
+) -> List[Tuple[str, List[str], List[str]]]:
new_results = []
- for ref, hyp in results:
+ for key, ref, hyp in results:
new_ref = asr_text_post_processing(" ".join(ref)).split()
new_hyp = asr_text_post_processing(" ".join(hyp)).split()
- new_results.append((new_ref, new_hyp))
+ new_results.append((key, new_ref, new_hyp))
return new_results
@@ -408,7 +408,7 @@ def decode_dataset(
sos_id: int,
eos_id: int,
G: Optional[k2.Fsa] = None,
-) -> Dict[str, List[Tuple[List[str], List[str]]]]:
+) -> Dict[str, List[Tuple[str, List[str], List[str]]]]:
"""Decode dataset.
Args:
@@ -451,6 +451,7 @@ def decode_dataset(
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,
@@ -469,9 +470,9 @@ def decode_dataset(
for lm_scale, hyps in hyps_dict.items():
this_batch = []
assert len(hyps) == len(texts)
- for hyp_words, ref_text in zip(hyps, texts):
+ for cut_id, hyp_words, ref_text in zip(cut_ids, hyps, texts):
ref_words = ref_text.split()
- this_batch.append((ref_words, hyp_words))
+ this_batch.append((cut_id, ref_words, hyp_words))
results[lm_scale].extend(this_batch)
else:
@@ -501,7 +502,7 @@ def decode_dataset(
def save_results(
params: AttributeDict,
test_set_name: str,
- results_dict: Dict[str, List[Tuple[List[str], List[str]]]],
+ results_dict: Dict[str, List[Tuple[str, List[str], List[str]]]],
):
if params.method == "attention-decoder":
# Set it to False since there are too many logs.
@@ -512,6 +513,7 @@ def save_results(
for key, results in results_dict.items():
recog_path = params.exp_dir / f"recogs-{test_set_name}-{key}.txt"
results = post_processing(results)
+ results = sorted(results)
store_transcripts(filename=recog_path, texts=results)
if enable_log:
logging.info(f"The transcripts are stored in {recog_path}")
@@ -676,6 +678,8 @@ def main():
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
gigaspeech = GigaSpeechAsrDataModule(args)
dev_cuts = gigaspeech.dev_cuts()
diff --git a/egs/gigaspeech/ASR/local/compute_fbank_gigaspeech_dev_test.py b/egs/gigaspeech/ASR/local/compute_fbank_gigaspeech_dev_test.py
index 9f1039893..8209ee3ec 100755
--- a/egs/gigaspeech/ASR/local/compute_fbank_gigaspeech_dev_test.py
+++ b/egs/gigaspeech/ASR/local/compute_fbank_gigaspeech_dev_test.py
@@ -20,11 +20,7 @@ import logging
from pathlib import Path
import torch
-from lhotse import (
- CutSet,
- KaldifeatFbank,
- KaldifeatFbankConfig,
-)
+from lhotse import CutSet, KaldifeatFbank, KaldifeatFbankConfig
# Torch's multithreaded behavior needs to be disabled or
# it wastes a lot of CPU and slow things down.
@@ -69,6 +65,7 @@ def compute_fbank_gigaspeech_dev_test():
storage_path=f"{in_out_dir}/feats_{partition}",
num_workers=num_workers,
batch_duration=batch_duration,
+ overwrite=True,
)
cut_set = cut_set.trim_to_supervisions(
keep_overlapping=False, min_duration=None
diff --git a/egs/gigaspeech/ASR/local/compute_fbank_gigaspeech_splits.py b/egs/gigaspeech/ASR/local/compute_fbank_gigaspeech_splits.py
index 9dd3c046d..6410249db 100755
--- a/egs/gigaspeech/ASR/local/compute_fbank_gigaspeech_splits.py
+++ b/egs/gigaspeech/ASR/local/compute_fbank_gigaspeech_splits.py
@@ -22,11 +22,7 @@ from datetime import datetime
from pathlib import Path
import torch
-from lhotse import (
- CutSet,
- KaldifeatFbank,
- KaldifeatFbankConfig,
-)
+from lhotse import CutSet, KaldifeatFbank, KaldifeatFbankConfig
# Torch's multithreaded behavior needs to be disabled or
# it wastes a lot of CPU and slow things down.
@@ -120,6 +116,7 @@ def compute_fbank_gigaspeech_splits(args):
storage_path=f"{output_dir}/feats_XL_{idx}",
num_workers=args.num_workers,
batch_duration=args.batch_duration,
+ overwrite=True,
)
logging.info("About to split cuts into smaller chunks.")
diff --git a/egs/gigaspeech/ASR/local/compute_fbank_musan.py b/egs/gigaspeech/ASR/local/compute_fbank_musan.py
deleted file mode 100755
index 562872993..000000000
--- a/egs/gigaspeech/ASR/local/compute_fbank_musan.py
+++ /dev/null
@@ -1,103 +0,0 @@
-#!/usr/bin/env python3
-# Copyright 2021 Johns Hopkins University (Piotr Żelasko)
-# Copyright 2021 Xiaomi Corp. (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.
-
-import logging
-from pathlib import Path
-
-import torch
-from lhotse import (
- CutSet,
- KaldifeatFbank,
- KaldifeatFbankConfig,
- combine,
-)
-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)
-
-
-def compute_fbank_musan():
- src_dir = Path("data/manifests")
- output_dir = Path("data/fbank")
-
- # number of workers in dataloader
- num_workers = 10
-
- # number of seconds in a batch
- batch_duration = 600
-
- dataset_parts = (
- "music",
- "speech",
- "noise",
- )
-
- manifests = read_manifests_if_cached(
- prefix="musan", dataset_parts=dataset_parts, output_dir=src_dir
- )
- assert manifests is not None
-
- musan_cuts_path = output_dir / "cuts_musan.json.gz"
-
- if musan_cuts_path.is_file():
- logging.info(f"{musan_cuts_path} already exists - skipping")
- return
-
- logging.info("Extracting features for Musan")
-
- device = torch.device("cpu")
- if torch.cuda.is_available():
- device = torch.device("cuda", 0)
- extractor = KaldifeatFbank(KaldifeatFbankConfig(device=device))
-
- logging.info(f"device: {device}")
-
- musan_cuts = (
- CutSet.from_manifests(
- recordings=combine(
- part["recordings"] for part in manifests.values()
- )
- )
- .cut_into_windows(10.0)
- .filter(lambda c: c.duration > 5)
- .compute_and_store_features_batch(
- extractor=extractor,
- storage_path=f"{output_dir}/feats_musan",
- num_workers=num_workers,
- batch_duration=batch_duration,
- )
- )
- musan_cuts.to_json(musan_cuts_path)
-
-
-def main():
- formatter = (
- "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
- )
- logging.basicConfig(format=formatter, level=logging.INFO)
-
- compute_fbank_musan()
-
-
-if __name__ == "__main__":
- main()
diff --git a/egs/gigaspeech/ASR/local/compute_fbank_musan.py b/egs/gigaspeech/ASR/local/compute_fbank_musan.py
new file mode 120000
index 000000000..5833f2484
--- /dev/null
+++ b/egs/gigaspeech/ASR/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/gigaspeech/ASR/local/preprocess_gigaspeech.py b/egs/gigaspeech/ASR/local/preprocess_gigaspeech.py
index 0cec82ad5..48d10a157 100755
--- a/egs/gigaspeech/ASR/local/preprocess_gigaspeech.py
+++ b/egs/gigaspeech/ASR/local/preprocess_gigaspeech.py
@@ -62,6 +62,13 @@ def preprocess_giga_speech():
)
assert manifests is not None
+ assert len(manifests) == len(dataset_parts), (
+ len(manifests),
+ len(dataset_parts),
+ list(manifests.keys()),
+ dataset_parts,
+ )
+
for partition, m in manifests.items():
logging.info(f"Processing {partition}")
raw_cuts_path = output_dir / f"cuts_{partition}_raw.jsonl.gz"
diff --git a/egs/gigaspeech/ASR/pruned_transducer_stateless2/asr_datamodule.py b/egs/gigaspeech/ASR/pruned_transducer_stateless2/asr_datamodule.py
index ff3d3b07a..c87686e1e 100644
--- a/egs/gigaspeech/ASR/pruned_transducer_stateless2/asr_datamodule.py
+++ b/egs/gigaspeech/ASR/pruned_transducer_stateless2/asr_datamodule.py
@@ -23,9 +23,8 @@ from pathlib import Path
from typing import Any, Dict, Optional
import torch
-from lhotse import CutSet, Fbank, FbankConfig, load_manifest
+from lhotse import CutSet, Fbank, FbankConfig, load_manifest, load_manifest_lazy
from lhotse.dataset import (
- BucketingSampler,
CutConcatenate,
CutMix,
DynamicBucketingSampler,
@@ -218,7 +217,7 @@ class GigaSpeechAsrDataModule:
logging.info("Enable MUSAN")
logging.info("About to get Musan cuts")
cuts_musan = load_manifest(
- self.args.manifest_dir / "cuts_musan.json.gz"
+ self.args.manifest_dir / "musan_cuts.jsonl.gz"
)
transforms.append(
CutMix(
@@ -358,7 +357,7 @@ class GigaSpeechAsrDataModule:
cut_transforms=transforms,
return_cuts=self.args.return_cuts,
)
- valid_sampler = BucketingSampler(
+ valid_sampler = DynamicBucketingSampler(
cuts_valid,
max_duration=self.args.max_duration,
shuffle=False,
@@ -382,8 +381,10 @@ class GigaSpeechAsrDataModule:
else PrecomputedFeatures(),
return_cuts=self.args.return_cuts,
)
- sampler = BucketingSampler(
- cuts, max_duration=self.args.max_duration, shuffle=False
+ sampler = DynamicBucketingSampler(
+ cuts,
+ max_duration=self.args.max_duration,
+ shuffle=False,
)
logging.debug("About to create test dataloader")
test_dl = DataLoader(
@@ -404,7 +405,9 @@ class GigaSpeechAsrDataModule:
@lru_cache()
def dev_cuts(self) -> CutSet:
logging.info("About to get dev cuts")
- cuts_valid = load_manifest(self.args.manifest_dir / "cuts_DEV.jsonl.gz")
+ cuts_valid = load_manifest_lazy(
+ self.args.manifest_dir / "cuts_DEV.jsonl.gz"
+ )
if self.args.small_dev:
return cuts_valid.subset(first=1000)
else:
@@ -413,4 +416,4 @@ class GigaSpeechAsrDataModule:
@lru_cache()
def test_cuts(self) -> CutSet:
logging.info("About to get test cuts")
- return load_manifest(self.args.manifest_dir / "cuts_TEST.jsonl.gz")
+ return load_manifest_lazy(self.args.manifest_dir / "cuts_TEST.jsonl.gz")
diff --git a/egs/gigaspeech/ASR/pruned_transducer_stateless2/decode.py b/egs/gigaspeech/ASR/pruned_transducer_stateless2/decode.py
index ce5116336..5849a3471 100755
--- a/egs/gigaspeech/ASR/pruned_transducer_stateless2/decode.py
+++ b/egs/gigaspeech/ASR/pruned_transducer_stateless2/decode.py
@@ -203,13 +203,13 @@ def get_parser():
def post_processing(
- results: List[Tuple[List[str], List[str]]],
-) -> List[Tuple[List[str], List[str]]]:
+ results: List[Tuple[str, List[str], List[str]]],
+) -> List[Tuple[str, List[str], List[str]]]:
new_results = []
- for ref, hyp in results:
+ for key, ref, hyp in results:
new_ref = asr_text_post_processing(" ".join(ref)).split()
new_hyp = asr_text_post_processing(" ".join(hyp)).split()
- new_results.append((new_ref, new_hyp))
+ new_results.append((key, new_ref, new_hyp))
return new_results
@@ -340,7 +340,7 @@ def decode_dataset(
model: nn.Module,
sp: spm.SentencePieceProcessor,
decoding_graph: Optional[k2.Fsa] = None,
-) -> Dict[str, List[Tuple[List[str], List[str]]]]:
+) -> Dict[str, List[Tuple[str, List[str], List[str]]]]:
"""Decode dataset.
Args:
@@ -374,6 +374,7 @@ def decode_dataset(
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,
@@ -386,9 +387,9 @@ def decode_dataset(
for name, hyps in hyps_dict.items():
this_batch = []
assert len(hyps) == len(texts)
- for hyp_words, ref_text in zip(hyps, texts):
+ for cut_id, hyp_words, ref_text in zip(cut_ids, hyps, texts):
ref_words = ref_text.split()
- this_batch.append((ref_words, hyp_words))
+ this_batch.append((cut_id, ref_words, hyp_words))
results[name].extend(this_batch)
@@ -406,7 +407,7 @@ def decode_dataset(
def save_results(
params: AttributeDict,
test_set_name: str,
- results_dict: Dict[str, List[Tuple[List[int], List[int]]]],
+ results_dict: Dict[str, List[Tuple[str, List[str], List[str]]]],
):
test_set_wers = dict()
for key, results in results_dict.items():
@@ -414,6 +415,7 @@ def save_results(
params.res_dir / f"recogs-{test_set_name}-{key}-{params.suffix}.txt"
)
results = post_processing(results)
+ results = sorted(results)
store_transcripts(filename=recog_path, texts=results)
logging.info(f"The transcripts are stored in {recog_path}")
@@ -544,6 +546,8 @@ def main():
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
gigaspeech = GigaSpeechAsrDataModule(args)
dev_cuts = gigaspeech.dev_cuts()
diff --git a/egs/gigaspeech/ASR/pruned_transducer_stateless2/export.py b/egs/gigaspeech/ASR/pruned_transducer_stateless2/export.py
index 6b3a7a9ff..cff9c7377 100755
--- a/egs/gigaspeech/ASR/pruned_transducer_stateless2/export.py
+++ b/egs/gigaspeech/ASR/pruned_transducer_stateless2/export.py
@@ -131,8 +131,6 @@ def main():
args = get_parser().parse_args()
args.exp_dir = Path(args.exp_dir)
- assert args.jit is False, "Support torchscript will be added later"
-
params = get_params()
params.update(vars(args))
@@ -191,6 +189,11 @@ def main():
model.eval()
if params.jit:
+ # 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"
diff --git a/egs/librispeech/ASR/README.md b/egs/librispeech/ASR/README.md
index e2aaa9d7e..570d1ba1f 100644
--- a/egs/librispeech/ASR/README.md
+++ b/egs/librispeech/ASR/README.md
@@ -23,7 +23,10 @@ The following table lists the differences among them.
| `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_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) |
The decoder in `transducer_stateless` is modified from the paper
[Rnn-Transducer with Stateless Prediction Network](https://ieeexplore.ieee.org/document/9054419/).
diff --git a/egs/librispeech/ASR/RESULTS.md b/egs/librispeech/ASR/RESULTS.md
index f5a32f13b..92323a556 100644
--- a/egs/librispeech/ASR/RESULTS.md
+++ b/egs/librispeech/ASR/RESULTS.md
@@ -1,14 +1,1115 @@
## Results
+### LibriSpeech BPE training results (Pruned Stateless LSTM RNN-T + gradient filter)
+
+#### [lstm_transducer_stateless3](./lstm_transducer_stateless3)
+
+It implements LSTM model with mechanisms in reworked model for streaming ASR.
+Gradient filter is applied inside each lstm module to stabilize the training.
+
+See for more details.
+
+##### training on full librispeech
+
+This model contains 12 encoder layers (LSTM module + Feedforward module). The number of model parameters is 84689496.
+
+The WERs are:
+
+| | test-clean | test-other | comment | decoding mode |
+|-------------------------------------|------------|------------|----------------------|----------------------|
+| greedy search (max sym per frame 1) | 3.66 | 9.51 | --epoch 40 --avg 15 | simulated streaming |
+| greedy search (max sym per frame 1) | 3.66 | 9.48 | --epoch 40 --avg 15 | streaming |
+| fast beam search | 3.55 | 9.33 | --epoch 40 --avg 15 | simulated streaming |
+| fast beam search | 3.57 | 9.25 | --epoch 40 --avg 15 | streaming |
+| modified beam search | 3.55 | 9.28 | --epoch 40 --avg 15 | simulated streaming |
+| modified beam search | 3.54 | 9.25 | --epoch 40 --avg 15 | streaming |
+
+Note: `simulated streaming` indicates feeding full utterance during decoding, while `streaming` indicates feeding certain number of frames at each time.
+
+
+The training command is:
+
+```bash
+./lstm_transducer_stateless3/train.py \
+ --world-size 4 \
+ --num-epochs 40 \
+ --start-epoch 1 \
+ --exp-dir lstm_transducer_stateless3/exp \
+ --full-libri 1 \
+ --max-duration 500 \
+ --master-port 12325 \
+ --num-encoder-layers 12 \
+ --grad-norm-threshold 25.0 \
+ --rnn-hidden-size 1024
+```
+
+The tensorboard log can be found at
+
+
+The simulated streaming decoding command using greedy search, fast beam search, and modified beam search is:
+```bash
+for decoding_method in greedy_search fast_beam_search modified_beam_search; do
+ ./lstm_transducer_stateless3/decode.py \
+ --epoch 40 \
+ --avg 15 \
+ --exp-dir lstm_transducer_stateless3/exp \
+ --max-duration 600 \
+ --num-encoder-layers 12 \
+ --rnn-hidden-size 1024 \
+ --decoding-method $decoding_method \
+ --use-averaged-model True \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8 \
+ --beam-size 4
+done
+```
+
+The streaming decoding command using greedy search, fast beam search, and modified beam search is:
+```bash
+for decoding_method in greedy_search fast_beam_search modified_beam_search; do
+ ./lstm_transducer_stateless3/streaming_decode.py \
+ --epoch 40 \
+ --avg 15 \
+ --exp-dir lstm_transducer_stateless3/exp \
+ --max-duration 600 \
+ --num-encoder-layers 12 \
+ --rnn-hidden-size 1024 \
+ --decoding-method $decoding_method \
+ --use-averaged-model True \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8 \
+ --beam-size 4
+done
+```
+
+Pretrained models, training logs, decoding logs, and decoding results
+are available at
+
+
+
+### LibriSpeech BPE training results (Pruned Stateless LSTM RNN-T + multi-dataset)
+
+#### [lstm_transducer_stateless2](./lstm_transducer_stateless2)
+
+See for more details.
+
+The WERs are:
+
+| | test-clean | test-other | comment |
+|-------------------------------------|------------|------------|-------------------------|
+| 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 |
+| 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 |
+| fast_beam_search | 2.77 | 7.29 | --iter 472000 --avg 18 |
+
+
+The training command is:
+
+```bash
+#!/usr/bin/env bash
+
+export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
+
+./lstm_transducer_stateless2/train.py \
+ --world-size 8 \
+ --num-epochs 35 \
+ --start-epoch 1 \
+ --full-libri 1 \
+ --exp-dir lstm_transducer_stateless2/exp \
+ --max-duration 500 \
+ --use-fp16 0 \
+ --lr-epochs 10 \
+ --num-workers 2 \
+ --giga-prob 0.9
+```
+**Note**: It was killed manually after getting `epoch-18.pt`. Also, we resumed
+training after getting `epoch-9.pt`.
+
+The tensorboard log can be found at
+
+
+The decoding command is
+```bash
+for m in greedy_search fast_beam_search modified_beam_search; do
+ for iter in 472000; do
+ for avg in 8 10 12 14 16 18; do
+ ./lstm_transducer_stateless2/decode.py \
+ --iter $iter \
+ --avg $avg \
+ --exp-dir lstm_transducer_stateless2/exp \
+ --max-duration 600 \
+ --num-encoder-layers 12 \
+ --rnn-hidden-size 1024 \
+ --decoding-method $m \
+ --use-averaged-model True \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8 \
+ --beam-size 4
+ done
+ done
+done
+```
+
+Pretrained models, training logs, decoding logs, and decoding results
+are available at
+
+
+
+### LibriSpeech BPE training results (Pruned Stateless LSTM RNN-T)
+
+#### [lstm_transducer_stateless](./lstm_transducer_stateless)
+
+It implements LSTM model with mechanisms in reworked model for streaming ASR.
+
+See for more details.
+
+##### training on full librispeech
+
+This model contains 12 encoder layers (LSTM module + Feedforward module). The number of model parameters is 84689496.
+
+The WERs are:
+
+| | test-clean | test-other | comment | decoding mode |
+|-------------------------------------|------------|------------|----------------------|----------------------|
+| greedy search (max sym per frame 1) | 3.81 | 9.73 | --epoch 35 --avg 15 | simulated streaming |
+| greedy search (max sym per frame 1) | 3.78 | 9.79 | --epoch 35 --avg 15 | streaming |
+| fast beam search | 3.74 | 9.59 | --epoch 35 --avg 15 | simulated streaming |
+| fast beam search | 3.73 | 9.61 | --epoch 35 --avg 15 | streaming |
+| modified beam search | 3.64 | 9.55 | --epoch 35 --avg 15 | simulated streaming |
+| modified beam search | 3.65 | 9.51 | --epoch 35 --avg 15 | streaming |
+
+Note: `simulated streaming` indicates feeding full utterance during decoding, while `streaming` indicates feeding certain number of frames at each time.
+
+The training command is:
+
+```bash
+./lstm_transducer_stateless/train.py \
+ --world-size 4 \
+ --num-epochs 35 \
+ --start-epoch 1 \
+ --exp-dir lstm_transducer_stateless/exp \
+ --full-libri 1 \
+ --max-duration 500 \
+ --master-port 12321 \
+ --num-encoder-layers 12 \
+ --rnn-hidden-size 1024
+```
+
+The tensorboard log can be found at
+
+
+The simulated streaming decoding command using greedy search, fast beam search, and modified beam search is:
+```bash
+for decoding_method in greedy_search fast_beam_search modified_beam_search; do
+ ./lstm_transducer_stateless/decode.py \
+ --epoch 35 \
+ --avg 15 \
+ --exp-dir lstm_transducer_stateless/exp \
+ --max-duration 600 \
+ --num-encoder-layers 12 \
+ --rnn-hidden-size 1024 \
+ --decoding-method $decoding_method \
+ --use-averaged-model True \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8 \
+ --beam-size 4
+done
+```
+
+The streaming decoding command using greedy search, fast beam search, and modified beam search is:
+```bash
+for decoding_method in greedy_search fast_beam_search modified_beam_search; do
+ ./lstm_transducer_stateless/streaming_decode.py \
+ --epoch 35 \
+ --avg 15 \
+ --exp-dir lstm_transducer_stateless/exp \
+ --max-duration 600 \
+ --num-encoder-layers 12 \
+ --rnn-hidden-size 1024 \
+ --decoding-method $decoding_method \
+ --use-averaged-model True \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8 \
+ --beam-size 4
+done
+```
+
+Pretrained models, training logs, decoding logs, and decoding results
+are available at
+
+
+
+#### LibriSpeech BPE training results (Pruned Stateless Conv-Emformer RNN-T 2)
+
+[conv_emformer_transducer_stateless2](./conv_emformer_transducer_stateless2)
+
+It implements [Emformer](https://arxiv.org/abs/2010.10759) augmented with convolution module and simplified memory bank for streaming ASR.
+It is modified from [torchaudio](https://github.com/pytorch/audio).
+
+See for more details.
+
+##### With lower latency setup, training on full librispeech
+
+In this model, the lengths of chunk and right context are 32 frames (i.e., 0.32s) and 8 frames (i.e., 0.08s), respectively.
+
+The WERs are:
+
+| | test-clean | test-other | comment | decoding mode |
+|-------------------------------------|------------|------------|----------------------|----------------------|
+| greedy search (max sym per frame 1) | 3.5 | 9.09 | --epoch 30 --avg 10 | simulated streaming |
+| greedy search (max sym per frame 1) | 3.57 | 9.1 | --epoch 30 --avg 10 | streaming |
+| fast beam search | 3.5 | 8.91 | --epoch 30 --avg 10 | simulated streaming |
+| fast beam search | 3.54 | 8.91 | --epoch 30 --avg 10 | streaming |
+| modified beam search | 3.43 | 8.86 | --epoch 30 --avg 10 | simulated streaming |
+| modified beam search | 3.48 | 8.88 | --epoch 30 --avg 10 | streaming |
+
+The training command is:
+
+```bash
+./conv_emformer_transducer_stateless2/train.py \
+ --world-size 6 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --full-libri 1 \
+ --max-duration 280 \
+ --master-port 12321 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32
+```
+
+The tensorboard log can be found at
+
+
+The simulated streaming decoding command using greedy search is:
+```bash
+./conv_emformer_transducer_stateless2/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --max-duration 300 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method greedy_search \
+ --use-averaged-model True
+```
+
+The simulated streaming decoding command using fast beam search is:
+```bash
+./conv_emformer_transducer_stateless2/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --max-duration 300 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method fast_beam_search \
+ --use-averaged-model True \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8
+```
+
+The simulated streaming decoding command using modified beam search is:
+```bash
+./conv_emformer_transducer_stateless2/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --max-duration 300 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method modified_beam_search \
+ --use-averaged-model True \
+ --beam-size 4
+```
+
+The streaming decoding command using greedy search is:
+```bash
+./conv_emformer_transducer_stateless2/streaming_decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --num-decode-streams 2000 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method greedy_search \
+ --use-averaged-model True
+```
+
+The streaming decoding command using fast beam search is:
+```bash
+./conv_emformer_transducer_stateless2/streaming_decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --num-decode-streams 2000 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method fast_beam_search \
+ --use-averaged-model True \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8
+```
+
+The streaming decoding command using modified beam search is:
+```bash
+./conv_emformer_transducer_stateless2/streaming_decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --num-decode-streams 2000 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method modified_beam_search \
+ --use-averaged-model True \
+ --beam-size 4
+```
+
+Pretrained models, training logs, decoding logs, and decoding results
+are available at
+
+
+##### With higher latency setup, training on full librispeech
+
+In this model, the lengths of chunk and right context are 64 frames (i.e., 0.64s) and 16 frames (i.e., 0.16s), respectively.
+
+The WERs are:
+
+| | test-clean | test-other | comment | decoding mode |
+|-------------------------------------|------------|------------|----------------------|----------------------|
+| greedy search (max sym per frame 1) | 3.3 | 8.71 | --epoch 30 --avg 10 | simulated streaming |
+| greedy search (max sym per frame 1) | 3.35 | 8.65 | --epoch 30 --avg 10 | streaming |
+| fast beam search | 3.27 | 8.58 | --epoch 30 --avg 10 | simulated streaming |
+| fast beam search | 3.31 | 8.48 | --epoch 30 --avg 10 | streaming |
+| modified beam search | 3.26 | 8.56 | --epoch 30 --avg 10 | simulated streaming |
+| modified beam search | 3.29 | 8.47 | --epoch 30 --avg 10 | streaming |
+
+The training command is:
+
+```bash
+./conv_emformer_transducer_stateless2/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --full-libri 1 \
+ --max-duration 280 \
+ --master-port 12321 \
+ --num-encoder-layers 12 \
+ --chunk-length 64 \
+ --cnn-module-kernel 31 \
+ --left-context-length 64 \
+ --right-context-length 16 \
+ --memory-size 32
+```
+
+The tensorboard log can be found at
+
+
+The simulated streaming decoding command using greedy search is:
+```bash
+./conv_emformer_transducer_stateless2/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --max-duration 300 \
+ --num-encoder-layers 12 \
+ --chunk-length 64 \
+ --cnn-module-kernel 31 \
+ --left-context-length 64 \
+ --right-context-length 16 \
+ --memory-size 32 \
+ --decoding-method greedy_search \
+ --use-averaged-model True
+```
+
+The simulated streaming decoding command using fast beam search is:
+```bash
+./conv_emformer_transducer_stateless2/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --max-duration 300 \
+ --num-encoder-layers 12 \
+ --chunk-length 64 \
+ --cnn-module-kernel 31 \
+ --left-context-length 64 \
+ --right-context-length 16 \
+ --memory-size 32 \
+ --decoding-method fast_beam_search \
+ --use-averaged-model True \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8
+```
+
+The simulated streaming decoding command using modified beam search is:
+```bash
+./conv_emformer_transducer_stateless2/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --max-duration 300 \
+ --num-encoder-layers 12 \
+ --chunk-length 64 \
+ --cnn-module-kernel 31 \
+ --left-context-length 64 \
+ --right-context-length 16 \
+ --memory-size 32 \
+ --decoding-method modified_beam_search \
+ --use-averaged-model True \
+ --beam-size 4
+```
+
+The streaming decoding command using greedy search is:
+```bash
+./conv_emformer_transducer_stateless2/streaming_decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --num-decode-streams 2000 \
+ --num-encoder-layers 12 \
+ --chunk-length 64 \
+ --cnn-module-kernel 31 \
+ --left-context-length 64 \
+ --right-context-length 16 \
+ --memory-size 32 \
+ --decoding-method greedy_search \
+ --use-averaged-model True
+```
+
+The streaming decoding command using fast beam search is:
+```bash
+./conv_emformer_transducer_stateless2/streaming_decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --num-decode-streams 2000 \
+ --num-encoder-layers 12 \
+ --chunk-length 64 \
+ --cnn-module-kernel 31 \
+ --left-context-length 64 \
+ --right-context-length 16 \
+ --memory-size 32 \
+ --decoding-method fast_beam_search \
+ --use-averaged-model True \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8
+```
+
+The streaming decoding command using modified beam search is:
+```bash
+./conv_emformer_transducer_stateless2/streaming_decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --num-decode-streams 2000 \
+ --num-encoder-layers 12 \
+ --chunk-length 64 \
+ --cnn-module-kernel 31 \
+ --left-context-length 64 \
+ --right-context-length 16 \
+ --memory-size 32 \
+ --decoding-method modified_beam_search \
+ --use-averaged-model True \
+ --beam-size 4
+```
+
+Pretrained models, training logs, decoding logs, and decoding results
+are available at
+
+
+
+### LibriSpeech BPE training results (Pruned Stateless Streaming Conformer RNN-T)
+
+#### [pruned_transducer_stateless](./pruned_transducer_stateless)
+
+See for more details.
+
+##### Training on full librispeech
+The WERs are (the number in the table formatted as test-clean & test-other):
+
+We only trained 25 epochs for saving time, if you want to get better results you can train more epochs.
+
+| decoding method | left context | chunk size = 2 | chunk size = 4 | chunk size = 8 | chunk size = 16|
+|----------------------|--------------|----------------|----------------|----------------|----------------|
+| greedy search | 32 | 4.74 & 11.38 | 4.57 & 10.86 | 4.18 & 10.37 | 3.87 & 9.85 |
+| greedy search | 64 | 4.74 & 11.25 | 4.48 & 10.72 | 4.1 & 10.24 | 3.85 & 9.73 |
+| fast beam search | 32 | 4.75 & 11.1 | 4.48 & 10.65 | 4.12 & 10.18 | 3.95 & 9.67 |
+| fast beam search | 64 | 4.7 & 11 | 4.37 & 10.49 | 4.07 & 10.04 | 3.89 & 9.53 |
+| modified beam search | 32 | 4.64 & 10.94 | 4.38 & 10.51 | 4.11 & 10.14 | 3.87 & 9.61 |
+| modified beam search | 64 | 4.59 & 10.81 | 4.29 & 10.39 | 4.02 & 10.02 | 3.84 & 9.43 |
+
+**NOTE:** The WERs in table above were decoded with simulate streaming method (i.e. using masking strategy), see commands below. We also have [real streaming decoding](./pruned_transducer_stateless/streaming_decode.py) script which should produce almost the same results. We tried adding right context in the real streaming decoding, but it seemed not to benefit the performance for all the models, the reasons might be the training and decoding mismatching.
+
+The training command is:
+
+```bash
+./pruned_transducer_stateless/train.py \
+ --exp-dir pruned_transducer_stateless/exp \
+ --full-libri 1 \
+ --dynamic-chunk-training 1 \
+ --causal-convolution 1 \
+ --short-chunk-size 20 \
+ --num-left-chunks 4 \
+ --max-duration 300 \
+ --world-size 4 \
+ --start-epoch 0 \
+ --num-epochs 25
+```
+
+You can find the tensorboard log here
+
+The decoding command is:
+```bash
+decoding_method="greedy_search" # "fast_beam_search", "modified_beam_search"
+
+for chunk in 2 4 8 16; do
+ for left in 32 64; do
+ ./pruned_transducer_stateless/decode.py \
+ --simulate-streaming 1 \
+ --decode-chunk-size ${chunk} \
+ --left-context ${left} \
+ --causal-convolution 1 \
+ --epoch 24 \
+ --avg 10 \
+ --exp-dir ./pruned_transducer_stateless/exp \
+ --max-sym-per-frame 1 \
+ --max-duration 1000 \
+ --decoding-method ${decoding_method}
+ done
+done
+```
+
+Pre-trained models, training and decoding logs, and decoding results are available at
+
+#### [pruned_transducer_stateless2](./pruned_transducer_stateless2)
+
+See for more details.
+
+##### Training on full librispeech
+The WERs are (the number in the table formatted as test-clean & test-other):
+
+We only trained 25 epochs for saving time, if you want to get better results you can train more epochs.
+
+| decoding method | left context | chunk size = 2 | chunk size = 4 | chunk size = 8 | chunk size = 16|
+|----------------------|--------------|----------------|----------------|----------------|----------------|
+| greedy search | 32 | 4.2 & 10.64 | 3.97 & 10.03 | 3.83 & 9.58 | 3.7 & 9.11 |
+| greedy search | 64 | 4.16 & 10.5 | 3.93 & 9.99 | 3.73 & 9.45 | 3.63 & 9.04 |
+| fast beam search | 32 | 4.13 & 10.3 | 3.93 & 9.82 | 3.8 & 9.35 | 3.62 & 8.93 |
+| fast beam search | 64 | 4.13 & 10.22 | 3.89 & 9.68 | 3.73 & 9.27 | 3.52 & 8.82 |
+| modified beam search | 32 | 4.02 & 10.22 | 3.9 & 9.71 | 3.74 & 9.33 | 3.59 & 8.87 |
+| modified beam search | 64 | 4.05 & 10.08 | 3.81 & 9.67 | 3.68 & 9.21 | 3.56 & 8.77 |
+
+**NOTE:** The WERs in table above were decoded with simulate streaming method (i.e. using masking strategy), see commands below. We also have [real streaming decoding](./pruned_transducer_stateless2/streaming_decode.py) script which should produce almost the same results. We tried adding right context in the real streaming decoding, but it seemed not to benefit the performance for all the models, the reasons might be the training and decoding mismatching.
+
+The training command is:
+
+```bash
+./pruned_transducer_stateless2/train.py \
+ --exp-dir pruned_transducer_stateless2/exp \
+ --full-libri 1 \
+ --dynamic-chunk-training 1 \
+ --causal-convolution 1 \
+ --short-chunk-size 20 \
+ --num-left-chunks 4 \
+ --max-duration 300 \
+ --world-size 4 \
+ --start-epoch 0 \
+ --num-epochs 25
+```
+
+You can find the tensorboard log here
+
+The decoding command is:
+```bash
+decoding_method="greedy_search" # "fast_beam_search", "modified_beam_search"
+
+for chunk in 2 4 8 16; do
+ for left in 32 64; do
+ ./pruned_transducer_stateless2/decode.py \
+ --simulate-streaming 1 \
+ --decode-chunk-size ${chunk} \
+ --left-context ${left} \
+ --causal-convolution 1 \
+ --epoch 24 \
+ --avg 10 \
+ --exp-dir ./pruned_transducer_stateless2/exp \
+ --max-sym-per-frame 1 \
+ --max-duration 1000 \
+ --decoding-method ${decoding_method}
+ done
+done
+```
+
+Pre-trained models, training and decoding logs, and decoding results are available at
+
+#### [pruned_transducer_stateless3](./pruned_transducer_stateless3)
+
+See for more details.
+
+##### Training on full librispeech (**Use giga_prob = 0.5**)
+
+The WERs are (the number in the table formatted as test-clean & test-other):
+
+| decoding method | left context | chunk size = 2 | chunk size = 4 | chunk size = 8 | chunk size = 16|
+|----------------------|--------------|----------------|----------------|----------------|----------------|
+| greedy search | 32 | 3.7 & 9.53 | 3.45 & 8.88 | 3.28 & 8.45 | 3.13 & 7.93 |
+| greedy search | 64 | 3.69 & 9.36 | 3.39 & 8.68 | 3.28 & 8.19 | 3.08 & 7.83 |
+| fast beam search | 32 | 3.71 & 9.18 | 3.36 & 8.65 | 3.23 & 8.23 | 3.17 & 7.78 |
+| fast beam search | 64 | 3.61 & 9.03 | 3.46 & 8.43 | 3.2 & 8.0 | 3.11 & 7.63 |
+| modified beam search | 32 | 3.56 & 9.08 | 3.34 & 8.58 | 3.21 & 8.14 | 3.06 & 7.73 |
+| modified beam search | 64 | 3.55 & 8.86 | 3.29 & 8.34 | 3.16 & 8.01 | 3.05 & 7.57 |
+
+**NOTE:** The WERs in table above were decoded with simulate streaming method (i.e. using masking strategy), see commands below. We also have [real streaming decoding](./pruned_transducer_stateless3/streaming_decode.py) script which should produce almost the same results. We tried adding right context in the real streaming decoding, but it seemed not to benefit the performance for all the models, the reasons might be the training and decoding mismatching.
+
+The training command is (Note: this model was trained with mix-precision training):
+
+```bash
+./pruned_transducer_stateless3/train.py \
+ --exp-dir pruned_transducer_stateless3/exp \
+ --full-libri 1 \
+ --dynamic-chunk-training 1 \
+ --causal-convolution 1 \
+ --short-chunk-size 32 \
+ --num-left-chunks 4 \
+ --max-duration 300 \
+ --world-size 4 \
+ --use-fp16 1 \
+ --start-epoch 0 \
+ --num-epochs 37 \
+ --num-workers 2 \
+ --giga-prob 0.5
+```
+
+You can find the tensorboard log here
+
+The decoding command is:
+```bash
+decoding_method="greedy_search" # "fast_beam_search", "modified_beam_search"
+
+for chunk in 2 4 8 16; do
+ for left in 32 64; do
+ ./pruned_transducer_stateless3/decode.py \
+ --simulate-streaming 1 \
+ --decode-chunk-size ${chunk} \
+ --left-context ${left} \
+ --causal-convolution 1 \
+ --epoch 36 \
+ --avg 8 \
+ --exp-dir ./pruned_transducer_stateless3/exp \
+ --max-sym-per-frame 1 \
+ --max-duration 1000 \
+ --decoding-method ${decoding_method}
+ done
+done
+```
+
+Pre-trained models, training and decoding logs, and decoding results are available at
+
+##### Training on full librispeech (**Use giga_prob = 0.9**)
+
+The WERs are (the number in the table formatted as test-clean & test-other):
+
+| decoding method | left context | chunk size = 2 | chunk size = 4 | chunk size = 8 | chunk size = 16|
+|----------------------|--------------|----------------|----------------|----------------|----------------|
+| greedy search | 32 | 3.25 & 8.2 | 3.07 & 7.67 | 2.91 & 7.28 | 2.8 & 6.89 |
+| greedy search | 64 | 3.22 & 8.12 | 3.05 & 7.59 | 2.91 & 7.07 | 2.78 & 6.81 |
+| fast beam search | 32 | 3.26 & 8.2 | 3.06 & 7.56 | 2.98 & 7.08 | 2.77 & 6.75 |
+| fast beam search | 64 | 3.24 & 8.09 | 3.06 & 7.43 | 2.88 & 7.03 | 2.73 & 6.68 |
+| modified beam search | 32 | 3.13 & 7.91 | 2.99 & 7.45 | 2.83 & 6.98 | 2.68 & 6.75 |
+| modified beam search | 64 | 3.08 & 7.8 | 2.97 & 7.37 | 2.81 & 6.82 | 2.66 & 6.67 |
+
+**NOTE:** The WERs in table above were decoded with simulate streaming method (i.e. using masking strategy), see commands below. We also have [real streaming decoding](./pruned_transducer_stateless3/streaming_decode.py) script which should produce almost the same results. We tried adding right context in the real streaming decoding, but it seemed not to benefit the performance for all the models, the reasons might be the training and decoding mismatching.
+
+The training command is:
+
+```bash
+./pruned_transducer_stateless3/train.py \
+ --exp-dir pruned_transducer_stateless3/exp \
+ --full-libri 1 \
+ --dynamic-chunk-training 1 \
+ --causal-convolution 1 \
+ --short-chunk-size 25 \
+ --num-left-chunks 8 \
+ --max-duration 300 \
+ --world-size 8 \
+ --start-epoch 0 \
+ --num-epochs 26 \
+ --num-workers 2 \
+ --giga-prob 0.9
+```
+
+You can find the tensorboard log here
+
+The decoding command is:
+```bash
+decoding_method="greedy_search" # "fast_beam_search", "modified_beam_search"
+
+for chunk in 2 4 8 16; do
+ for left in 32 64; do
+ ./pruned_transducer_stateless3/decode.py \
+ --simulate-streaming 1 \
+ --decode-chunk-size ${chunk} \
+ --left-context ${left} \
+ --causal-convolution 1 \
+ --epoch 25 \
+ --avg 12 \
+ --exp-dir ./pruned_transducer_stateless3/exp \
+ --max-sym-per-frame 1 \
+ --max-duration 1000 \
+ --decoding-method ${decoding_method}
+ done
+done
+```
+
+Pre-trained models, training and decoding logs, and decoding results are available at
+
+#### [pruned_transducer_stateless4](./pruned_transducer_stateless4)
+
+See for more details.
+
+##### Training on full librispeech
+The WERs are (the number in the table formatted as test-clean & test-other):
+
+We only trained 25 epochs for saving time, if you want to get better results you can train more epochs.
+
+| decoding method | left context | chunk size = 2 | chunk size = 4 | chunk size = 8 | chunk size = 16|
+|----------------------|--------------|----------------|----------------|----------------|----------------|
+| greedy search | 32 | 3.96 & 10.45 | 3.73 & 9.97 | 3.54 & 9.56 | 3.45 & 9.08 |
+| greedy search | 64 | 3.9 & 10.34 | 3.7 & 9.9 | 3.53 & 9.41 | 3.39 & 9.03 |
+| fast beam search | 32 | 3.9 & 10.09 | 3.69 & 9.65 | 3.58 & 9.28 | 3.46 & 8.91 |
+| fast beam search | 64 | 3.82 & 10.03 | 3.67 & 9.56 | 3.51 & 9.18 | 3.43 & 8.78 |
+| modified beam search | 32 | 3.78 & 10.0 | 3.63 & 9.54 | 3.43 & 9.29 | 3.39 & 8.84 |
+| modified beam search | 64 | 3.76 & 9.95 | 3.54 & 9.48 | 3.4 & 9.13 | 3.33 & 8.74 |
+
+**NOTE:** The WERs in table above were decoded with simulate streaming method (i.e. using masking strategy), see commands below. We also have [real streaming decoding](./pruned_transducer_stateless4/streaming_decode.py) script which should produce almost the same results. We tried adding right context in the real streaming decoding, but it seemed not to benefit the performance for all the models, the reasons might be the training and decoding mismatching.
+
+The training command is:
+
+```bash
+./pruned_transducer_stateless4/train.py \
+ --exp-dir pruned_transducer_stateless4/exp \
+ --full-libri 1 \
+ --dynamic-chunk-training 1 \
+ --causal-convolution 1 \
+ --short-chunk-size 20 \
+ --num-left-chunks 4 \
+ --max-duration 300 \
+ --world-size 4 \
+ --start-epoch 1 \
+ --num-epochs 25
+```
+
+You can find the tensorboard log here
+
+The decoding command is:
+```bash
+decoding_method="greedy_search" # "fast_beam_search", "modified_beam_search"
+
+for chunk in 2 4 8 16; do
+ for left in 32 64; do
+ ./pruned_transducer_stateless4/decode.py \
+ --simulate-streaming 1 \
+ --decode-chunk-size ${chunk} \
+ --left-context ${left} \
+ --causal-convolution 1 \
+ --epoch 25 \
+ --avg 3 \
+ --exp-dir ./pruned_transducer_stateless4/exp \
+ --max-sym-per-frame 1 \
+ --max-duration 1000 \
+ --decoding-method ${decoding_method}
+ done
+done
+```
+
+Pre-trained models, training and decoding logs, and decoding results are available at
+
+#### [pruned_transducer_stateless5](./pruned_transducer_stateless5)
+
+See for more details.
+
+##### Training on full librispeech
+The WERs are (the number in the table formatted as test-clean & test-other):
+
+We only trained 25 epochs for saving time, if you want to get better results you can train more epochs.
+
+| decoding method | left context | chunk size = 2 | chunk size = 4 | chunk size = 8 | chunk size = 16|
+|----------------------|--------------|----------------|----------------|----------------|----------------|
+| greedy search | 32 | 3.93 & 9.88 | 3.64 & 9.43 | 3.51 & 8.92 | 3.26 & 8.37 |
+| greedy search | 64 | 4.84 & 9.81 | 3.59 & 9.27 | 3.44 & 8.83 | 3.23 & 8.33 |
+| fast beam search | 32 | 3.86 & 9.77 | 3.67 & 9.3 | 3.5 & 8.83 | 3.27 & 8.33 |
+| fast beam search | 64 | 3.79 & 9.68 | 3.57 & 9.21 | 3.41 & 8.72 | 3.25 & 8.27 |
+| modified beam search | 32 | 3.84 & 9.71 | 3.66 & 9.38 | 3.47 & 8.86 | 3.26 & 8.42 |
+| modified beam search | 64 | 3.81 & 9.59 | 3.58 & 9.2 | 3.44 & 8.74 | 3.23 & 8.35 |
+
+
+**NOTE:** The WERs in table above were decoded with simulate streaming method (i.e. using masking strategy), see commands below. We also have [real streaming decoding](./pruned_transducer_stateless5/streaming_decode.py) script which should produce almost the same results. We tried adding right context in the real streaming decoding, but it seemed not to benefit the performance for all the models, the reasons might be the training and decoding mismatching.
+
+The training command is:
+
+```bash
+./pruned_transducer_stateless5/train.py \
+ --exp-dir pruned_transducer_stateless5/exp \
+ --num-encoder-layers 18 \
+ --dim-feedforward 2048 \
+ --nhead 8 \
+ --encoder-dim 512 \
+ --decoder-dim 512 \
+ --joiner-dim 512 \
+ --full-libri 1 \
+ --dynamic-chunk-training 1 \
+ --causal-convolution 1 \
+ --short-chunk-size 20 \
+ --num-left-chunks 4 \
+ --max-duration 300 \
+ --world-size 4 \
+ --start-epoch 1 \
+ --num-epochs 25
+```
+
+You can find the tensorboard log here
+
+The decoding command is:
+```bash
+decoding_method="greedy_search" # "fast_beam_search", "modified_beam_search"
+
+for chunk in 2 4 8 16; do
+ for left in 32 64; do
+ ./pruned_transducer_stateless5/decode.py \
+ --num-encoder-layers 18 \
+ --dim-feedforward 2048 \
+ --nhead 8 \
+ --encoder-dim 512 \
+ --decoder-dim 512 \
+ --joiner-dim 512 \
+ --simulate-streaming 1 \
+ --decode-chunk-size ${chunk} \
+ --left-context ${left} \
+ --causal-convolution 1 \
+ --epoch 25 \
+ --avg 3 \
+ --exp-dir ./pruned_transducer_stateless5/exp \
+ --max-sym-per-frame 1 \
+ --max-duration 1000 \
+ --decoding-method ${decoding_method}
+ done
+done
+```
+
+Pre-trained models, training and decoding logs, and decoding results are available at
+
+
+### LibriSpeech BPE training results (Pruned Stateless Conv-Emformer RNN-T)
+
+#### [conv_emformer_transducer_stateless](./conv_emformer_transducer_stateless)
+
+It implements [Emformer](https://arxiv.org/abs/2010.10759) augmented with convolution module for streaming ASR.
+It is modified from [torchaudio](https://github.com/pytorch/audio).
+
+See for more details.
+
+##### Training on full librispeech
+
+In this model, the lengths of chunk and right context are 32 frames (i.e., 0.32s) and 8 frames (i.e., 0.08s), respectively.
+
+The WERs are:
+
+| | test-clean | test-other | comment | decoding mode |
+|-------------------------------------|------------|------------|----------------------|----------------------|
+| greedy search (max sym per frame 1) | 3.63 | 9.61 | --epoch 30 --avg 10 | simulated streaming |
+| greedy search (max sym per frame 1) | 3.64 | 9.65 | --epoch 30 --avg 10 | streaming |
+| fast beam search | 3.61 | 9.4 | --epoch 30 --avg 10 | simulated streaming |
+| fast beam search | 3.58 | 9.5 | --epoch 30 --avg 10 | streaming |
+| modified beam search | 3.56 | 9.41 | --epoch 30 --avg 10 | simulated streaming |
+| modified beam search | 3.54 | 9.46 | --epoch 30 --avg 10 | streaming |
+
+The training command is:
+
+```bash
+./conv_emformer_transducer_stateless/train.py \
+ --world-size 6 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --exp-dir conv_emformer_transducer_stateless/exp \
+ --full-libri 1 \
+ --max-duration 300 \
+ --master-port 12321 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32
+```
+
+The tensorboard log can be found at
+
+
+The simulated streaming decoding command using greedy search is:
+```bash
+./conv_emformer_transducer_stateless/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless/exp \
+ --max-duration 300 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method greedy_search \
+ --use-averaged-model True
+```
+
+The simulated streaming decoding command using fast beam search is:
+```bash
+./conv_emformer_transducer_stateless/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless/exp \
+ --max-duration 300 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method fast_beam_search \
+ --use-averaged-model True \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8
+```
+
+The simulated streaming decoding command using modified beam search is:
+```bash
+./conv_emformer_transducer_stateless/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless/exp \
+ --max-duration 300 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method modified_beam_search \
+ --use-averaged-model True \
+ --beam-size 4
+```
+
+The streaming decoding command using greedy search is:
+```bash
+./conv_emformer_transducer_stateless/streaming_decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless/exp \
+ --num-decode-streams 2000 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method greedy_search \
+ --use-averaged-model True
+```
+
+The streaming decoding command using fast beam search is:
+```bash
+./conv_emformer_transducer_stateless/streaming_decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless/exp \
+ --num-decode-streams 2000 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method fast_beam_search \
+ --use-averaged-model True \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8
+```
+
+The streaming decoding command using modified beam search is:
+```bash
+./conv_emformer_transducer_stateless/streaming_decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless/exp \
+ --num-decode-streams 2000 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method modified_beam_search \
+ --use-averaged-model True \
+ --beam-size 4
+```
+
+Pretrained models, training logs, decoding logs, and decoding results
+are available at
+
+
### LibriSpeech BPE training results (Pruned Stateless Emformer RNN-T)
-[pruned_stateless_emformer_rnnt2](./pruned_stateless_emformer_rnnt2)
+#### [pruned_stateless_emformer_rnnt2](./pruned_stateless_emformer_rnnt2)
Use .
Use [Emformer](https://arxiv.org/abs/2010.10759) from [torchaudio](https://github.com/pytorch/audio)
for streaming ASR. The Emformer model is imported from torchaudio without modifications.
+You can use to deploy it.
+
| | test-clean | test-other | comment |
|-------------------------------------|------------|------------|----------------------------------------|
| greedy search (max sym per frame 1) | 4.28 | 11.42 | --epoch 39 --avg 6 --max-duration 600 |
@@ -68,7 +1169,7 @@ results at:
### LibriSpeech BPE training results (Pruned Stateless Transducer 5)
-[pruned_transducer_stateless5](./pruned_transducer_stateless5)
+#### [pruned_transducer_stateless5](./pruned_transducer_stateless5)
Same as `Pruned Stateless Transducer 2` but with more layers.
@@ -81,15 +1182,15 @@ The notations `large` and `medium` below are from the [Conformer](https://arxiv.
paper, where the large model has about 118 M parameters and the medium model
has 30.8 M parameters.
-#### Large
+##### Large
Number of model parameters 118129516 (i.e, 118.13 M).
| | test-clean | test-other | comment |
|-------------------------------------|------------|------------|----------------------------------------|
-| greedy search (max sym per frame 1) | 2.39 | 5.57 | --epoch 39 --avg 7 --max-duration 600 |
-| modified beam search | 2.35 | 5.50 | --epoch 39 --avg 7 --max-duration 600 |
-| fast beam search | 2.38 | 5.50 | --epoch 39 --avg 7 --max-duration 600 |
+| greedy search (max sym per frame 1) | 2.43 | 5.72 | --epoch 30 --avg 10 --max-duration 600 |
+| modified beam search | 2.43 | 5.69 | --epoch 30 --avg 10 --max-duration 600 |
+| fast beam search | 2.43 | 5.67 | --epoch 30 --avg 10 --max-duration 600 |
The training commands are:
@@ -98,8 +1199,8 @@ export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
./pruned_transducer_stateless5/train.py \
--world-size 8 \
- --num-epochs 40 \
- --start-epoch 0 \
+ --num-epochs 30 \
+ --start-epoch 1 \
--full-libri 1 \
--exp-dir pruned_transducer_stateless5/exp-L \
--max-duration 300 \
@@ -113,15 +1214,15 @@ export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
```
The tensorboard log can be found at
-
+
The decoding commands are:
```bash
for method in greedy_search modified_beam_search fast_beam_search; do
./pruned_transducer_stateless5/decode.py \
- --epoch 39 \
- --avg 7 \
+ --epoch 30 \
+ --avg 10 \
--exp-dir ./pruned_transducer_stateless5/exp-L \
--max-duration 600 \
--decoding-method $method \
@@ -131,24 +1232,25 @@ for method in greedy_search modified_beam_search fast_beam_search; do
--nhead 8 \
--encoder-dim 512 \
--decoder-dim 512 \
- --joiner-dim 512
+ --joiner-dim 512 \
+ --use-averaged-model True
done
```
You can find a pretrained model, training logs, decoding logs, and decoding
results at:
-
+
-#### Medium
+##### Medium
Number of model parameters 30896748 (i.e, 30.9 M).
| | test-clean | test-other | comment |
|-------------------------------------|------------|------------|-----------------------------------------|
-| greedy search (max sym per frame 1) | 2.88 | 6.69 | --epoch 39 --avg 17 --max-duration 600 |
-| modified beam search | 2.83 | 6.59 | --epoch 39 --avg 17 --max-duration 600 |
-| fast beam search | 2.83 | 6.61 | --epoch 39 --avg 17 --max-duration 600 |
+| greedy search (max sym per frame 1) | 2.87 | 6.92 | --epoch 30 --avg 10 --max-duration 600 |
+| modified beam search | 2.83 | 6.75 | --epoch 30 --avg 10 --max-duration 600 |
+| fast beam search | 2.81 | 6.76 | --epoch 30 --avg 10 --max-duration 600 |
The training commands are:
@@ -157,8 +1259,8 @@ export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
./pruned_transducer_stateless5/train.py \
--world-size 8 \
- --num-epochs 40 \
- --start-epoch 0 \
+ --num-epochs 30 \
+ --start-epoch 1 \
--full-libri 1 \
--exp-dir pruned_transducer_stateless5/exp-M \
--max-duration 300 \
@@ -172,15 +1274,15 @@ export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
```
The tensorboard log can be found at
-
+
The decoding commands are:
```bash
for method in greedy_search modified_beam_search fast_beam_search; do
./pruned_transducer_stateless5/decode.py \
- --epoch 39 \
- --avg 17 \
+ --epoch 30 \
+ --avg 10 \
--exp-dir ./pruned_transducer_stateless5/exp-M \
--max-duration 600 \
--decoding-method $method \
@@ -190,35 +1292,36 @@ for method in greedy_search modified_beam_search fast_beam_search; do
--nhead 4 \
--encoder-dim 256 \
--decoder-dim 512 \
- --joiner-dim 512
+ --joiner-dim 512 \
+ --use-averaged-model True
done
```
You can find a pretrained model, training logs, decoding logs, and decoding
results at:
-
+
-#### Baseline-2
+##### Baseline-2
-It has 88.98 M parameters. Compared to the model in pruned_transducer_stateless2, its has more
+It has 87.8 M parameters. Compared to the model in pruned_transducer_stateless2, its has more
layers (24 v.s 12) but a narrower model (1536 feedforward dim and 384 encoder dim vs 2048 feed forward dim and 512 encoder dim).
| | test-clean | test-other | comment |
|-------------------------------------|------------|------------|-----------------------------------------|
-| greedy search (max sym per frame 1) | 2.41 | 5.70 | --epoch 31 --avg 17 --max-duration 600 |
-| modified beam search | 2.41 | 5.69 | --epoch 31 --avg 17 --max-duration 600 |
-| fast beam search | 2.41 | 5.69 | --epoch 31 --avg 17 --max-duration 600 |
+| 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 |
+| fast beam search | 2.5 | 5.72 | --epoch 30 --avg 10 --max-duration 600 |
```bash
export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
./pruned_transducer_stateless5/train.py \
--world-size 8 \
- --num-epochs 40 \
- --start-epoch 0 \
+ --num-epochs 30 \
+ --start-epoch 1 \
--full-libri 1 \
- --exp-dir pruned_transducer_stateless5/exp \
+ --exp-dir pruned_transducer_stateless5/exp-B \
--max-duration 300 \
--use-fp16 0 \
--num-encoder-layers 24 \
@@ -230,19 +1333,16 @@ export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
```
The tensorboard log can be found at
-
-
-**Caution**: The training script is updated so that epochs are counted from 1
-after the training.
+
The decoding commands are:
```bash
for method in greedy_search modified_beam_search fast_beam_search; do
./pruned_transducer_stateless5/decode.py \
- --epoch 31 \
- --avg 17 \
- --exp-dir ./pruned_transducer_stateless5/exp-M \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir ./pruned_transducer_stateless5/exp-B \
--max-duration 600 \
--decoding-method $method \
--max-sym-per-frame 1 \
@@ -251,24 +1351,25 @@ for method in greedy_search modified_beam_search fast_beam_search; do
--nhead 8 \
--encoder-dim 384 \
--decoder-dim 512 \
- --joiner-dim 512
+ --joiner-dim 512 \
+ --use-averaged-model True
done
```
You can find a pretrained model, training logs, decoding logs, and decoding
results at:
-
+
### LibriSpeech BPE training results (Pruned Stateless Transducer 4)
-[pruned_transducer_stateless4](./pruned_transducer_stateless4)
+#### [pruned_transducer_stateless4](./pruned_transducer_stateless4)
This version saves averaged model during training, and decodes with averaged model.
See for details about the idea of model averaging.
-#### Training on full librispeech
+##### Training on full librispeech
See
@@ -278,12 +1379,12 @@ The WERs are:
| | test-clean | test-other | comment |
|-------------------------------------|------------|------------|-------------------------------------------------------------------------------|
-| greedy search (max sym per frame 1) | 2.75 | 6.74 | --epoch 30 --avg 6 --use_averaged_model False |
-| greedy search (max sym per frame 1) | 2.69 | 6.64 | --epoch 30 --avg 6 --use_averaged_model True |
-| fast beam search | 2.72 | 6.67 | --epoch 30 --avg 6 --use_averaged_model False |
-| fast beam search | 2.66 | 6.6 | --epoch 30 --avg 6 --use_averaged_model True |
-| modified beam search | 2.67 | 6.68 | --epoch 30 --avg 6 --use_averaged_model False |
-| modified beam search | 2.62 | 6.57 | --epoch 30 --avg 6 --use_averaged_model True |
+| greedy search (max sym per frame 1) | 2.75 | 6.74 | --epoch 30 --avg 6 --use-averaged-model False |
+| greedy search (max sym per frame 1) | 2.69 | 6.64 | --epoch 30 --avg 6 --use-averaged-model True |
+| fast beam search | 2.72 | 6.67 | --epoch 30 --avg 6 --use-averaged-model False |
+| fast beam search | 2.66 | 6.6 | --epoch 30 --avg 6 --use-averaged-model True |
+| modified beam search | 2.67 | 6.68 | --epoch 30 --avg 6 --use-averaged-model False |
+| modified beam search | 2.62 | 6.57 | --epoch 30 --avg 6 --use-averaged-model True |
The training command is:
@@ -344,7 +1445,7 @@ Pretrained models, training logs, decoding logs, and decoding results
are available at
-#### Training on train-clean-100
+##### Training on train-clean-100
See
@@ -381,7 +1482,7 @@ The tensorboard log can be found at
### LibriSpeech BPE training results (Pruned Stateless Transducer 3, 2022-04-29)
-[pruned_transducer_stateless3](./pruned_transducer_stateless3)
+#### [pruned_transducer_stateless3](./pruned_transducer_stateless3)
Same as `Pruned Stateless Transducer 2` but using the XL subset from
[GigaSpeech](https://github.com/SpeechColab/GigaSpeech) as extra training data.
@@ -595,10 +1696,10 @@ can be found at
### LibriSpeech BPE training results (Pruned Transducer 2)
-[pruned_transducer_stateless2](./pruned_transducer_stateless2)
+#### [pruned_transducer_stateless2](./pruned_transducer_stateless2)
This is with a reworked version of the conformer encoder, with many changes.
-#### Training on fulll librispeech
+##### Training on full librispeech
Using commit `34aad74a2c849542dd5f6359c9e6b527e8782fd6`.
See
@@ -619,9 +1720,25 @@ The WERs are:
The train and decode commands are:
-`python3 ./pruned_transducer_stateless2/train.py --exp-dir=pruned_transducer_stateless2/exp --world-size 8 --num-epochs 26 --full-libri 1 --max-duration 300`
+```bash
+python3 ./pruned_transducer_stateless2/train.py \
+ --exp-dir=pruned_transducer_stateless2/exp \
+ --world-size 8 \
+ --num-epochs 26 \
+ --full-libri 1 \
+ --max-duration 300
+```
+
and:
-`python3 ./pruned_transducer_stateless2/decode.py --exp-dir pruned_transducer_stateless2/exp --epoch 25 --avg 8 --bpe-model ./data/lang_bpe_500/bpe.model --max-duration 600`
+
+```bash
+python3 ./pruned_transducer_stateless2/decode.py \
+ --exp-dir pruned_transducer_stateless2/exp \
+ --epoch 25 \
+ --avg 8 \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --max-duration 600
+```
The Tensorboard log is at (apologies, log starts
only from epoch 3).
@@ -631,12 +1748,29 @@ can be found at
-#### Training on train-clean-100:
+##### Training on train-clean-100:
Trained with 1 job:
-`python3 ./pruned_transducer_stateless2/train.py --exp-dir=pruned_transducer_stateless2/exp_100h_ws1 --world-size 1 --num-epochs 40 --full-libri 0 --max-duration 300`
+```
+python3 ./pruned_transducer_stateless2/train.py \
+ --exp-dir=pruned_transducer_stateless2/exp_100h_ws1 \
+ --world-size 1 \
+ --num-epochs 40 \
+ --full-libri 0 \
+ --max-duration 300
+```
+
and decoded with:
-`python3 ./pruned_transducer_stateless2/decode.py --exp-dir pruned_transducer_stateless2/exp_100h_ws1 --epoch 19 --avg 8 --bpe-model ./data/lang_bpe_500/bpe.model --max-duration 600`.
+
+```
+python3 ./pruned_transducer_stateless2/decode.py \
+ --exp-dir pruned_transducer_stateless2/exp_100h_ws1 \
+ --epoch 19 \
+ --avg 8 \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --max-duration 600
+```
+
The Tensorboard log is at (learning rate
schedule is not visible due to a since-fixed bug).
@@ -650,9 +1784,26 @@ schedule is not visible due to a since-fixed bug).
| fast beam search | 6.53 | 16.82 | --epoch 39 --avg 10 --decoding-method fast_beam_search |
Trained with 2 jobs:
-`python3 ./pruned_transducer_stateless2/train.py --exp-dir=pruned_transducer_stateless2/exp_100h_ws2 --world-size 2 --num-epochs 40 --full-libri 0 --max-duration 300`
+
+```bash
+python3 ./pruned_transducer_stateless2/train.py \
+ --exp-dir=pruned_transducer_stateless2/exp_100h_ws2 \
+ --world-size 2 \
+ --num-epochs 40 \
+ --full-libri 0 \
+ --max-duration 300
+```
+
and decoded with:
-`python3 ./pruned_transducer_stateless2/decode.py --exp-dir pruned_transducer_stateless2/exp_100h_ws2 --epoch 19 --avg 8 --bpe-model ./data/lang_bpe_500/bpe.model --max-duration 600`.
+
+```
+python3 ./pruned_transducer_stateless2/decode.py \
+ --exp-dir pruned_transducer_stateless2/exp_100h_ws2 \
+ --epoch 19 \
+ --avg 8 \
+ --bpe-model ./data/lang_bpe_500/bpe.model \
+ --max-duration 600
+```
The Tensorboard log is at
(learning rate schedule is not visible due to a since-fixed bug).
@@ -665,9 +1816,26 @@ The Tensorboard log is at
@@ -684,7 +1852,16 @@ The Tensorboard log is at .
Train command was
-`python3 ./pruned_transducer_stateless2/train.py --exp-dir=pruned_transducer_stateless2/exp_100h_fp16 --world-size 1 --num-epochs 40 --full-libri 0 --max-duration 300 --use-fp16 True`
+
+```
+python3 ./pruned_transducer_stateless2/train.py \
+ --exp-dir=pruned_transducer_stateless2/exp_100h_fp16 \
+ --world-size 1 \
+ --num-epochs 40 \
+ --full-libri 0 \
+ --max-duration 300 \
+ --use-fp16 True
+```
The Tensorboard log is at
@@ -698,7 +1875,16 @@ The Tensorboard log is at .
Train command was
-`python3 ./pruned_transducer_stateless2/train.py --exp-dir=pruned_transducer_stateless2/exp_100h_fp16 --world-size 1 --num-epochs 40 --full-libri 0 --max-duration 500 --use-fp16 True`
+
+```
+python3 ./pruned_transducer_stateless2/train.py \
+ --exp-dir=pruned_transducer_stateless2/exp_100h_fp16 \
+ --world-size 1 \
+ --num-epochs 40 \
+ --full-libri 0 \
+ --max-duration 500 \
+ --use-fp16 True
+```
The Tensorboard log is at
@@ -710,7 +1896,6 @@ The Tensorboard log is at
+### LibriSpeech BPE training results (Conformer-CTC 2)
+
+#### [conformer_ctc2](./conformer_ctc2)
+
+#### 2022-07-21
+
+It implements a 'reworked' version of CTC attention model.
+As demenstrated by pruned_transducer_stateless2, reworked Conformer model has superior performance compared to the original Conformer.
+So in this modified version of CTC attention model, it has the reworked Conformer as the encoder and the reworked Transformer as the decoder.
+conformer_ctc2 also integrates with the idea of the 'averaging models' in pruned_transducer_stateless4.
+
+The WERs on comparisons with a baseline model, for the librispeech test dataset, are listed as below.
+
+The baseline model is the original conformer CTC attention model trained with icefall/egs/librispeech/ASR/conformer_ctc.
+The model is downloaded from .
+This model has 12 layers of Conformer encoder layers and 6 Transformer decoder layers.
+Number of model parameters is 109,226,120.
+It has been trained with 90 epochs with full Librispeech dataset.
+
+For this reworked CTC attention model, it has 12 layers of reworked Conformer encoder layers and 6 reworked Transformer decoder layers.
+Number of model parameters is 103,071,035.
+With full Librispeech data set, it was trained for **only** 30 epochs because the reworked model would converge much faster.
+Please refer to to see the loss convergence curve.
+Please find the above trained model at in huggingface.
+
+The decoding configuration for the reworked model is --epoch 30, --avg 8, --use-averaged-model True, which is the best after searching.
+
+| WER | reworked ctc attention | with --epoch 30 --avg 8 --use-averaged-model True | | ctc attention| with --epoch 77 --avg 55 | |
+|------------------------|-------|------|------|------|------|-----|
+| test sets | test-clean | test-other | Avg | test-clean | test-other | Avg |
+| ctc-greedy-search | 2.98% | 7.14%| 5.06%| 2.90%| 7.47%| 5.19%|
+| ctc-decoding | 2.98% | 7.14%| 5.06%| 2.90%| 7.47%| 5.19%|
+| 1best | 2.93% | 6.37%| 4.65%| 2.70%| 6.49%| 4.60%|
+| nbest | 2.94% | 6.39%| 4.67%| 2.70%| 6.48%| 4.59%|
+| nbest-rescoring | 2.68% | 5.77%| 4.23%| 2.55%| 6.07%| 4.31%|
+| whole-lattice-rescoring| 2.66% | 5.76%| 4.21%| 2.56%| 6.04%| 4.30%|
+| attention-decoder | 2.59% | 5.54%| 4.07%| 2.41%| 5.77%| 4.09%|
+| nbest-oracle | 1.53% | 3.47%| 2.50%| 1.69%| 4.02%| 2.86%|
+| rnn-lm | 2.37% | 4.98%| 3.68%| 2.31%| 5.35%| 3.83%|
+
+
+
+conformer_ctc2 also implements the CTC greedy search decoding, it has the identical WERs with the CTC-decoding method.
+For other decoding methods, the average WER of the two test sets with the two models is similar.
+Except for the 1best and nbest methods, the overall performance of reworked model is better than the baseline model.
+
+
+To reproduce the above result, use the following commands.
+
+The training commands are:
+
+```bash
+ WORLD_SIZE=8
+ export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
+ ./conformer_ctc2/train.py \
+ --manifest-dir data/fbank \
+ --exp-dir conformer_ctc2/exp \
+ --full-libri 1 \
+ --spec-aug-time-warp-factor 80 \
+ --max-duration 300 \
+ --world-size ${WORLD_SIZE} \
+ --start-epoch 1 \
+ --num-epochs 30 \
+ --att-rate 0.7 \
+ --num-decoder-layers 6
+```
+
+
+And the following commands are for decoding:
+
+```bash
+
+
+for method in ctc-greedy-search ctc-decoding 1best nbest-oracle; do
+ python3 ./conformer_ctc2/decode.py \
+ --exp-dir conformer_ctc2/exp \
+ --use-averaged-model True --epoch 30 --avg 8 --max-duration 200 --method $method
+done
+
+for method in nbest nbest-rescoring whole-lattice-rescoring attention-decoder ; do
+ python3 ./conformer_ctc2/decode.py \
+ --exp-dir conformer_ctc2/exp \
+ --use-averaged-model True --epoch 30 --avg 8 --max-duration 20 --method $method
+done
+
+rnn_dir=$(git rev-parse --show-toplevel)/icefall/rnn_lm
+./conformer_ctc2/decode.py \
+ --exp-dir conformer_ctc2/exp \
+ --lang-dir data/lang_bpe_500 \
+ --lm-dir data/lm \
+ --max-duration 30 \
+ --concatenate-cuts 0 \
+ --bucketing-sampler 1 \
+ --num-paths 1000 \
+ --use-averaged-model True \
+ --epoch 30 \
+ --avg 8 \
+ --nbest-scale 0.5 \
+ --rnn-lm-exp-dir ${rnn_dir}/exp \
+ --rnn-lm-epoch 29 \
+ --rnn-lm-avg 3 \
+ --rnn-lm-embedding-dim 2048 \
+ --rnn-lm-hidden-dim 2048 \
+ --rnn-lm-num-layers 3 \
+ --rnn-lm-tie-weights true \
+ --method rnn-lm
+```
+
+You can find the RNN-LM pre-trained model at
+
+
+
### LibriSpeech BPE training results (Conformer-CTC)
#### 2021-11-09
-The best WER, as of 2021-11-09, for the librispeech test dataset is below
-(using HLG decoding + n-gram LM rescoring + attention decoder rescoring):
+The best WER, as of 2022-06-20, for the librispeech test dataset is below
+(using HLG decoding + n-gram LM rescoring + attention decoder rescoring + rnn lm rescoring):
| | test-clean | test-other |
|-----|------------|------------|
-| WER | 2.42 | 5.73 |
+| WER | 2.32 | 5.39 |
Scale values used in n-gram LM rescoring and attention rescoring for the best WERs are:
-| ngram_lm_scale | attention_scale |
-|----------------|-----------------|
-| 2.0 | 2.0 |
+
+| ngram_lm_scale | attention_scale | rnn_lm_scale |
+|----------------|-----------------|--------------|
+| 0.3 | 2.1 | 2.2 |
To reproduce the above result, use the following commands for training:
@@ -1168,11 +2466,27 @@ export CUDA_VISIBLE_DEVICES="0,1,2,3"
--start-epoch 0 \
--num-epochs 90
# Note: It trains for 90 epochs, but the best WER is at epoch-77.pt
+
+# Train the RNN-LM
+cd icefall
+export CUDA_VISIBLE_DEVICES="0,1,2,3"
+./rnn_lm/train.py \
+ --exp-dir rnn_lm/exp_2048_3_tied \
+ --start-epoch 0 \
+ --world-size 4 \
+ --num-epochs 30 \
+ --use-fp16 1 \
+ --embedding-dim 2048 \
+ --hidden-dim 2048 \
+ --num-layers 3 \
+ --batch-size 500 \
+ --tie-weights true
```
and the following command for decoding
```
+rnn_dir=$(git rev-parse --show-toplevel)/icefall/rnn_lm
./conformer_ctc/decode.py \
--exp-dir conformer_ctc/exp_500_att0.8 \
--lang-dir data/lang_bpe_500 \
@@ -1182,13 +2496,23 @@ and the following command for decoding
--num-paths 1000 \
--epoch 77 \
--avg 55 \
- --method attention-decoder \
- --nbest-scale 0.5
+ --nbest-scale 0.5 \
+ --rnn-lm-exp-dir ${rnn_dir}/exp_2048_3_tied \
+ --rnn-lm-epoch 29 \
+ --rnn-lm-avg 3 \
+ --rnn-lm-embedding-dim 2048 \
+ --rnn-lm-hidden-dim 2048 \
+ --rnn-lm-num-layers 3 \
+ --rnn-lm-tie-weights true \
+ --method rnn-lm
```
-You can find the pre-trained model by visiting
+You can find the Conformer-CTC pre-trained model by visiting
+and the RNN-LM pre-trained model:
+
+
The tensorboard log for training is available at
diff --git a/egs/librispeech/ASR/conformer_ctc/ali.py b/egs/librispeech/ASR/conformer_ctc/ali.py
index 42fa2308e..2828e309e 100755
--- a/egs/librispeech/ASR/conformer_ctc/ali.py
+++ b/egs/librispeech/ASR/conformer_ctc/ali.py
@@ -96,14 +96,14 @@ def get_parser():
- labels_xxx.h5
- aux_labels_xxx.h5
- - cuts_xxx.json.gz
+ - librispeech_cuts_xxx.jsonl.gz
where xxx is the value of `--dataset`. For instance, if
`--dataset` is `train-clean-100`, it will contain 3 files:
- `labels_train-clean-100.h5`
- `aux_labels_train-clean-100.h5`
- - `cuts_train-clean-100.json.gz`
+ - `librispeech_cuts_train-clean-100.jsonl.gz`
Note: Both labels_xxx.h5 and aux_labels_xxx.h5 contain framewise
alignment. The difference is that labels_xxx.h5 contains repeats.
@@ -289,7 +289,9 @@ def main():
out_labels_ali_filename = out_dir / f"labels_{params.dataset}.h5"
out_aux_labels_ali_filename = out_dir / f"aux_labels_{params.dataset}.h5"
- out_manifest_filename = out_dir / f"cuts_{params.dataset}.json.gz"
+ out_manifest_filename = (
+ out_dir / f"librispeech_cuts_{params.dataset}.jsonl.gz"
+ )
for f in (
out_labels_ali_filename,
diff --git a/egs/librispeech/ASR/conformer_ctc/conformer.py b/egs/librispeech/ASR/conformer_ctc/conformer.py
index 871712a46..6fac07f93 100644
--- a/egs/librispeech/ASR/conformer_ctc/conformer.py
+++ b/egs/librispeech/ASR/conformer_ctc/conformer.py
@@ -253,7 +253,9 @@ class ConformerEncoderLayer(nn.Module):
residual = src
if self.normalize_before:
src = self.norm_conv(src)
- src = residual + self.dropout(self.conv_module(src))
+ src = residual + self.dropout(
+ self.conv_module(src, src_key_padding_mask=src_key_padding_mask)
+ )
if not self.normalize_before:
src = self.norm_conv(src)
@@ -369,7 +371,7 @@ class RelPositionalEncoding(torch.nn.Module):
):
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
+ # 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 Tensor:
+ 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).
Returns:
Tensor: Output tensor (#time, batch, channels).
@@ -908,6 +915,8 @@ class ConvolutionModule(nn.Module):
x = nn.functional.glu(x, dim=1) # (batch, channels, time)
# 1D Depthwise Conv
+ if src_key_padding_mask is not None:
+ x.masked_fill_(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
x = self.depthwise_conv(x)
if self.use_batchnorm:
x = self.norm(x)
diff --git a/egs/librispeech/ASR/conformer_ctc/decode.py b/egs/librispeech/ASR/conformer_ctc/decode.py
index 177e33a6e..3f3b1acda 100755
--- a/egs/librispeech/ASR/conformer_ctc/decode.py
+++ b/egs/librispeech/ASR/conformer_ctc/decode.py
@@ -30,7 +30,7 @@ from asr_datamodule import LibriSpeechAsrDataModule
from conformer import Conformer
from icefall.bpe_graph_compiler import BpeCtcTrainingGraphCompiler
-from icefall.checkpoint import average_checkpoints, load_checkpoint
+from icefall.checkpoint import load_checkpoint
from icefall.decode import (
get_lattice,
nbest_decoding,
@@ -38,15 +38,19 @@ from icefall.decode import (
one_best_decoding,
rescore_with_attention_decoder,
rescore_with_n_best_list,
+ rescore_with_rnn_lm,
rescore_with_whole_lattice,
)
from icefall.env import get_env_info
from icefall.lexicon import Lexicon
+from icefall.rnn_lm.model import RnnLmModel
from icefall.utils import (
AttributeDict,
get_texts,
+ load_averaged_model,
setup_logger,
store_transcripts,
+ str2bool,
write_error_stats,
)
@@ -93,7 +97,9 @@ def get_parser():
is the decoding result.
- (5) attention-decoder. Extract n paths from the LM rescored
lattice, the path with the highest score is the decoding result.
- - (6) nbest-oracle. Its WER is the lower bound of any n-best
+ - (6) rnn-lm. Rescoring with attention-decoder and RNN LM. We assume
+ you have trained an RNN LM using ./rnn_lm/train.py
+ - (7) nbest-oracle. Its WER is the lower bound of any n-best
rescoring method can achieve. Useful for debugging n-best
rescoring method.
""",
@@ -105,7 +111,7 @@ def get_parser():
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
+ nbest, nbest-rescoring, attention-decoder, rnn-lm, and nbest-oracle
""",
)
@@ -116,7 +122,7 @@ def get_parser():
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
+ nbest, nbest-rescoring, attention-decoder, rnn-lm, and nbest-oracle
A smaller value results in more unique paths.
""",
)
@@ -139,11 +145,67 @@ def get_parser():
"--lm-dir",
type=str,
default="data/lm",
- help="""The LM dir.
+ help="""The n-gram LM dir.
It should contain either G_4_gram.pt or G_4_gram.fst.txt
""",
)
+ 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",
+ type=str2bool,
+ default=False,
+ help="""True to share the weights between the input embedding layer and the
+ last output linear layer
+ """,
+ )
+
return parser
@@ -173,6 +235,7 @@ def get_params() -> AttributeDict:
def decode_one_batch(
params: AttributeDict,
model: nn.Module,
+ rnn_lm_model: Optional[nn.Module],
HLG: Optional[k2.Fsa],
H: Optional[k2.Fsa],
bpe_model: Optional[spm.SentencePieceProcessor],
@@ -205,6 +268,8 @@ def decode_one_batch(
model:
The neural model.
+ rnn_lm_model:
+ The neural model for RNN LM.
HLG:
The decoding graph. Used only when params.method is NOT ctc-decoding.
H:
@@ -330,6 +395,7 @@ def decode_one_batch(
"nbest-rescoring",
"whole-lattice-rescoring",
"attention-decoder",
+ "rnn-lm",
]
lm_scale_list = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]
@@ -357,8 +423,6 @@ def decode_one_batch(
G_with_epsilon_loops=G,
lm_scale_list=None,
)
- # TODO: pass `lattice` instead of `rescored_lattice` to
- # `rescore_with_attention_decoder`
best_path_dict = rescore_with_attention_decoder(
lattice=rescored_lattice,
@@ -370,6 +434,26 @@ def decode_one_batch(
eos_id=eos_id,
nbest_scale=params.nbest_scale,
)
+ elif params.method == "rnn-lm":
+ # lattice uses a 3-gram Lm. We rescore it with a 4-gram LM.
+ rescored_lattice = rescore_with_whole_lattice(
+ lattice=lattice,
+ G_with_epsilon_loops=G,
+ lm_scale_list=None,
+ )
+
+ best_path_dict = rescore_with_rnn_lm(
+ lattice=rescored_lattice,
+ num_paths=params.num_paths,
+ rnn_lm_model=rnn_lm_model,
+ model=model,
+ memory=memory,
+ memory_key_padding_mask=memory_key_padding_mask,
+ sos_id=sos_id,
+ eos_id=eos_id,
+ blank_id=0,
+ nbest_scale=params.nbest_scale,
+ )
else:
assert False, f"Unsupported decoding method: {params.method}"
@@ -388,6 +472,7 @@ def decode_dataset(
dl: torch.utils.data.DataLoader,
params: AttributeDict,
model: nn.Module,
+ rnn_lm_model: Optional[nn.Module],
HLG: Optional[k2.Fsa],
H: Optional[k2.Fsa],
bpe_model: Optional[spm.SentencePieceProcessor],
@@ -395,7 +480,7 @@ def decode_dataset(
sos_id: int,
eos_id: int,
G: Optional[k2.Fsa] = None,
-) -> Dict[str, List[Tuple[List[str], List[str]]]]:
+) -> Dict[str, List[Tuple[str, List[str], List[str]]]]:
"""Decode dataset.
Args:
@@ -405,6 +490,8 @@ def decode_dataset(
It is returned by :func:`get_params`.
model:
The neural model.
+ rnn_lm_model:
+ The neural model for RNN LM.
HLG:
The decoding graph. Used only when params.method is NOT ctc-decoding.
H:
@@ -438,10 +525,12 @@ def decode_dataset(
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,
+ rnn_lm_model=rnn_lm_model,
HLG=HLG,
H=H,
bpe_model=bpe_model,
@@ -456,9 +545,9 @@ def decode_dataset(
for lm_scale, hyps in hyps_dict.items():
this_batch = []
assert len(hyps) == len(texts)
- for hyp_words, ref_text in zip(hyps, texts):
+ for cut_id, hyp_words, ref_text in zip(cut_ids, hyps, texts):
ref_words = ref_text.split()
- this_batch.append((ref_words, hyp_words))
+ this_batch.append((cut_id, ref_words, hyp_words))
results[lm_scale].extend(this_batch)
else:
@@ -488,9 +577,9 @@ def decode_dataset(
def save_results(
params: AttributeDict,
test_set_name: str,
- results_dict: Dict[str, List[Tuple[List[int], List[int]]]],
+ results_dict: Dict[str, List[Tuple[str, List[int], List[int]]]],
):
- if params.method == "attention-decoder":
+ if params.method in ("attention-decoder", "rnn-lm"):
# Set it to False since there are too many logs.
enable_log = False
else:
@@ -498,6 +587,7 @@ def save_results(
test_set_wers = dict()
for key, results in results_dict.items():
recog_path = params.exp_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}")
@@ -566,6 +656,10 @@ def main():
sos_id = graph_compiler.sos_id
eos_id = graph_compiler.eos_id
+ params.num_classes = num_classes
+ params.sos_id = sos_id
+ params.eos_id = eos_id
+
if params.method == "ctc-decoding":
HLG = None
H = k2.ctc_topo(
@@ -590,6 +684,7 @@ def main():
"nbest-rescoring",
"whole-lattice-rescoring",
"attention-decoder",
+ "rnn-lm",
):
if not (params.lm_dir / "G_4_gram.pt").is_file():
logging.info("Loading G_4_gram.fst.txt")
@@ -621,7 +716,11 @@ def main():
d = torch.load(params.lm_dir / "G_4_gram.pt", map_location=device)
G = k2.Fsa.from_dict(d)
- if params.method in ["whole-lattice-rescoring", "attention-decoder"]:
+ if params.method in [
+ "whole-lattice-rescoring",
+ "attention-decoder",
+ "rnn-lm",
+ ]:
# Add epsilon self-loops to G as we will compose
# it with the whole lattice later
G = k2.add_epsilon_self_loops(G)
@@ -648,20 +747,42 @@ def main():
if params.avg == 1:
load_checkpoint(f"{params.exp_dir}/epoch-{params.epoch}.pt", model)
else:
- start = params.epoch - params.avg + 1
- filenames = []
- for i in range(start, params.epoch + 1):
- if start >= 0:
- filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
- logging.info(f"averaging {filenames}")
- model.to(device)
- model.load_state_dict(average_checkpoints(filenames, device=device))
+ model = load_averaged_model(
+ params.exp_dir, model, params.epoch, params.avg, 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}")
+ rnn_lm_model = None
+ if params.method == "rnn-lm":
+ rnn_lm_model = RnnLmModel(
+ vocab_size=params.num_classes,
+ 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.rnn_lm_avg == 1:
+ load_checkpoint(
+ f"{params.rnn_lm_exp_dir}/epoch-{params.rnn_lm_epoch}.pt",
+ rnn_lm_model,
+ )
+ rnn_lm_model.to(device)
+ else:
+ rnn_lm_model = load_averaged_model(
+ params.rnn_lm_exp_dir,
+ rnn_lm_model,
+ params.rnn_lm_epoch,
+ params.rnn_lm_avg,
+ device,
+ )
+ rnn_lm_model.eval()
+
+ # we need cut ids to display recognition results.
+ args.return_cuts = True
librispeech = LibriSpeechAsrDataModule(args)
test_clean_cuts = librispeech.test_clean_cuts()
@@ -678,6 +799,7 @@ def main():
dl=test_dl,
params=params,
model=model,
+ rnn_lm_model=rnn_lm_model,
HLG=HLG,
H=H,
bpe_model=bpe_model,
diff --git a/egs/librispeech/ASR/conformer_ctc/pretrained.py b/egs/librispeech/ASR/conformer_ctc/pretrained.py
index 28724e1eb..a2c0a5486 100755
--- a/egs/librispeech/ASR/conformer_ctc/pretrained.py
+++ b/egs/librispeech/ASR/conformer_ctc/pretrained.py
@@ -322,7 +322,7 @@ def main():
H = k2.ctc_topo(
max_token=max_token_id,
- modified=False,
+ modified=params.num_classes > 500,
device=device,
)
diff --git a/egs/librispeech/ASR/conformer_ctc/train.py b/egs/librispeech/ASR/conformer_ctc/train.py
index b81bd6330..6419f6816 100755
--- a/egs/librispeech/ASR/conformer_ctc/train.py
+++ b/egs/librispeech/ASR/conformer_ctc/train.py
@@ -17,6 +17,17 @@
# 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 \
+ --exp-dir ./conformer_ctc/exp \
+ --world-size 4 \
+ --full-libri 1 \
+ --max-duration 200 \
+ --num-epochs 20
+"""
+
import argparse
import logging
from pathlib import Path
@@ -29,6 +40,7 @@ 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.utils import fix_random_seed
from torch import Tensor
from torch.nn.parallel import DistributedDataParallel as DDP
@@ -435,6 +447,17 @@ def compute_loss(
info["loss"] = loss.detach().cpu().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"] = supervisions["num_frames"].sum().item()
+ # averaged padding proportion over utterances
+ info["utt_pad_proportion"] = (
+ ((feature.size(1) - supervisions["num_frames"]) / feature.size(1))
+ .sum()
+ .item()
+ )
+
return loss, info
@@ -676,6 +699,20 @@ def run(rank, world_size, args):
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()
diff --git a/egs/librispeech/ASR/conformer_ctc2/__init__.py b/egs/librispeech/ASR/conformer_ctc2/__init__.py
new file mode 120000
index 000000000..b24e5e357
--- /dev/null
+++ b/egs/librispeech/ASR/conformer_ctc2/__init__.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/__init__.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conformer_ctc2/asr_datamodule.py b/egs/librispeech/ASR/conformer_ctc2/asr_datamodule.py
new file mode 120000
index 000000000..a074d6085
--- /dev/null
+++ b/egs/librispeech/ASR/conformer_ctc2/asr_datamodule.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/asr_datamodule.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conformer_ctc2/attention.py b/egs/librispeech/ASR/conformer_ctc2/attention.py
new file mode 100644
index 000000000..1375d7245
--- /dev/null
+++ b/egs/librispeech/ASR/conformer_ctc2/attention.py
@@ -0,0 +1,252 @@
+# 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.
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn.init import xavier_normal_
+
+from scaling import ScaledLinear
+
+
+class MultiheadAttention(nn.Module):
+ r"""Allows the model to jointly attend to information
+ from different representation subspaces.
+ See `Attention Is All You Need `_.
+
+ .. math::
+ \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
+
+ where :math:`head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`.
+
+ 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`` (no dropout).
+ 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``.
+ kdim: Total number of features for keys. Default: ``None`` (uses ``kdim=embed_dim``).
+ vdim: Total number of features for values. Default: ``None`` (uses ``vdim=embed_dim``).
+ batch_first: If ``True``, then the input and output tensors are provided
+ as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
+
+ Examples::
+
+ >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
+ >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
+ """
+ __constants__ = ["batch_first"]
+ bias_k: Optional[torch.Tensor]
+ bias_v: Optional[torch.Tensor]
+
+ def __init__(
+ self,
+ embed_dim,
+ num_heads,
+ dropout=0.0,
+ bias=True,
+ add_bias_kv=False,
+ add_zero_attn=False,
+ kdim=None,
+ vdim=None,
+ batch_first=False,
+ device=None,
+ dtype=None,
+ ) -> None:
+ factory_kwargs = {"device": device, "dtype": dtype}
+ super(MultiheadAttention, self).__init__()
+ self.embed_dim = embed_dim
+ self.kdim = kdim if kdim is not None else embed_dim
+ self.vdim = vdim if vdim is not None else embed_dim
+ self._qkv_same_embed_dim = (
+ self.kdim == embed_dim and self.vdim == embed_dim
+ )
+
+ self.num_heads = num_heads
+ self.dropout = dropout
+ self.batch_first = batch_first
+ self.head_dim = embed_dim // num_heads
+ assert (
+ self.head_dim * num_heads == self.embed_dim
+ ), "embed_dim must be divisible by num_heads"
+
+ if self._qkv_same_embed_dim is False:
+ self.q_proj_weight = ScaledLinear(embed_dim, embed_dim, bias=bias)
+ self.k_proj_weight = ScaledLinear(self.kdim, embed_dim, bias=bias)
+ self.v_proj_weight = ScaledLinear(self.vdim, embed_dim, bias=bias)
+ self.register_parameter("in_proj_weight", None)
+ else:
+ self.in_proj_weight = ScaledLinear(
+ embed_dim, 3 * embed_dim, bias=bias
+ )
+ self.register_parameter("q_proj_weight", None)
+ self.register_parameter("k_proj_weight", None)
+ self.register_parameter("v_proj_weight", None)
+
+ if not bias:
+ self.register_parameter("in_proj_bias", None)
+
+ self.out_proj = ScaledLinear(embed_dim, embed_dim, bias=bias)
+
+ if add_bias_kv:
+ self.bias_k = nn.Parameter(
+ torch.empty((1, 1, embed_dim), **factory_kwargs)
+ )
+ self.bias_v = nn.Parameter(
+ torch.empty((1, 1, embed_dim), **factory_kwargs)
+ )
+ else:
+ self.bias_k = self.bias_v = None
+
+ self.add_zero_attn = add_zero_attn
+
+ self._reset_parameters()
+
+ def _reset_parameters(self):
+ if self.bias_k is not None:
+ xavier_normal_(self.bias_k)
+ if self.bias_v is not None:
+ xavier_normal_(self.bias_v)
+
+ def __setstate__(self, state):
+ # Support loading old MultiheadAttention checkpoints generated by v1.1.0
+ if "_qkv_same_embed_dim" not in state:
+ state["_qkv_same_embed_dim"] = True
+
+ super(MultiheadAttention, self).__setstate__(state)
+
+ def forward(
+ self,
+ query: Tensor,
+ key: Tensor,
+ value: Tensor,
+ key_padding_mask: Optional[Tensor] = None,
+ need_weights: bool = True,
+ attn_mask: Optional[Tensor] = None,
+ ) -> Tuple[Tensor, Optional[Tensor]]:
+ r"""
+ Args:
+ query: Query embeddings of shape :math:`(L, N, E_q)` when ``batch_first=False`` or :math:`(N, L, E_q)`
+ when ``batch_first=True``, where :math:`L` is the target sequence length, :math:`N` is the batch size,
+ and :math:`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 :math:`(S, N, E_k)` when ``batch_first=False`` or :math:`(N, S, E_k)` when
+ ``batch_first=True``, where :math:`S` is the source sequence length, :math:`N` is the batch size, and
+ :math:`E_k` is the key embedding dimension ``kdim``. See "Attention Is All You Need" for more details.
+ value: Value embeddings of shape :math:`(S, N, E_v)` when ``batch_first=False`` or :math:`(N, S, E_v)` when
+ ``batch_first=True``, where :math:`S` is the source sequence length, :math:`N` is the batch size, and
+ :math:`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 :math:`(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 specified, 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
+ :math:`(L, S)` or :math:`(N\cdot\text{num\_heads}, L, S)`, where :math:`N` is the batch size,
+ :math:`L` is the target sequence length, and :math:`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.
+
+ Outputs:
+ - **attn_output** - Attention outputs of shape :math:`(L, N, E)` when ``batch_first=False`` or
+ :math:`(N, L, E)` when ``batch_first=True``, where :math:`L` is the target sequence length, :math:`N` is
+ the batch size, and :math:`E` is the embedding dimension ``embed_dim``.
+ - **attn_output_weights** - Attention output weights of shape :math:`(N, L, S)`, where :math:`N` is the batch
+ size, :math:`L` is the target sequence length, and :math:`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)]
+
+ if not self._qkv_same_embed_dim:
+ q_proj_weight = (
+ self.q_proj_weight.get_weight()
+ if self.q_proj_weight is not None
+ else None
+ )
+ k_proj_weight = (
+ self.k_proj_weight.get_weight()
+ if self.k_proj_weight is not None
+ else None
+ )
+ v_proj_weight = (
+ self.v_proj_weight.get_weight()
+ if self.v_proj_weight is not None
+ else None
+ )
+ (
+ attn_output,
+ attn_output_weights,
+ ) = nn.functional.multi_head_attention_forward(
+ query,
+ key,
+ value,
+ self.embed_dim,
+ self.num_heads,
+ self.in_proj_weight.get_weight(),
+ self.in_proj_weight.get_bias(),
+ self.bias_k,
+ self.bias_v,
+ self.add_zero_attn,
+ 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,
+ use_separate_proj_weight=True,
+ q_proj_weight=q_proj_weight,
+ k_proj_weight=k_proj_weight,
+ v_proj_weight=v_proj_weight,
+ )
+ else:
+ (
+ attn_output,
+ attn_output_weights,
+ ) = nn.functional.multi_head_attention_forward(
+ query,
+ key,
+ value,
+ self.embed_dim,
+ self.num_heads,
+ self.in_proj_weight.get_weight(),
+ self.in_proj_weight.get_bias(),
+ self.bias_k,
+ self.bias_v,
+ self.add_zero_attn,
+ 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,
+ )
+ if self.batch_first:
+ return attn_output.transpose(1, 0), attn_output_weights
+ else:
+ return attn_output, attn_output_weights
diff --git a/egs/librispeech/ASR/conformer_ctc2/conformer.py b/egs/librispeech/ASR/conformer_ctc2/conformer.py
new file mode 100644
index 000000000..b906d2650
--- /dev/null
+++ b/egs/librispeech/ASR/conformer_ctc2/conformer.py
@@ -0,0 +1,973 @@
+#!/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 Optional, Tuple
+
+import torch
+from scaling import (
+ ActivationBalancer,
+ BasicNorm,
+ DoubleSwish,
+ ScaledConv1d,
+ ScaledLinear,
+)
+from torch import Tensor, nn
+from subsampling import Conv2dSubsampling
+
+from transformer import Supervisions, Transformer, encoder_padding_mask
+
+
+class Conformer(Transformer):
+ """
+ Args:
+ num_features (int): Number of input features
+ num_classes (int): Number of output classes
+ subsampling_factor (int): subsampling factor of encoder (the convolution layers before transformers)
+ d_model (int): attention dimension, also the output dimension
+ nhead (int): number of head
+ dim_feedforward (int): feedforward dimention
+ num_encoder_layers (int): number of encoder layers
+ num_decoder_layers (int): number of decoder layers
+ dropout (float): dropout rate
+ layer_dropout (float): layer-dropout rate.
+ cnn_module_kernel (int): Kernel size of convolution module
+ vgg_frontend (bool): whether to use vgg frontend.
+ """
+
+ def __init__(
+ self,
+ num_features: int,
+ num_classes: int,
+ subsampling_factor: int = 4,
+ d_model: int = 256,
+ nhead: int = 4,
+ dim_feedforward: int = 2048,
+ num_encoder_layers: int = 12,
+ num_decoder_layers: int = 6,
+ dropout: float = 0.1,
+ layer_dropout: float = 0.075,
+ cnn_module_kernel: int = 31,
+ ) -> None:
+ super(Conformer, self).__init__(
+ num_features=num_features,
+ num_classes=num_classes,
+ subsampling_factor=subsampling_factor,
+ d_model=d_model,
+ nhead=nhead,
+ dim_feedforward=dim_feedforward,
+ num_encoder_layers=num_encoder_layers,
+ num_decoder_layers=num_decoder_layers,
+ dropout=dropout,
+ 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,
+ nhead,
+ dim_feedforward,
+ dropout,
+ layer_dropout,
+ cnn_module_kernel,
+ )
+ self.encoder = ConformerEncoder(encoder_layer, num_encoder_layers)
+
+ 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:
+ Tensor: Predictor tensor of dimension (input_length, batch_size, d_model).
+ Tensor: Mask tensor of dimension (batch_size, input_length)
+ """
+ x = self.encoder_embed(x)
+ x, pos_emb = self.encoder_pos(x)
+ x = x.permute(1, 0, 2) # (N, T, C) -> (T, N, C)
+ mask = encoder_padding_mask(x.size(0), supervisions)
+ if mask is not None:
+ mask = mask.to(x.device)
+
+ # Caution: We assume the subsampling factor is 4!
+
+ x = self.encoder(
+ x, pos_emb, src_key_padding_mask=mask, warmup=warmup
+ ) # (T, N, C)
+
+ # x = x.permute(1, 0, 2) # (T, N, C) ->(N, T, C)
+
+ # return x, lengths
+ return x, mask
+
+
+class ConformerEncoderLayer(nn.Module):
+ """
+ ConformerEncoderLayer is made up of self-attn, feedforward and convolution networks.
+ See: "Conformer: Convolution-augmented Transformer for Speech Recognition"
+
+ Args:
+ d_model: the number of expected features in the input (required).
+ nhead: the number of heads in the multiheadattention models (required).
+ dim_feedforward: the dimension of the feedforward network model (default=2048).
+ dropout: the dropout value (default=0.1).
+ cnn_module_kernel (int): Kernel size of convolution module.
+
+ 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,
+ layer_dropout: float = 0.075,
+ cnn_module_kernel: int = 31,
+ ) -> None:
+ super(ConformerEncoderLayer, self).__init__()
+
+ self.layer_dropout = layer_dropout
+
+ 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),
+ 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: Tensor,
+ pos_emb: Tensor,
+ src_mask: Optional[Tensor] = None,
+ src_key_padding_mask: Optional[Tensor] = None,
+ warmup: float = 1.0,
+ ) -> 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).
+ warmup: controls selective bypass of of layers; if < 1.0, we will
+ bypass layers more frequently.
+
+ 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
+
+ warmup_scale = min(0.1 + 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 0.1
+ )
+ 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, src_key_padding_mask=src_key_padding_mask)
+ )
+
+ # 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):
+ r"""ConformerEncoder is a stack of N encoder layers
+
+ Args:
+ encoder_layer: an instance of the ConformerEncoderLayer() class (required).
+ num_layers: the number of sub-encoder-layers in the encoder (required).
+
+ 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) -> None:
+ 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: Tensor,
+ pos_emb: Tensor,
+ mask: Optional[Tensor] = None,
+ src_key_padding_mask: Optional[Tensor] = None,
+ warmup: float = 1.0,
+ ) -> Tensor:
+ r"""Pass the input through the encoder layers in turn.
+
+ Args:
+ src: the sequence to the encoder (required).
+ pos_emb: Positional embedding tensor (required).
+ mask: the mask for the src sequence (optional).
+ src_key_padding_mask: the mask for the src keys per batch (optional).
+
+ Shape:
+ src: (S, N, E).
+ pos_emb: (N, 2*S-1, E)
+ mask: (S, S).
+ src_key_padding_mask: (N, S).
+ S is the source sequence length, T is the target sequence length, N is the batch size, E is the feature number
+
+ """
+ output = src
+
+ for i, mod in enumerate(self.layers):
+ output = mod(
+ output,
+ pos_emb,
+ src_mask=mask,
+ src_key_padding_mask=src_key_padding_mask,
+ warmup=warmup,
+ )
+
+ return output
+
+
+class RelPositionalEncoding(torch.nn.Module):
+ """Relative positional encoding module.
+
+ See : Appendix B in "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context"
+ Modified from https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/transformer/embedding.py
+
+ Args:
+ d_model: Embedding dimension.
+ dropout_rate: Dropout rate.
+ max_len: Maximum input length.
+
+ """
+
+ def __init__(
+ self, d_model: int, dropout_rate: float, max_len: int = 5000
+ ) -> None:
+ """Construct an PositionalEncoding object."""
+ super(RelPositionalEncoding, self).__init__()
+ self.d_model = d_model
+ self.dropout = torch.nn.Dropout(p=dropout_rate)
+ self.pe = None
+ self.extend_pe(torch.tensor(0.0).expand(1, max_len))
+
+ def extend_pe(self, x: Tensor) -> None:
+ """Reset the positional encodings."""
+ if self.pe is not None:
+ # self.pe contains both positive and negative parts
+ # the length of self.pe is 2 * input_len - 1
+ if self.pe.size(1) >= x.size(1) * 2 - 1:
+ # Note: TorchScript doesn't implement operator== for torch.Device
+ if self.pe.dtype != x.dtype or str(self.pe.device) != str(
+ x.device
+ ):
+ self.pe = self.pe.to(dtype=x.dtype, device=x.device)
+ return
+ # Suppose `i` means to the position of query vecotr and `j` means the
+ # position of key vector. We use position relative positions when keys
+ # are to the left (i>j) and negative relative positions otherwise (i Tuple[Tensor, Tensor]:
+ """Add positional encoding.
+
+ Args:
+ x (torch.Tensor): Input tensor (batch, time, `*`).
+
+ Returns:
+ torch.Tensor: Encoded tensor (batch, time, `*`).
+ torch.Tensor: Encoded tensor (batch, 2*time-1, `*`).
+
+ """
+ self.extend_pe(x)
+ 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):
+ 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 = True,
+ attn_mask: Optional[Tensor] = None,
+ ) -> Tuple[Tensor, Optional[Tensor]]:
+ r"""
+ Args:
+ query, key, value: map a query and a set of key-value pairs to an output.
+ pos_emb: Positional embedding tensor
+ key_padding_mask: if provided, specified padding elements in the key will
+ be ignored by the attention. When given a binary mask and a value is True,
+ the corresponding value on the attention layer will be ignored. When given
+ a byte mask and a value is non-zero, the corresponding value on the attention
+ layer will be ignored
+ need_weights: output attn_output_weights.
+ attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
+ the batches while a 3D mask allows to specify a different mask for the entries of each batch.
+
+ Shape:
+ - Inputs:
+ - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
+ the embedding dimension.
+ - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
+ the embedding dimension.
+ - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
+ the embedding dimension.
+ - pos_emb: :math:`(N, 2*L-1, E)` where L is the target sequence length, N is the batch size, E is
+ the embedding dimension.
+ - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
+ If a ByteTensor is provided, the non-zero positions will be ignored while the position
+ with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the
+ value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+ - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
+ 3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
+ S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked
+ positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
+ while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
+ is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+ is provided, it will be added to the attention weight.
+
+ - Outputs:
+ - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
+ E is the embedding dimension.
+ - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
+ L is the target sequence length, S is the source sequence length.
+ """
+ return self.multi_head_attention_forward(
+ query,
+ key,
+ value,
+ pos_emb,
+ self.embed_dim,
+ self.num_heads,
+ self.in_proj.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: Tensor) -> Tensor:
+ """Compute relative positional encoding.
+
+ Args:
+ x: Input tensor (batch, head, time1, 2*time1-1).
+ time1 means the length of query vector.
+
+ Returns:
+ Tensor: tensor of shape (batch, head, time1, time2)
+ (note: time2 has the same value as time1, but it is for
+ the key, while time1 is for the query).
+ """
+ (batch_size, num_heads, time1, n) = x.shape
+ assert n == 2 * time1 - 1
+ # Note: TorchScript requires explicit arg for stride()
+ batch_stride = x.stride(0)
+ head_stride = x.stride(1)
+ time1_stride = x.stride(2)
+ n_stride = x.stride(3)
+ return x.as_strided(
+ (batch_size, num_heads, time1, time1),
+ (batch_stride, head_stride, time1_stride - n_stride, n_stride),
+ storage_offset=n_stride * (time1 - 1),
+ )
+
+ def multi_head_attention_forward(
+ self,
+ query: Tensor,
+ key: Tensor,
+ value: Tensor,
+ pos_emb: Tensor,
+ embed_dim_to_check: int,
+ num_heads: int,
+ in_proj_weight: Tensor,
+ in_proj_bias: Tensor,
+ dropout_p: float,
+ out_proj_weight: Tensor,
+ out_proj_bias: Tensor,
+ training: bool = True,
+ key_padding_mask: Optional[Tensor] = None,
+ need_weights: bool = True,
+ attn_mask: Optional[Tensor] = None,
+ ) -> Tuple[Tensor, Optional[Tensor]]:
+ r"""
+ Args:
+ query, key, value: map a query and a set of key-value pairs to an output.
+ pos_emb: Positional embedding tensor
+ embed_dim_to_check: total dimension of the model.
+ num_heads: parallel attention heads.
+ in_proj_weight, in_proj_bias: input projection weight and bias.
+ dropout_p: probability of an element to be zeroed.
+ out_proj_weight, out_proj_bias: the output projection weight and bias.
+ training: apply dropout if is ``True``.
+ key_padding_mask: if provided, specified padding elements in the key will
+ be ignored by the attention. This is an binary mask. When the value is True,
+ the corresponding value on the attention layer will be filled with -inf.
+ need_weights: output attn_output_weights.
+ attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
+ the batches while a 3D mask allows to specify a different mask for the entries of each batch.
+
+ Shape:
+ Inputs:
+ - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
+ the embedding dimension.
+ - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
+ the embedding dimension.
+ - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
+ the embedding dimension.
+ - pos_emb: :math:`(N, 2*L-1, E)` or :math:`(1, 2*L-1, E)` where L is the target sequence
+ length, N is the batch size, E is the embedding dimension.
+ - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
+ If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions
+ will be unchanged. If a BoolTensor is provided, the positions with the
+ value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+ - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
+ 3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
+ S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked
+ positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
+ while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
+ are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+ is provided, it will be added to the attention weight.
+
+ Outputs:
+ - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
+ E is the embedding dimension.
+ - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
+ L is the target sequence length, S is the source sequence length.
+ """
+
+ tgt_len, bsz, embed_dim = query.size()
+ assert embed_dim == embed_dim_to_check
+ assert key.size(0) == value.size(0) and key.size(1) == value.size(1)
+
+ head_dim = embed_dim // num_heads
+ assert (
+ head_dim * num_heads == embed_dim
+ ), "embed_dim must be divisible by num_heads"
+
+ scaling = float(head_dim) ** -0.5
+
+ if torch.equal(query, key) and torch.equal(key, value):
+ # self-attention
+ q, k, v = nn.functional.linear(
+ query, in_proj_weight, in_proj_bias
+ ).chunk(3, dim=-1)
+
+ elif torch.equal(key, value):
+ # encoder-decoder attention
+ # This is inline in_proj function with in_proj_weight and in_proj_bias
+ _b = in_proj_bias
+ _start = 0
+ _end = embed_dim
+ _w = in_proj_weight[_start:_end, :]
+ if _b is not None:
+ _b = _b[_start:_end]
+ q = nn.functional.linear(query, _w, _b)
+
+ # This is inline in_proj function with in_proj_weight and in_proj_bias
+ _b = in_proj_bias
+ _start = embed_dim
+ _end = None
+ _w = in_proj_weight[_start:, :]
+ if _b is not None:
+ _b = _b[_start:]
+ k, v = nn.functional.linear(key, _w, _b).chunk(2, dim=-1)
+
+ else:
+ # This is inline in_proj function with in_proj_weight and in_proj_bias
+ _b = in_proj_bias
+ _start = 0
+ _end = embed_dim
+ _w = in_proj_weight[_start:_end, :]
+ if _b is not None:
+ _b = _b[_start:_end]
+ q = nn.functional.linear(query, _w, _b)
+
+ # This is inline in_proj function with in_proj_weight and in_proj_bias
+ _b = in_proj_bias
+ _start = embed_dim
+ _end = embed_dim * 2
+ _w = in_proj_weight[_start:_end, :]
+ if _b is not None:
+ _b = _b[_start:_end]
+ k = nn.functional.linear(key, _w, _b)
+
+ # This is inline in_proj function with in_proj_weight and in_proj_bias
+ _b = in_proj_bias
+ _start = embed_dim * 2
+ _end = None
+ _w = in_proj_weight[_start:, :]
+ if _b is not None:
+ _b = _b[_start:]
+ v = nn.functional.linear(value, _w, _b)
+
+ if 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:
+ 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):
+ """ConvolutionModule in Conformer model.
+ Modified from https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/conformer/convolution.py
+
+ Args:
+ channels (int): The number of channels of conv layers.
+ kernel_size (int): Kernerl size of conv layers.
+ bias (bool): Whether to use bias in conv layers (default=True).
+
+ """
+
+ def __init__(
+ self, channels: int, kernel_size: int, bias: bool = True
+ ) -> None:
+ """Construct an ConvolutionModule object."""
+ super(ConvolutionModule, self).__init__()
+ # kernerl_size should be a odd number for 'SAME' padding
+ assert (kernel_size - 1) % 2 == 0
+
+ self.pointwise_conv1 = 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: 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).
+
+ Returns:
+ Tensor: Output tensor (#time, batch, channels).
+
+ """
+ # exchange the temporal dimension and the feature dimension
+ x = x.permute(1, 2, 0) # (#batch, channels, time).
+
+ # GLU mechanism
+ x = self.pointwise_conv1(x) # (batch, 2*channels, time)
+
+ x = self.deriv_balancer1(x)
+ x = nn.functional.glu(x, dim=1) # (batch, channels, time)
+
+ # 1D Depthwise Conv
+ if src_key_padding_mask is not None:
+ 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.activation(x)
+
+ x = self.pointwise_conv2(x) # (batch, channel, time)
+
+ return x.permute(2, 0, 1)
+
+
+if __name__ == "__main__":
+ feature_dim = 50
+ c = Conformer(num_features=feature_dim, d_model=128, nhead=4)
+ batch_size = 5
+ seq_len = 20
+ # 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),
+ warmup=0.5,
+ )
diff --git a/egs/librispeech/ASR/conformer_ctc2/decode.py b/egs/librispeech/ASR/conformer_ctc2/decode.py
new file mode 100755
index 000000000..97f2f2d39
--- /dev/null
+++ b/egs/librispeech/ASR/conformer_ctc2/decode.py
@@ -0,0 +1,999 @@
+#!/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 LibriSpeechAsrDataModule
+from conformer import Conformer
+
+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_rnn_lm,
+ rescore_with_whole_lattice,
+)
+from icefall.env import get_env_info
+from icefall.lexicon import Lexicon
+from icefall.rnn_lm.model import RnnLmModel
+from icefall.utils import (
+ AttributeDict,
+ get_texts,
+ load_averaged_model,
+ 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=77,
+ 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) rnn-lm. Rescoring with attention-decoder and RNN LM. We assume
+ you have trained an RNN LM using ./rnn_lm/train.py
+ - (8) 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-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(
+ "--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, rnn-lm, 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, rnn-lm, 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-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
+ """,
+ )
+
+ 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",
+ type=str2bool,
+ default=False,
+ help="""True to share the weights between the input embedding layer and the
+ last output linear layer
+ """,
+ )
+
+ return parser
+
+
+def get_params() -> AttributeDict:
+ params = AttributeDict(
+ {
+ # parameters for conformer
+ "subsampling_factor": 4,
+ "feature_dim": 80,
+ "nhead": 8,
+ "dim_feedforward": 2048,
+ "encoder_dim": 512,
+ "num_encoder_layers": 12,
+ # parameters for decoding
+ "search_beam": 20,
+ "output_beam": 8,
+ "min_active_states": 30,
+ "max_active_states": 10000,
+ "use_double_scores": True,
+ "env_info": get_env_info(),
+ }
+ )
+ return params
+
+
+def ctc_greedy_search(
+ nnet_output: torch.Tensor,
+ memory: torch.Tensor,
+ memory_key_padding_mask: torch.Tensor,
+) -> List[List[int]]:
+ """Apply CTC greedy search
+
+ Args:
+ speech (torch.Tensor): (batch, max_len, feat_dim)
+ speech_length (torch.Tensor): (batch, )
+ Returns:
+ List[List[int]]: best path result
+ """
+ batch_size = memory.shape[1]
+ # Let's assume B = batch_size
+ encoder_out = memory
+ encoder_mask = memory_key_padding_mask
+ maxlen = encoder_out.size(0)
+
+ ctc_probs = nnet_output # (B, maxlen, vocab_size)
+ topk_prob, topk_index = ctc_probs.topk(1, dim=2) # (B, maxlen, 1)
+ topk_index = topk_index.view(batch_size, maxlen) # (B, maxlen)
+ topk_index = topk_index.masked_fill_(encoder_mask, 0) # (B, maxlen)
+ hyps = [hyp.tolist() for hyp in topk_index]
+ scores = topk_prob.max(1)
+ hyps = [remove_duplicates_and_blank(hyp) for hyp in hyps]
+ return hyps, scores
+
+
+def remove_duplicates_and_blank(hyp: List[int]) -> List[int]:
+ # from https://github.com/wenet-e2e/wenet/blob/main/wenet/utils/common.py
+ new_hyp: List[int] = []
+ cur = 0
+ while cur < len(hyp):
+ if hyp[cur] != 0:
+ new_hyp.append(hyp[cur])
+ prev = cur
+ while cur < len(hyp) and hyp[cur] == hyp[prev]:
+ cur += 1
+ return new_hyp
+
+
+def decode_one_batch(
+ params: AttributeDict,
+ model: nn.Module,
+ rnn_lm_model: Optional[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.
+ rnn_lm_model:
+ The neural model for RNN LM.
+ 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="trunc",
+ ),
+ torch.div(
+ supervisions["num_frames"],
+ params.subsampling_factor,
+ rounding_mode="trunc",
+ ),
+ ),
+ 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'], ... ]
+ hyps = [s.split() for s in hyps]
+ key = "ctc-decoding"
+ return {key: hyps}
+
+ if params.method == "ctc-greedy-search":
+ hyps, _ = ctc_greedy_search(
+ nnet_output,
+ memory,
+ 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'], ... ]
+ hyps = [s.split() 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] for ids in hyps]
+ key = f"oracle_{params.num_paths}_nbest_scale_{params.nbest_scale}" # noqa
+ return {key: hyps}
+
+ if params.method in ["1best", "nbest"]:
+ if params.method == "1best":
+ best_path = one_best_decoding(
+ lattice=lattice, use_double_scores=params.use_double_scores
+ )
+ key = "no_rescore"
+ else:
+ best_path = nbest_decoding(
+ lattice=lattice,
+ num_paths=params.num_paths,
+ use_double_scores=params.use_double_scores,
+ 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.method in [
+ "nbest-rescoring",
+ "whole-lattice-rescoring",
+ "attention-decoder",
+ "rnn-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 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":
+ # lattice uses a 3-gram Lm. We rescore it with a 4-gram LM.
+ rescored_lattice = rescore_with_whole_lattice(
+ lattice=lattice,
+ G_with_epsilon_loops=G,
+ lm_scale_list=None,
+ )
+ # TODO: pass `lattice` instead of `rescored_lattice` to
+ # `rescore_with_attention_decoder`
+
+ best_path_dict = rescore_with_attention_decoder(
+ lattice=rescored_lattice,
+ num_paths=params.num_paths,
+ model=model,
+ memory=memory,
+ memory_key_padding_mask=memory_key_padding_mask,
+ sos_id=sos_id,
+ eos_id=eos_id,
+ nbest_scale=params.nbest_scale,
+ )
+ elif params.method == "rnn-lm":
+ # lattice uses a 3-gram Lm. We rescore it with a 4-gram LM.
+ rescored_lattice = rescore_with_whole_lattice(
+ lattice=lattice,
+ G_with_epsilon_loops=G,
+ lm_scale_list=None,
+ )
+
+ best_path_dict = rescore_with_rnn_lm(
+ lattice=rescored_lattice,
+ num_paths=params.num_paths,
+ rnn_lm_model=rnn_lm_model,
+ model=model,
+ memory=memory,
+ memory_key_padding_mask=memory_key_padding_mask,
+ sos_id=sos_id,
+ eos_id=eos_id,
+ blank_id=0,
+ nbest_scale=params.nbest_scale,
+ )
+ else:
+ assert False, 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] 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,
+ rnn_lm_model: Optional[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.
+ rnn_lm_model:
+ The neural model for RNN LM.
+ 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,
+ rnn_lm_model=rnn_lm_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]]]],
+):
+ if params.method in ("attention-decoder", "rnn-lm"):
+ # 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.exp_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.exp_dir / f"errs-{test_set_name}-{key}.txt"
+ with open(errs_filename, "w") as f:
+ wer = write_error_stats(
+ f, f"{test_set_name}-{key}", results, 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.exp_dir / f"wer-summary-{test_set_name}.txt"
+ with open(errs_info, "w") as f:
+ print("settings\tWER", file=f)
+ for key, val in test_set_wers:
+ print("{}\t{}".format(key, val), file=f)
+
+ s = "\nFor {}, WER of different settings are:\n".format(test_set_name)
+ note = "\tbest for {}".format(test_set_name)
+ for key, val in test_set_wers:
+ s += "{}\t{}{}\n".format(key, val, note)
+ note = ""
+ logging.info(s)
+
+
+@torch.no_grad()
+def main():
+ parser = get_parser()
+ LibriSpeechAsrDataModule.add_arguments(parser)
+ args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
+ args.lang_dir = Path(args.lang_dir)
+ args.lm_dir = Path(args.lm_dir)
+
+ 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
+
+ params.num_classes = num_classes
+ params.sos_id = sos_id
+ params.eos_id = eos_id
+
+ if params.method == "ctc-decoding" or params.method == "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",
+ "attention-decoder",
+ "rnn-lm",
+ ):
+ 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.method in [
+ "whole-lattice-rescoring",
+ "attention-decoder",
+ "rnn-lm",
+ ]:
+ # Add epsilon self-loops to G as we will compose
+ # it with the whole lattice later
+ G = k2.add_epsilon_self_loops(G)
+ G = k2.arc_sort(G)
+ G = G.to(device)
+
+ # G.lm_scores is used to replace HLG.lm_scores during
+ # LM rescoring.
+ G.lm_scores = G.scores.clone()
+ else:
+ G = None
+
+ model = Conformer(
+ num_features=params.feature_dim,
+ nhead=params.nhead,
+ d_model=params.encoder_dim,
+ num_classes=num_classes,
+ subsampling_factor=params.subsampling_factor,
+ 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}")
+
+ rnn_lm_model = None
+ if params.method == "rnn-lm":
+ rnn_lm_model = RnnLmModel(
+ vocab_size=params.num_classes,
+ 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.rnn_lm_avg == 1:
+ load_checkpoint(
+ f"{params.rnn_lm_exp_dir}/epoch-{params.rnn_lm_epoch}.pt",
+ rnn_lm_model,
+ )
+ rnn_lm_model.to(device)
+ else:
+ rnn_lm_model = load_averaged_model(
+ params.rnn_lm_exp_dir,
+ rnn_lm_model,
+ params.rnn_lm_epoch,
+ params.rnn_lm_avg,
+ device,
+ )
+ rnn_lm_model.eval()
+
+ # 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,
+ rnn_lm_model=rnn_lm_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)
+torch.set_num_interop_threads(1)
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/librispeech/ASR/conformer_ctc2/export.py b/egs/librispeech/ASR/conformer_ctc2/export.py
new file mode 100755
index 000000000..584b3c3fc
--- /dev/null
+++ b/egs/librispeech/ASR/conformer_ctc2/export.py
@@ -0,0 +1,281 @@
+#!/usr/bin/env python3
+#
+# Copyright 2021 Xiaomi Corporation (Author: 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.
+
+# 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/librispeech/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 decode import get_params
+
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from conformer import Conformer
+
+from icefall.utils import str2bool
+from icefall.lexicon import Lexicon
+
+
+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(
+ "--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-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(
+ "--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.
+ """,
+ )
+
+ return parser
+
+
+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,
+ nhead=params.nhead,
+ d_model=params.encoder_dim,
+ num_classes=num_classes,
+ subsampling_factor=params.subsampling_factor,
+ 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"
+ 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.eval()
+
+ 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/egs/librispeech/ASR/conformer_ctc2/label_smoothing.py b/egs/librispeech/ASR/conformer_ctc2/label_smoothing.py
new file mode 120000
index 000000000..08734abd7
--- /dev/null
+++ b/egs/librispeech/ASR/conformer_ctc2/label_smoothing.py
@@ -0,0 +1 @@
+../conformer_ctc/label_smoothing.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conformer_ctc2/optim.py b/egs/librispeech/ASR/conformer_ctc2/optim.py
new file mode 120000
index 000000000..e2deb4492
--- /dev/null
+++ b/egs/librispeech/ASR/conformer_ctc2/optim.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/optim.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conformer_ctc2/scaling.py b/egs/librispeech/ASR/conformer_ctc2/scaling.py
new file mode 120000
index 000000000..09d802cc4
--- /dev/null
+++ b/egs/librispeech/ASR/conformer_ctc2/scaling.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/scaling.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conformer_ctc2/subsampling.py b/egs/librispeech/ASR/conformer_ctc2/subsampling.py
new file mode 100644
index 000000000..3fcb4196f
--- /dev/null
+++ b/egs/librispeech/ASR/conformer_ctc2/subsampling.py
@@ -0,0 +1,121 @@
+#!/usr/bin/env python3
+# Copyright (c) 2021 University of Chinese Academy of Sciences (author: Han Zhu)
+# 2022 Xiaomi Corporation (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 torch
+from scaling import (
+ ActivationBalancer,
+ BasicNorm,
+ DoubleSwish,
+ ScaledConv2d,
+ ScaledLinear,
+)
+from torch import nn
+
+
+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-1)//2 - 1)//2, which approximates T' == T//4
+
+ 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,
+ ) -> 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-1)//2 - 1)//2, out_channels)
+ layer1_channels:
+ Number of channels in layer1
+ layer1_channels:
+ Number of channels in layer2
+ """
+ assert in_channels >= 7
+ super().__init__()
+
+ self.conv = nn.Sequential(
+ ScaledConv2d(
+ in_channels=1,
+ out_channels=layer1_channels,
+ kernel_size=3,
+ padding=1,
+ ),
+ ActivationBalancer(channel_dim=1),
+ DoubleSwish(),
+ ScaledConv2d(
+ in_channels=layer1_channels,
+ out_channels=layer2_channels,
+ kernel_size=3,
+ stride=2,
+ ),
+ ActivationBalancer(channel_dim=1),
+ DoubleSwish(),
+ ScaledConv2d(
+ in_channels=layer2_channels,
+ out_channels=layer3_channels,
+ kernel_size=3,
+ stride=2,
+ ),
+ ActivationBalancer(channel_dim=1),
+ DoubleSwish(),
+ )
+ self.out = ScaledLinear(
+ layer3_channels * (((in_channels - 1) // 2 - 1) // 2), out_channels
+ )
+ # set learn_eps=False because out_norm is preceded by `out`, and `out`
+ # itself has learned scale, so the extra degree of freedom is not
+ # needed.
+ self.out_norm = BasicNorm(out_channels, learn_eps=False)
+ # constrain median of output to be close to zero.
+ self.out_balancer = ActivationBalancer(
+ channel_dim=-1, min_positive=0.45, max_positive=0.55
+ )
+
+ 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-1)//2 - 1)//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-1)//2 - 1)//2, ((idim-1)//2 - 1)//2)
+ b, c, t, f = x.size()
+ x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f))
+ # Now x is of shape (N, ((T-1)//2 - 1))//2, odim)
+ x = self.out_norm(x)
+ x = self.out_balancer(x)
+ return x
diff --git a/egs/librispeech/ASR/conformer_ctc2/train.py b/egs/librispeech/ASR/conformer_ctc2/train.py
new file mode 100755
index 000000000..9d9c2af1f
--- /dev/null
+++ b/egs/librispeech/ASR/conformer_ctc2/train.py
@@ -0,0 +1,1128 @@
+#!/usr/bin/env python3
+# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang,
+# Wei Kang,
+# Mingshuang Luo,
+# Zengwei Yao,
+# 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.
+"""
+Usage:
+
+export CUDA_VISIBLE_DEVICES="0,1,2,3"
+
+./conformer_ctc2/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --exp-dir conformer_ctc2/exp \
+ --full-libri 1 \
+ --max-duration 300
+
+# For mix precision training:
+
+./conformer_ctc2/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --use-fp16 1 \
+ --exp-dir conformer_ctc2/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 torch
+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 optim import Eden, Eve
+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.bpe_graph_compiler import BpeCtcTrainingGraphCompiler
+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.graph_compiler import CtcTrainingGraphCompiler
+from icefall.lexicon import Lexicon
+from icefall.utils import (
+ AttributeDict,
+ MetricsTracker,
+ encode_supervisions,
+ setup_logger,
+ str2bool,
+)
+
+LRSchedulerType = Union[
+ torch.optim.lr_scheduler._LRScheduler, optim.LRScheduler
+]
+
+
+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="conformer_ctc2/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(
+ "--bpe-model",
+ type=str,
+ default="data/lang_bpe_500/bpe.model",
+ help="Path to the 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(
+ "--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(
+ "--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(
+ "--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=8000,
+ 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=20,
+ 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.",
+ )
+
+ 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.
+
+ - beam_size: It is used in k2.ctc_loss
+
+ - reduction: It is used in k2.ctc_loss
+
+ - use_double_scores: It is used in k2.ctc_loss
+
+ - 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": 1,
+ "reset_interval": 200,
+ "valid_interval": 3000, # For the 100h subset, use 800
+ # parameters for conformer
+ "feature_dim": 80,
+ "subsampling_factor": 4,
+ "encoder_dim": 512,
+ "nhead": 8,
+ "dim_feedforward": 2048,
+ "num_encoder_layers": 12,
+ # parameters for ctc loss
+ "beam_size": 10,
+ "reduction": "sum",
+ "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: 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],
+ batch: dict,
+ graph_compiler: BpeCtcTrainingGraphCompiler,
+ 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
+ )
+ # logging.info('feature shape: {}'.format(feature.shape))
+ # logging.info('nnet_output shape: {}'.format(nnet_output.shape))
+ # logging.info('encoder_memory shape: {}'.format(encoder_memory.shape))
+ # logging.info('memory_mask shape: {}'.format(memory_mask.shape))
+ # after the main warmup step, we keep pruned_loss_scale small
+ # for the same amount of time (model_warm_step), to avoid
+ # overwhelming the simple_loss and causing it to diverge,
+ # in case it had not fully learned the alignment yet.
+
+ # NOTE: We need `encode_supervisions` to sort sequences with
+ # different duration in decreasing order, required by
+ # `k2.intersect_dense` called in `k2.ctc_loss`
+ supervision_segments, texts = encode_supervisions(
+ supervisions, subsampling_factor=params.subsampling_factor
+ )
+
+ if isinstance(graph_compiler, BpeCtcTrainingGraphCompiler):
+ # Works with a BPE model
+ token_ids = graph_compiler.texts_to_ids(texts)
+ decoding_graph = graph_compiler.compile(token_ids)
+ elif isinstance(graph_compiler, CtcTrainingGraphCompiler):
+ # Works with a phone lexicon
+ decoding_graph = graph_compiler.compile(texts)
+ else:
+ raise ValueError(
+ f"Unsupported type of graph compiler: {type(graph_compiler)}"
+ )
+
+ 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,
+ )
+ loss = (1.0 - params.att_rate) * ctc_loss + params.att_rate * att_loss
+ else:
+ loss = ctc_loss
+ att_loss = torch.tensor([0])
+
+ assert loss.requires_grad == is_training
+
+ info = MetricsTracker()
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ info["frames"] = (
+ (feature_lens // params.subsampling_factor).sum().item()
+ )
+ info["ctc_loss"] = ctc_loss.detach().cpu().item()
+ if params.att_rate != 0.0:
+ info["att_loss"] = att_loss.detach().cpu().item()
+
+ # Note: We use reduction=sum while computing the loss.
+ info["loss"] = loss.detach().cpu().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()
+ )
+
+ return loss, info
+
+
+def compute_validation_loss(
+ params: AttributeDict,
+ model: Union[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_idx, batch in enumerate(valid_dl):
+ loss, loss_info = compute_loss(
+ params=params,
+ model=model,
+ batch=batch,
+ graph_compiler=graph_compiler,
+ 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,
+ graph_compiler: BpeCtcTrainingGraphCompiler,
+ scheduler: LRSchedulerType,
+ 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.
+ graph_compiler:
+ It is used to convert transcripts to FSAs.
+ 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"])
+ batch_name = batch["supervisions"]["uttid"]
+
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, loss_info = compute_loss(
+ params=params,
+ model=model,
+ batch=batch,
+ graph_compiler=graph_compiler,
+ 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()
+
+ try:
+ # loss.backward()
+ scaler.scale(loss).backward()
+ except RuntimeError as e:
+ if "CUDA out of memory" in str(e):
+ logging.error(
+ f"failing batch size:{batch_size} "
+ f"failing batch names {batch_name}"
+ )
+ raise
+
+ scheduler.step_batch(params.batch_idx_train)
+ scaler.step(optimizer)
+ scaler.update()
+ optimizer.zero_grad()
+
+ if params.print_diagnostics and batch_idx == 30:
+ 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 % 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 loss_info["ctc_loss"] == float("inf") or loss_info[
+ "att_loss"
+ ] == float("inf"):
+ logging.error(
+ "Your loss contains inf, something goes wrong"
+ f"failing batch names {batch_name}"
+ )
+ 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))
+ 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")
+ 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)
+
+ if "lang_bpe" in str(params.lang_dir):
+ graph_compiler = BpeCtcTrainingGraphCompiler(
+ params.lang_dir,
+ device=device,
+ sos_token="",
+ eos_token="",
+ )
+ elif "lang_phone" in str(params.lang_dir):
+ assert params.att_rate == 0, (
+ "Attention decoder training does not support phone lang dirs "
+ "at this time due to a missing symbol. Set --att-rate=0 "
+ "for pure CTC training when using a phone-based lang dir."
+ )
+ assert params.num_decoder_layers == 0, (
+ "Attention decoder training does not support phone lang dirs "
+ "at this time due to a missing symbol. "
+ "Set --num-decoder-layers=0 for pure CTC training when using "
+ "a phone-based lang dir."
+ )
+ graph_compiler = CtcTrainingGraphCompiler(
+ lexicon,
+ device=device,
+ )
+ # Manually add the sos/eos ID with their default values
+ # from the BPE recipe which we're adapting here.
+ graph_compiler.sos_id = 1
+ graph_compiler.eos_id = 1
+ else:
+ raise ValueError(
+ f"Unsupported type of lang dir (we expected it to have "
+ f"'lang_bpe' or 'lang_phone' in its name): {params.lang_dir}"
+ )
+
+ logging.info("About to create model")
+ model = Conformer(
+ num_features=params.feature_dim,
+ nhead=params.nhead,
+ d_model=params.encoder_dim,
+ num_classes=num_classes,
+ subsampling_factor=params.subsampling_factor,
+ num_encoder_layers=params.num_encoder_layers,
+ num_decoder_layers=params.num_decoder_layers,
+ )
+
+ print(model)
+
+ 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)
+
+ 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 = Eve(model.parameters(), lr=params.initial_lr)
+
+ 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:
+ diagnostic = diagnostics.attach_diagnostics(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
+ return 1.0 <= c.duration <= 20.0
+
+ def remove_invalid_utt_ctc(c: Cut):
+ # Caution: We assume the subsampling factor is 4!
+ # num_tokens = len(sp.encode(c.supervisions[0].text, out_type=int))
+ num_tokens = len(graph_compiler.texts_to_ids(c.supervisions[0].text))
+ min_output_input_ratio = 0.0005
+ max_output_input_ratio = 0.1
+ return (
+ min_output_input_ratio
+ < num_tokens / float(c.features.num_frames)
+ < max_output_input_ratio
+ )
+
+ train_cuts = train_cuts.filter(remove_short_and_long_utt)
+ train_cuts = train_cuts.filter(remove_invalid_utt_ctc)
+
+ 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 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)
+ 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,
+ graph_compiler=graph_compiler,
+ scheduler=scheduler,
+ 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[nn.Module, DDP],
+ train_dl: torch.utils.data.DataLoader,
+ optimizer: torch.optim.Optimizer,
+ graph_compiler: BpeCtcTrainingGraphCompiler,
+ 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:
+ # warmup = 0.0 is so that the derivs for the pruned loss stay zero
+ # (i.e. are not remembered by the decaying-average in adam), because
+ # we want to avoid these params being subject to shrinkage in adam.
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, _ = compute_loss(
+ params=params,
+ model=model,
+ batch=batch,
+ graph_compiler=graph_compiler,
+ is_training=True,
+ warmup=0.0,
+ )
+ loss.backward()
+ optimizer.step()
+ optimizer.zero_grad()
+ except RuntimeError 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]}) ..."
+ )
+ raise
+
+
+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/conformer_ctc2/transformer.py b/egs/librispeech/ASR/conformer_ctc2/transformer.py
new file mode 100644
index 000000000..fa179acc0
--- /dev/null
+++ b/egs/librispeech/ASR/conformer_ctc2/transformer.py
@@ -0,0 +1,1092 @@
+# 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 label_smoothing import LabelSmoothingLoss
+from subsampling import Conv2dSubsampling
+from attention import MultiheadAttention
+from torch.nn.utils.rnn import pad_sequence
+
+from scaling import (
+ ActivationBalancer,
+ BasicNorm,
+ DoubleSwish,
+ ScaledLinear,
+ ScaledEmbedding,
+)
+
+
+# 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,
+ ) -> 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 (float): layer-dropout rate.
+ """
+ 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,
+ )
+
+ self.encoder = TransformerEncoder(
+ encoder_layer=encoder_layer,
+ num_layers=num_encoder_layers,
+ )
+
+ # 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,
+ )
+
+ self.decoder_output_layer = ScaledLinear(
+ d_model, self.decoder_num_class, bias=True
+ )
+
+ self.decoder_criterion = LabelSmoothingLoss()
+ 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, T, 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, T, C)
+ - Encoder output with shape (T, 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, T).
+ 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, T, 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.
+ Returns:
+ Return a tuple with two tensors:
+ - The encoder output, with shape (T, N, C)
+ - encoder padding mask, with shape (N, T).
+ 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, T, C) -> (T, 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
+ ) # (T, 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 (T, N, C)
+
+ Returns:
+ Return a tensor that can be used for CTC decoding.
+ Its shape is (N, T, C)
+ """
+ x = self.encoder_output_layer(x)
+ x = x.permute(1, 0, 2) # (T, N, C) ->(N, T, C)
+ x = nn.functional.log_softmax(x, dim=-1) # (N, T, 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,
+ ) -> torch.Tensor:
+ """
+ Args:
+ memory:
+ It's the output of the encoder with shape (T, N, C)
+ memory_key_padding_mask:
+ The padding mask from the encoder.
+ 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
+
+ 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,
+ ) # (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,
+ ) -> torch.Tensor:
+ """
+ Args:
+ memory:
+ It's the output of the encoder with shape (T, N, C)
+ memory_key_padding_mask:
+ The padding mask from the encoder.
+ 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.
+ 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) # (B, T) -> (B, T, F)
+ tgt = self.decoder_pos(tgt)
+ tgt = tgt.permute(1, 0, 2) # (B, T, F) -> (T, B, F)
+ pred_pad = self.decoder(
+ tgt=tgt,
+ memory=memory,
+ tgt_mask=tgt_mask,
+ tgt_key_padding_mask=tgt_key_padding_mask,
+ memory_key_padding_mask=memory_key_padding_mask,
+ ) # (T, B, F)
+ pred_pad = pred_pad.permute(1, 0, 2) # (T, B, F) -> (B, T, F)
+ pred_pad = self.decoder_output_layer(pred_pad) # (B, T, F)
+ # nll: negative log-likelihood
+ nll = torch.nn.functional.cross_entropy(
+ pred_pad.view(-1, self.decoder_num_class),
+ ys_out_pad.view(-1),
+ ignore_index=-1,
+ reduction="none",
+ )
+
+ nll = nll.view(pred_pad.shape[0], -1)
+
+ return nll
+
+
+class TransformerEncoderLayer(nn.Module):
+ """
+ Modified from torch.nn.TransformerEncoderLayer.
+
+ Args:
+ d_model:
+ the number of expected features in the input (required).
+ nhead:
+ the number of heads in the multiheadattention models (required).
+ dim_feedforward:
+ the dimension of the feedforward network model (default=2048).
+ dropout:
+ the dropout value (default=0.1).
+ activation:
+ the activation function of intermediate layer, relu or
+ gelu (default=relu).
+
+ Examples::
+ >>> 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,
+ layer_dropout: float = 0.075,
+ activation: str = "relu",
+ ) -> None:
+ super(TransformerEncoderLayer, self).__init__()
+
+ self.layer_dropout = layer_dropout
+
+ self.self_attn = MultiheadAttention(d_model, nhead, dropout=0.0)
+ # 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 __setstate__(self, state):
+ # if "activation" not in state:
+ # state["activation"] = nn.functional.relu
+ # super(TransformerEncoderLayer, self).__setstate__(state)
+
+ 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 (required).
+ src_mask: the mask for the src sequence (optional).
+ src_key_padding_mask: the mask for the src keys per batch (optional)
+ warmup: controls selective bypass of of layers; if < 1.0, we will
+ bypass layers more frequently.
+
+ Shape:
+ src: (S, N, E).
+ src_mask: (S, S).
+ src_key_padding_mask: (N, S).
+ S is the source sequence length, T is the target sequence length,
+ N is the batch size, E is the feature number
+ """
+ src_orig = src
+
+ warmup_scale = min(0.1 + 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 0.1
+ )
+ else:
+ alpha = 1.0
+
+ # src_att = self.self_attn(src, src, src, src_mask)
+ 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 - alpha) * src_orig
+
+ return src
+
+
+class TransformerDecoderLayer(nn.Module):
+ """
+ Modified from torch.nn.TransformerDecoderLayer.
+ Add support of normalize_before,
+ i.e., use layer_norm before the first block.
+
+ Args:
+ d_model:
+ the number of expected features in the input (required).
+ nhead:
+ the number of heads in the multiheadattention models (required).
+ dim_feedforward:
+ the dimension of the feedforward network model (default=2048).
+ dropout:
+ the dropout value (default=0.1).
+ activation:
+ the activation function of intermediate layer, relu or
+ gelu (default=relu).
+
+ Examples::
+ >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
+ >>> memory = torch.rand(10, 32, 512)
+ >>> tgt = torch.rand(20, 32, 512)
+ >>> out = decoder_layer(tgt, memory)
+ """
+
+ def __init__(
+ self,
+ d_model: int,
+ nhead: int,
+ dim_feedforward: int = 2048,
+ dropout: float = 0.1,
+ layer_dropout: float = 0.075,
+ # activation: str = "relu",
+ normalize_before: bool = True,
+ ) -> None:
+ super(TransformerDecoderLayer, self).__init__()
+ self.layer_dropout = layer_dropout
+ self.self_attn = MultiheadAttention(d_model, nhead, dropout=0.0)
+ self.src_attn = MultiheadAttention(d_model, nhead, dropout=0.0)
+ # 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 __setstate__(self, state):
+ # if "activation" not in state:
+ # state["activation"] = nn.functional.relu
+ # super(TransformerDecoderLayer, self).__setstate__(state)
+
+ 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 (required).
+ memory:
+ the sequence from the last layer of the encoder (required).
+ tgt_mask:
+ the mask for the tgt sequence (optional).
+ memory_mask:
+ the mask for the memory sequence (optional).
+ tgt_key_padding_mask:
+ the mask for the tgt keys per batch (optional).
+ memory_key_padding_mask:
+ the mask for the memory keys per batch (optional).
+ warmup: controls selective bypass of of layers; if < 1.0, we will
+ bypass layers more frequently.
+
+
+
+ Shape:
+ tgt: (T, N, E).
+ memory: (S, N, E).
+ tgt_mask: (T, T).
+ memory_mask: (T, S).
+ tgt_key_padding_mask: (N, T).
+ memory_key_padding_mask: (N, S).
+ S is the source sequence length, T is the target sequence length,
+ N is the batch size, E is the feature number
+ """
+ tgt_orig = tgt
+
+ warmup_scale = min(0.1 + 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 0.1
+ )
+ else:
+ alpha = 1.0
+
+ # tgt_att = self.self_attn(tgt, tgt, tgt, tgt_mask)
+ 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, memory_mask)
+ 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 - alpha) * tgt_orig
+
+ return tgt
+
+
+def _get_activation_fn(activation: str):
+ if activation == "relu":
+ return nn.functional.relu
+ elif activation == "gelu":
+ return nn.functional.gelu
+
+ raise RuntimeError(
+ "activation should be relu/gelu, not {}".format(activation)
+ )
+
+
+class TransformerEncoder(nn.Module):
+ r"""TransformerEncoder is a stack of N encoder layers
+
+ Args:
+ encoder_layer: an instance of the TransformerEncoderLayer() class (required).
+ num_layers: the number of sub-encoder-layers in the encoder (required).
+
+ 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) -> None:
+ 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,
+ mask: Optional[torch.Tensor] = None,
+ src_key_padding_mask: Optional[torch.Tensor] = None,
+ warmup: float = 1.0,
+ ) -> torch.Tensor:
+ r"""Pass the input through the encoder layers in turn.
+
+ Args:
+ src: the sequence to the encoder (required).
+ mask: the mask for the src sequence (optional).
+ src_key_padding_mask: the mask for the src keys per batch (optional).
+
+ Shape:
+ src: (S, N, E).
+ mask: (S, S).
+ src_key_padding_mask: (N, S).
+ S is the source sequence length, T is the target sequence length, N is the batch size, E is the feature number
+
+ """
+ output = src
+
+ for i, mod in enumerate(self.layers):
+ output = mod(
+ output,
+ src_mask=mask,
+ src_key_padding_mask=src_key_padding_mask,
+ warmup=warmup,
+ )
+
+ return output
+
+
+class TransformerDecoder(nn.Module):
+ r"""TransformerDecoder is a stack of N decoder layers
+
+ Args:
+ decoder_layer: an instance of the TransformerDecoderLayer() class (required).
+ num_layers: the number of sub-decoder-layers in the decoder (required).
+
+ 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(10, 32, 512)
+ >>> out = transformer_decoder(tgt, memory)
+ """
+
+ def __init__(self, decoder_layer: nn.Module, num_layers: int) -> None:
+ super().__init__()
+ self.layers = nn.ModuleList(
+ [copy.deepcopy(decoder_layer) for i in range(num_layers)]
+ )
+ self.num_layers = num_layers
+
+ 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:
+ r"""Pass the input through the decoder layers in turn.
+
+ Args:
+ tgt: the sequence to the decoder (required).
+ memory: the sequence from the last layer of the encoder (required).
+ tgt_mask: the mask for the tgt sequence (optional).
+ memory_mask: the mask for the memory sequence (optional).
+ tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
+ memory_key_padding_mask: the mask for the memory keys per batch (optional).
+
+ Shape:
+ tgt: (S, N, E).
+ tgt_mask: (S, S).
+ tgt_key_padding_mask: (N, S).
+
+ """
+ output = tgt
+
+ 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,
+ )
+
+ 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).
+ Returns:
+ Return None.
+ """
+ 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:
+ Its shape is (N, T, C)
+
+ Returns:
+ Return a tensor of shape (N, T, C)
+ """
+ self.extend_pe(x)
+ x = x * self.xscale + self.pe[:, : x.size(1), :]
+ return self.dropout(x)
+
+
+class Noam(object):
+ """
+ Implements Noam optimizer.
+
+ Proposed in
+ "Attention Is All You Need", https://arxiv.org/pdf/1706.03762.pdf
+
+ Modified from
+ https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/transformer/optimizer.py # noqa
+
+ Args:
+ params:
+ iterable of parameters to optimize or dicts defining parameter groups
+ model_size:
+ attention dimension of the transformer model
+ factor:
+ learning rate factor
+ warm_step:
+ warmup steps
+ """
+
+ def __init__(
+ self,
+ params,
+ model_size: int = 256,
+ factor: float = 10.0,
+ warm_step: int = 25000,
+ weight_decay=0,
+ ) -> None:
+ """Construct an Noam object."""
+ self.optimizer = torch.optim.Adam(
+ params, lr=0, betas=(0.9, 0.98), eps=1e-9, weight_decay=weight_decay
+ )
+ self._step = 0
+ self.warmup = warm_step
+ self.factor = factor
+ self.model_size = model_size
+ self._rate = 0
+
+ @property
+ def param_groups(self):
+ """Return param_groups."""
+ return self.optimizer.param_groups
+
+ def step(self):
+ """Update parameters and rate."""
+ self._step += 1
+ rate = self.rate()
+ for p in self.optimizer.param_groups:
+ p["lr"] = rate
+ self._rate = rate
+ self.optimizer.step()
+
+ def rate(self, step=None):
+ """Implement `lrate` above."""
+ if step is None:
+ step = self._step
+ return (
+ self.factor
+ * self.model_size ** (-0.5)
+ * min(step ** (-0.5), step * self.warmup ** (-1.5))
+ )
+
+ def zero_grad(self):
+ """Reset gradient."""
+ self.optimizer.zero_grad()
+
+ def state_dict(self):
+ """Return state_dict."""
+ return {
+ "_step": self._step,
+ "warmup": self.warmup,
+ "factor": self.factor,
+ "model_size": self.model_size,
+ "_rate": self._rate,
+ "optimizer": self.optimizer.state_dict(),
+ }
+
+ def load_state_dict(self, state_dict):
+ """Load state_dict."""
+ for key, value in state_dict.items():
+ if key == "optimizer":
+ self.optimizer.load_state_dict(state_dict["optimizer"])
+ else:
+ setattr(self, key, value)
+
+
+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:
+ Tensor: Mask tensor of dimension (batch_size, input_length),
+ True denote the masked indices.
+ """
+ if supervisions is None:
+ return None
+
+ supervision_segments = torch.stack(
+ (
+ supervisions["sequence_idx"],
+ supervisions["start_frame"],
+ supervisions["num_frames"],
+ ),
+ 1,
+ ).to(torch.int32)
+
+ lengths = [
+ 0 for _ in range(int(supervision_segments[:, 0].max().item()) + 1)
+ ]
+ for idx in range(supervision_segments.size(0)):
+ # Note: TorchScript doesn't allow to unpack tensors as tuples
+ sequence_idx = supervision_segments[idx, 0].item()
+ start_frame = supervision_segments[idx, 1].item()
+ num_frames = supervision_segments[idx, 2].item()
+ lengths[sequence_idx] = start_frame + num_frames
+
+ lengths = [((i - 1) // 2 - 1) // 2 for i in lengths]
+ bs = int(len(lengths))
+ seq_range = torch.arange(0, max_len, dtype=torch.int64)
+ seq_range_expand = seq_range.unsqueeze(0).expand(bs, max_len)
+ # Note: TorchScript doesn't implement Tensor.new()
+ seq_length_expand = torch.tensor(
+ lengths, device=seq_range_expand.device, dtype=seq_range_expand.dtype
+ ).unsqueeze(-1)
+ mask = seq_range_expand >= seq_length_expand
+
+ return mask
+
+
+def decoder_padding_mask(
+ ys_pad: 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:
+ Tensor:
+ 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 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-list 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-list, 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/librispeech/ASR/conformer_mmi/conformer.py b/egs/librispeech/ASR/conformer_mmi/conformer.py
index b19b94db1..97c8d83a2 100644
--- a/egs/librispeech/ASR/conformer_mmi/conformer.py
+++ b/egs/librispeech/ASR/conformer_mmi/conformer.py
@@ -247,7 +247,9 @@ class ConformerEncoderLayer(nn.Module):
residual = src
if self.normalize_before:
src = self.norm_conv(src)
- src = residual + self.dropout(self.conv_module(src))
+ src = residual + self.dropout(
+ self.conv_module(src, src_key_padding_mask=src_key_padding_mask)
+ )
if not self.normalize_before:
src = self.norm_conv(src)
@@ -363,7 +365,7 @@ class RelPositionalEncoding(torch.nn.Module):
):
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
+ # 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 Tensor:
+ 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).
Returns:
Tensor: Output tensor (#time, batch, channels).
@@ -896,6 +903,8 @@ class ConvolutionModule(nn.Module):
x = nn.functional.glu(x, dim=1) # (batch, channels, time)
# 1D Depthwise Conv
+ if src_key_padding_mask is not None:
+ x.masked_fill_(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
x = self.depthwise_conv(x)
x = self.activation(self.norm(x))
diff --git a/egs/librispeech/ASR/conformer_mmi/decode.py b/egs/librispeech/ASR/conformer_mmi/decode.py
index a77168b62..fc9861489 100755
--- a/egs/librispeech/ASR/conformer_mmi/decode.py
+++ b/egs/librispeech/ASR/conformer_mmi/decode.py
@@ -404,7 +404,7 @@ def decode_dataset(
sos_id: int,
eos_id: int,
G: Optional[k2.Fsa] = None,
-) -> Dict[str, List[Tuple[List[str], List[str]]]]:
+) -> Dict[str, List[Tuple[str, List[str], List[str]]]]:
"""Decode dataset.
Args:
@@ -449,6 +449,7 @@ def decode_dataset(
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,
@@ -466,9 +467,9 @@ def decode_dataset(
for lm_scale, hyps in hyps_dict.items():
this_batch = []
assert len(hyps) == len(texts)
- for hyp_words, ref_text in zip(hyps, texts):
+ for cut_id, hyp_words, ref_text in zip(cut_ids, hyps, texts):
ref_words = ref_text.split()
- this_batch.append((ref_words, hyp_words))
+ this_batch.append((cut_id, ref_words, hyp_words))
results[lm_scale].extend(this_batch)
@@ -486,7 +487,7 @@ def decode_dataset(
def save_results(
params: AttributeDict,
test_set_name: str,
- results_dict: Dict[str, List[Tuple[List[int], List[int]]]],
+ results_dict: Dict[str, List[Tuple[str, List[str], List[str]]]],
):
if params.method == "attention-decoder":
# Set it to False since there are too many logs.
@@ -496,6 +497,7 @@ def save_results(
test_set_wers = dict()
for key, results in results_dict.items():
recog_path = params.exp_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}")
@@ -661,6 +663,8 @@ def main():
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)
# CAUTION: `test_sets` is for displaying only.
# If you want to skip test-clean, you have to skip
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless/asr_datamodule.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless/asr_datamodule.py
new file mode 120000
index 000000000..a074d6085
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless/asr_datamodule.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/asr_datamodule.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless/beam_search.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless/beam_search.py
new file mode 120000
index 000000000..8554e44cc
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless/beam_search.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/beam_search.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless/decode.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless/decode.py
new file mode 100755
index 000000000..620d69a19
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless/decode.py
@@ -0,0 +1,661 @@
+#!/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
+./conv_emformer_transducer_stateless/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless/exp \
+ --max-duration 300 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method greedy_search \
+ --use-averaged-model True
+
+(2) modified beam search
+./conv_emformer_transducer_stateless/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless/exp \
+ --max-duration 300 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method modified_beam_search \
+ --use-averaged-model True \
+ --beam-size 4
+
+(3) fast beam search
+./conv_emformer_transducer_stateless/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless/exp \
+ --max-duration 300 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method fast_beam_search \
+ --use-averaged-model True \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8
+"""
+
+
+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_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.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=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_stateless4/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="""Possible values are:
+ - greedy_search
+ - modified_beam_search
+ - fast_beam_search
+ """,
+ )
+
+ 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=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=8,
+ 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(
+ "--max-sym-per-frame",
+ type=int,
+ default=1,
+ help="""Maximum number of symbols per frame.
+ Used only when --decoding_method is greedy_search""",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def decode_one_batch(
+ params: AttributeDict,
+ model: nn.Module,
+ sp: spm.SentencePieceProcessor,
+ 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.
+ sp:
+ The BPE 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 = 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 += params.chunk_length
+ feature = torch.nn.functional.pad(
+ feature,
+ pad=(0, 0, 0, params.chunk_length),
+ 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 == "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 params.decoding_method == "fast_beam_search":
+ return {
+ (
+ f"beam_{params.beam}_"
+ f"max_contexts_{params.max_contexts}_"
+ f"max_states_{params.max_states}"
+ ): 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,
+ 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.
+ 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"]
+ 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,
+ 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",
+ "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}"
+ 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 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 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
+ start = params.epoch - params.avg
+ assert start >= 1
+ 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 params.decoding_method == "fast_beam_search":
+ 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}")
+
+ # 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,
+ 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/conv_emformer_transducer_stateless/decoder.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless/decoder.py
new file mode 120000
index 000000000..0793c5709
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless/decoder.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/decoder.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless/emformer.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless/emformer.py
new file mode 100644
index 000000000..8ca7d5568
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless/emformer.py
@@ -0,0 +1,1935 @@
+# Copyright 2022 Xiaomi Corporation (Author: 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.
+#
+# It is modified based on
+# 1) https://github.com/pytorch/audio/blob/main/torchaudio/models/emformer.py # noqa
+# 2) https://github.com/pytorch/audio/blob/main/torchaudio/prototype/models/conv_emformer.py # noqa
+
+import math
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from encoder_interface import EncoderInterface
+from scaling import (
+ ActivationBalancer,
+ BasicNorm,
+ DoubleSwish,
+ ScaledConv1d,
+ ScaledConv2d,
+ ScaledLinear,
+)
+
+from icefall.utils import make_pad_mask
+
+
+LOG_EPSILON = math.log(1e-10)
+
+
+def unstack_states(
+ states: Tuple[List[List[torch.Tensor]], List[torch.Tensor]]
+) -> List[Tuple[List[List[torch.Tensor]], List[torch.Tensor]]]:
+ """Unstack the emformer 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.
+
+ Args:
+ states:
+ A tuple of 2 elements.
+ ``states[0]`` is the attention caches of a batch of utterance.
+ ``states[1]`` is the convolution caches of a batch of utterance.
+ ``len(states[0])`` and ``len(states[1])`` both eqaul to number of layers. # noqa
+
+ Returns:
+ A list of states.
+ ``states[i]`` is a tuple of 2 elements of i-th utterance.
+ ``states[i][0]`` is the attention caches of i-th utterance.
+ ``states[i][1]`` is the convolution caches of i-th utterance.
+ ``len(states[i][0])`` and ``len(states[i][1])`` both eqaul to number of layers. # noqa
+ """
+
+ attn_caches, conv_caches = states
+ batch_size = conv_caches[0].size(0)
+ num_layers = len(attn_caches)
+
+ list_attn_caches = [None] * batch_size
+ for i in range(batch_size):
+ list_attn_caches[i] = [[] for _ in range(num_layers)]
+ for li, layer in enumerate(attn_caches):
+ for s in layer:
+ s_list = s.unbind(dim=1)
+ for bi, b in enumerate(list_attn_caches):
+ b[li].append(s_list[bi])
+
+ list_conv_caches = [None] * batch_size
+ for i in range(batch_size):
+ list_conv_caches[i] = [None] * num_layers
+ for li, layer in enumerate(conv_caches):
+ c_list = layer.unbind(dim=0)
+ for bi, b in enumerate(list_conv_caches):
+ b[li] = c_list[bi]
+
+ ans = [None] * batch_size
+ for i in range(batch_size):
+ ans[i] = [list_attn_caches[i], list_conv_caches[i]]
+
+ return ans
+
+
+def stack_states(
+ state_list: List[Tuple[List[List[torch.Tensor]], List[torch.Tensor]]]
+) -> Tuple[List[List[torch.Tensor]], List[torch.Tensor]]:
+ """Stack list of emformer states that correspond to separate utterances
+ into a single emformer state so that it can be used as an input for
+ emformer 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 emformer model for a single utterance.
+ ``states[i]`` is a tuple of 2 elements of i-th utterance.
+ ``states[i][0]`` is the attention caches of i-th utterance.
+ ``states[i][1]`` is the convolution caches of i-th utterance.
+ ``len(states[i][0])`` and ``len(states[i][1])`` both eqaul to number of layers. # noqa
+
+ 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)
+
+ attn_caches = []
+ for layer in state_list[0][0]:
+ if batch_size > 1:
+ # Note: We will stack attn_caches[layer][s][] later to get attn_caches[layer][s] # noqa
+ attn_caches.append([[s] for s in layer])
+ else:
+ attn_caches.append([s.unsqueeze(1) for s in layer])
+ for b, states in enumerate(state_list[1:], 1):
+ for li, layer in enumerate(states[0]):
+ for si, s in enumerate(layer):
+ attn_caches[li][si].append(s)
+ if b == batch_size - 1:
+ attn_caches[li][si] = torch.stack(
+ attn_caches[li][si], dim=1
+ )
+
+ conv_caches = []
+ for layer in state_list[0][1]:
+ if batch_size > 1:
+ # Note: We will stack conv_caches[layer][] later to get conv_caches[layer] # noqa
+ conv_caches.append([layer])
+ else:
+ conv_caches.append(layer.unsqueeze(0))
+ for b, states in enumerate(state_list[1:], 1):
+ for li, layer in enumerate(states[1]):
+ conv_caches[li].append(layer)
+ if b == batch_size - 1:
+ conv_caches[li] = torch.stack(conv_caches[li], dim=0)
+
+ return [attn_caches, conv_caches]
+
+
+class ConvolutionModule(nn.Module):
+ """ConvolutionModule.
+
+ Modified from https://github.com/pytorch/audio/blob/main/torchaudio/prototype/models/conv_emformer.py # noqa
+
+ Args:
+ chunk_length (int):
+ Length of each chunk.
+ right_context_length (int):
+ Length of right context.
+ channels (int):
+ The number of input channels and output 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,
+ chunk_length: int,
+ right_context_length: int,
+ channels: int,
+ kernel_size: int,
+ bias: bool = True,
+ ) -> None:
+ """Construct an ConvolutionModule object."""
+ super().__init__()
+ # kernerl_size should be an odd number for 'SAME' padding
+ assert (kernel_size - 1) % 2 == 0, kernel_size
+
+ self.chunk_length = chunk_length
+ self.right_context_length = right_context_length
+ self.channels = channels
+
+ 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
+ )
+
+ # make it causal by padding cached (kernel_size - 1) frames on the left
+ self.cache_size = kernel_size - 1
+ self.depthwise_conv = ScaledConv1d(
+ channels,
+ channels,
+ kernel_size,
+ stride=1,
+ padding=0,
+ 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 _split_right_context(
+ self,
+ pad_utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ ) -> torch.Tensor:
+ """
+ Args:
+ pad_utterance:
+ Its shape is (cache_size + U, B, D).
+ right_context:
+ Its shape is (R, B, D).
+
+ Returns:
+ Right context segments padding with corresponding context.
+ Its shape is (num_segs * B, D, cache_size + right_context_length).
+ """
+ U_, B, D = pad_utterance.size()
+ R = right_context.size(0)
+ assert self.right_context_length != 0
+ assert R % self.right_context_length == 0
+ num_chunks = R // self.right_context_length
+ right_context = right_context.reshape(
+ num_chunks, self.right_context_length, B, D
+ )
+ right_context = right_context.permute(0, 2, 1, 3).reshape(
+ num_chunks * B, self.right_context_length, D
+ )
+
+ intervals = torch.arange(
+ 0, self.chunk_length * (num_chunks - 1), self.chunk_length
+ )
+ first = torch.arange(
+ self.chunk_length, self.chunk_length + self.cache_size
+ )
+ indexes = intervals.unsqueeze(1) + first.unsqueeze(0)
+ indexes = torch.cat(
+ [indexes, torch.arange(U_ - self.cache_size, U_).unsqueeze(0)]
+ )
+ padding = pad_utterance[indexes] # (num_chunks, cache_size, B, D)
+ padding = padding.permute(0, 2, 1, 3).reshape(
+ num_chunks * B, self.cache_size, D
+ )
+
+ pad_right_context = torch.cat([padding, right_context], dim=1)
+ # (num_chunks * B, cache_size + right_context_length, D)
+ return pad_right_context.permute(0, 2, 1)
+
+ def _merge_right_context(
+ self, right_context: torch.Tensor, B: int
+ ) -> torch.Tensor:
+ """
+ Args:
+ right_context:
+ Right context segments.
+ It shape is (num_segs * B, D, right_context_length).
+ B:
+ Batch size.
+
+ Returns:
+ A tensor of shape (B, D, R), where
+ R = num_segs * right_context_length.
+ """
+ right_context = right_context.reshape(
+ -1, B, self.channels, self.right_context_length
+ )
+ right_context = right_context.permute(1, 2, 0, 3)
+ right_context = right_context.reshape(B, self.channels, -1)
+ return right_context
+
+ def forward(
+ self,
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Causal convolution module.
+
+ Args:
+ utterance (torch.Tensor):
+ Utterance tensor of shape (U, B, D).
+ right_context (torch.Tensor):
+ Right context tensor of shape (R, B, D).
+
+ Returns:
+ A tuple of 2 tensors:
+ - output utterance of shape (U, B, D).
+ - output right_context of shape (R, B, D).
+ """
+ U, B, D = utterance.size()
+ R, _, _ = right_context.size()
+
+ # point-wise conv and GLU mechanism
+ x = torch.cat([right_context, utterance], dim=0) # (R + U, B, D)
+ x = x.permute(1, 2, 0) # (B, D, R + U)
+ x = self.pointwise_conv1(x) # (B, 2 * D, R + U)
+ x = self.deriv_balancer1(x)
+ x = nn.functional.glu(x, dim=1) # (B, D, R + U)
+ utterance = x[:, :, R:] # (B, D, U)
+ right_context = x[:, :, :R] # (B, D, R)
+
+ # make causal convolution
+ cache = torch.zeros(
+ B, D, self.cache_size, device=x.device, dtype=x.dtype
+ )
+ pad_utterance = torch.cat(
+ [cache, utterance], dim=2
+ ) # (B, D, cache + U)
+
+ # depth-wise conv on utterance
+ utterance = self.depthwise_conv(pad_utterance) # (B, D, U)
+
+ if self.right_context_length > 0:
+ # depth-wise conv on right_context
+ pad_right_context = self._split_right_context(
+ pad_utterance.permute(2, 0, 1), right_context.permute(2, 0, 1)
+ ) # (num_segs * B, D, cache_size + right_context_length)
+ right_context = self.depthwise_conv(
+ pad_right_context
+ ) # (num_segs * B, D, right_context_length)
+ right_context = self._merge_right_context(
+ right_context, B
+ ) # (B, D, R)
+
+ x = torch.cat([right_context, utterance], dim=2) # (B, D, R + U)
+ x = self.deriv_balancer2(x)
+ x = self.activation(x)
+
+ # point-wise conv
+ x = self.pointwise_conv2(x) # (B, D, R + U)
+
+ right_context = x[:, :, :R] # (B, D, R)
+ utterance = x[:, :, R:] # (B, D, U)
+ return (
+ utterance.permute(2, 0, 1),
+ right_context.permute(2, 0, 1),
+ )
+
+ @torch.jit.export
+ def infer(
+ self,
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ cache: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Causal convolution module applied on both utterance and right_context.
+
+ Args:
+ utterance (torch.Tensor):
+ Utterance tensor of shape (U, B, D).
+ right_context (torch.Tensor):
+ Right context tensor of shape (R, B, D).
+ cache (torch.Tensor, optional):
+ Cached tensor for left padding of shape (B, D, cache_size).
+
+ Returns:
+ A tuple of 3 tensors:
+ - output utterance of shape (U, B, D).
+ - output right_context of shape (R, B, D).
+ - updated cache tensor of shape (B, D, cache_size).
+ """
+ U, B, D = utterance.size()
+ R, _, _ = right_context.size()
+
+ # point-wise conv
+ x = torch.cat([utterance, right_context], dim=0) # (U + R, B, D)
+ x = x.permute(1, 2, 0) # (B, D, U + R)
+ x = self.pointwise_conv1(x) # (B, 2 * D, U + R)
+ x = self.deriv_balancer1(x)
+ x = nn.functional.glu(x, dim=1) # (B, D, U + R)
+
+ # make causal convolution
+ assert cache.shape == (B, D, self.cache_size), cache.shape
+ x = torch.cat([cache, x], dim=2) # (B, D, cache_size + U + R)
+ # update cache
+ new_cache = x[:, :, -R - self.cache_size : -R]
+
+ # 1-D depth-wise conv
+ x = self.depthwise_conv(x) # (B, D, U + R)
+
+ x = self.deriv_balancer2(x)
+ x = self.activation(x)
+
+ # point-wise conv
+ x = self.pointwise_conv2(x) # (B, D, U + R)
+
+ utterance = x[:, :, :U] # (B, D, U)
+ right_context = x[:, :, U:] # (B, D, R)
+ return (
+ utterance.permute(2, 0, 1),
+ right_context.permute(2, 0, 1),
+ new_cache,
+ )
+
+
+class EmformerAttention(nn.Module):
+ r"""Emformer layer attention module.
+
+ Args:
+ embed_dim (int):
+ Embedding dimension.
+ nhead (int):
+ Number of attention heads in each Emformer layer.
+ dropout (float, optional):
+ Dropout probability. (Default: 0.0)
+ tanh_on_mem (bool, optional):
+ If ``True``, applies tanh to memory elements. (Default: ``False``)
+ negative_inf (float, optional):
+ Value to use for negative infinity in attention weights. (Default: -1e8)
+ """
+
+ def __init__(
+ self,
+ embed_dim: int,
+ nhead: int,
+ dropout: float = 0.0,
+ tanh_on_mem: bool = False,
+ negative_inf: float = -1e8,
+ ):
+ super().__init__()
+
+ if embed_dim % nhead != 0:
+ raise ValueError(
+ f"embed_dim ({embed_dim}) is not a multiple of"
+ f"nhead ({nhead})."
+ )
+
+ self.embed_dim = embed_dim
+ self.nhead = nhead
+ self.tanh_on_mem = tanh_on_mem
+ self.negative_inf = negative_inf
+ self.head_dim = embed_dim // nhead
+ self.dropout = dropout
+
+ self.emb_to_key_value = ScaledLinear(
+ embed_dim, 2 * embed_dim, bias=True
+ )
+ self.emb_to_query = ScaledLinear(embed_dim, embed_dim, bias=True)
+ self.out_proj = ScaledLinear(
+ embed_dim, embed_dim, bias=True, initial_scale=0.25
+ )
+
+ def _gen_attention_probs(
+ self,
+ attention_weights: torch.Tensor,
+ attention_mask: torch.Tensor,
+ padding_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """Given the entire attention weights, mask out unecessary connections
+ and optionally with padding positions, to obtain underlying chunk-wise
+ attention probabilities.
+
+ B: batch size;
+ Q: length of query;
+ KV: length of key and value.
+
+ Args:
+ attention_weights (torch.Tensor):
+ Attention weights computed on the entire concatenated tensor
+ with shape (B * nhead, Q, KV).
+ attention_mask (torch.Tensor):
+ Mask tensor where chunk-wise connections are filled with `False`,
+ and other unnecessary connections are filled with `True`,
+ with shape (Q, KV).
+ padding_mask (torch.Tensor, optional):
+ Mask tensor where the padding positions are fill with `True`,
+ and other positions are filled with `False`, with shapa `(B, KV)`.
+
+ Returns:
+ A tensor of shape (B * nhead, Q, KV).
+ """
+ attention_weights_float = attention_weights.float()
+ attention_weights_float = attention_weights_float.masked_fill(
+ attention_mask.unsqueeze(0), self.negative_inf
+ )
+ if padding_mask is not None:
+ Q = attention_weights.size(1)
+ B = attention_weights.size(0) // self.nhead
+ attention_weights_float = attention_weights_float.view(
+ B, self.nhead, Q, -1
+ )
+ attention_weights_float = attention_weights_float.masked_fill(
+ padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool),
+ self.negative_inf,
+ )
+ attention_weights_float = attention_weights_float.view(
+ B * self.nhead, Q, -1
+ )
+
+ attention_probs = nn.functional.softmax(
+ attention_weights_float, dim=-1
+ ).type_as(attention_weights)
+
+ attention_probs = nn.functional.dropout(
+ attention_probs, p=self.dropout, training=self.training
+ )
+ return attention_probs
+
+ def _forward_impl(
+ self,
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ summary: torch.Tensor,
+ memory: torch.Tensor,
+ attention_mask: torch.Tensor,
+ padding_mask: Optional[torch.Tensor] = None,
+ left_context_key: Optional[torch.Tensor] = None,
+ left_context_val: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Underlying chunk-wise attention implementation."""
+ U, B, _ = utterance.size()
+ R = right_context.size(0)
+ M = memory.size(0)
+ scaling = float(self.head_dim) ** -0.5
+
+ # compute query with [right_context, utterance, summary].
+ query = self.emb_to_query(
+ torch.cat([right_context, utterance, summary])
+ )
+ # compute key and value with [memory, right_context, utterance].
+ key, value = self.emb_to_key_value(
+ torch.cat([memory, right_context, utterance])
+ ).chunk(chunks=2, dim=2)
+
+ if left_context_key is not None and left_context_val is not None:
+ # now compute key and value with
+ # [memory, right context, left context, uttrance]
+ # this is used in inference mode
+ key = torch.cat([key[: M + R], left_context_key, key[M + R :]])
+ value = torch.cat(
+ [value[: M + R], left_context_val, value[M + R :]]
+ )
+ Q = query.size(0)
+ # KV = key.size(0)
+
+ reshaped_query, reshaped_key, reshaped_value = [
+ tensor.contiguous()
+ .view(-1, B * self.nhead, self.head_dim)
+ .transpose(0, 1)
+ for tensor in [query, key, value]
+ ] # (B * nhead, Q or KV, head_dim)
+ attention_weights = torch.bmm(
+ reshaped_query * scaling, reshaped_key.transpose(1, 2)
+ ) # (B * nhead, Q, KV)
+
+ # compute attention probabilities
+ attention_probs = self._gen_attention_probs(
+ attention_weights, attention_mask, padding_mask
+ )
+
+ # compute attention outputs
+ attention = torch.bmm(attention_probs, reshaped_value)
+ assert attention.shape == (B * self.nhead, Q, self.head_dim)
+ attention = (
+ attention.transpose(0, 1).contiguous().view(Q, B, self.embed_dim)
+ )
+
+ # apply output projection
+ outputs = self.out_proj(attention)
+
+ output_right_context_utterance = outputs[: R + U]
+ output_memory = outputs[R + U :]
+ if self.tanh_on_mem:
+ output_memory = torch.tanh(output_memory)
+ else:
+ output_memory = torch.clamp(output_memory, min=-10, max=10)
+
+ return output_right_context_utterance, output_memory, key, value
+
+ def forward(
+ self,
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ summary: torch.Tensor,
+ memory: torch.Tensor,
+ attention_mask: torch.Tensor,
+ padding_mask: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ # TODO: Modify docs.
+ """Forward pass for training and validation mode.
+
+ B: batch size;
+ D: embedding dimension;
+ R: length of the hard-copied right contexts;
+ U: length of full utterance;
+ S: length of summary vectors;
+ M: length of memory vectors.
+
+ It computes a `big` attention matrix on full utterance and
+ then utilizes a pre-computed mask to simulate chunk-wise attention.
+
+ It concatenates three blocks: hard-copied right contexts,
+ full utterance, and summary vectors, as a `big` block,
+ to compute the query tensor:
+ query = [right_context, utterance, summary],
+ with length Q = R + U + S.
+ It concatenates the three blocks: memory vectors,
+ hard-copied right contexts, and full utterance as another `big` block,
+ to compute the key and value tensors:
+ key & value = [memory, right_context, utterance],
+ with length KV = M + R + U.
+ Attention scores is computed with above `big` query and key.
+
+ Then the underlying chunk-wise attention is obtained by applying
+ the attention mask. Suppose
+ c_i: chunk at index i;
+ r_i: right context that c_i can use;
+ l_i: left context that c_i can use;
+ m_i: past memory vectors from previous layer that c_i can use;
+ s_i: summary vector of c_i;
+ The target chunk-wise attention is:
+ c_i, r_i (in query) -> l_i, c_i, r_i, m_i (in key);
+ s_i (in query) -> l_i, c_i, r_i (in key).
+
+ Args:
+ utterance (torch.Tensor):
+ Full utterance frames, with shape (U, B, D).
+ right_context (torch.Tensor):
+ Hard-copied right context frames, with shape (R, B, D),
+ where R = num_chunks * right_context_length
+ summary (torch.Tensor):
+ Summary elements with shape (S, B, D), where S = num_chunks.
+ It is an empty tensor without using memory.
+ memory (torch.Tensor):
+ Memory elements, with shape (M, B, D), where M = num_chunks - 1.
+ It is an empty tensor without using memory.
+ attention_mask (torch.Tensor):
+ Pre-computed attention mask to simulate underlying chunk-wise
+ attention, with shape (Q, KV).
+ padding_mask (torch.Tensor):
+ Padding mask of key tensor, with shape (B, KV).
+
+ Returns:
+ A tuple containing 2 tensors:
+ - output of right context and utterance, with shape (R + U, B, D).
+ - memory output, with shape (M, B, D), where M = S - 1 or M = 0.
+ """
+ (
+ output_right_context_utterance,
+ output_memory,
+ _,
+ _,
+ ) = self._forward_impl(
+ utterance,
+ right_context,
+ summary,
+ memory,
+ attention_mask,
+ padding_mask=padding_mask,
+ )
+ return output_right_context_utterance, output_memory[:-1]
+
+ @torch.jit.export
+ def infer(
+ self,
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ summary: torch.Tensor,
+ memory: torch.Tensor,
+ left_context_key: torch.Tensor,
+ left_context_val: torch.Tensor,
+ padding_mask: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Forward pass for inference.
+
+ B: batch size;
+ D: embedding dimension;
+ R: length of right context;
+ U: length of utterance, i.e., current chunk;
+ L: length of cached left context;
+ S: length of summary vectors, S = 1;
+ M: length of cached memory vectors.
+
+ It concatenates the right context, utterance (i.e., current chunk)
+ and summary vector of current chunk, to compute the query tensor:
+ query = [right_context, utterance, summary],
+ with length Q = R + U + S.
+ It concatenates the memory vectors, right context, left context, and
+ current chunk, to compute the key and value tensors:
+ key & value = [memory, right_context, left_context, utterance],
+ with length KV = M + R + L + U.
+
+ The chunk-wise attention is:
+ chunk, right context (in query) ->
+ left context, chunk, right context, memory vectors (in key);
+ summary (in query) -> left context, chunk, right context (in key).
+
+ Args:
+ utterance (torch.Tensor):
+ Current chunk frames, with shape (U, B, D), where U = chunk_length.
+ right_context (torch.Tensor):
+ Right context frames, with shape (R, B, D),
+ where R = right_context_length.
+ summary (torch.Tensor):
+ Summary vector with shape (1, B, D), or empty tensor.
+ memory (torch.Tensor):
+ Memory vectors, with shape (M, B, D), or empty tensor.
+ left_context_key (torch,Tensor):
+ Cached attention key of left context from preceding computation,
+ with shape (L, B, D).
+ left_context_val (torch.Tensor):
+ Cached attention value of left context from preceding computation,
+ with shape (L, B, D).
+ padding_mask (torch.Tensor):
+ Padding mask of key tensor, with shape (B, KV).
+
+ Returns:
+ A tuple containing 4 tensors:
+ - output of right context and utterance, with shape (R + U, B, D).
+ - memory output, with shape (1, B, D) or (0, B, D).
+ - attention key of left context and utterance, which would be cached
+ for next computation, with shape (L + U, B, D).
+ - attention value of left context and utterance, which would be
+ cached for next computation, with shape (L + U, B, D).
+ """
+ U = utterance.size(0)
+ R = right_context.size(0)
+ L = left_context_key.size(0)
+ S = summary.size(0)
+ M = memory.size(0)
+
+ # TODO: move it outside
+ # query = [right context, utterance, summary]
+ Q = R + U + S
+ # key, value = [memory, right context, left context, uttrance]
+ KV = M + R + L + U
+ attention_mask = torch.zeros(Q, KV).to(
+ dtype=torch.bool, device=utterance.device
+ )
+ # disallow attention bettween the summary vector with the memory bank
+ attention_mask[-1, :M] = True
+ (
+ output_right_context_utterance,
+ output_memory,
+ key,
+ value,
+ ) = self._forward_impl(
+ utterance,
+ right_context,
+ summary,
+ memory,
+ attention_mask,
+ padding_mask=padding_mask,
+ left_context_key=left_context_key,
+ left_context_val=left_context_val,
+ )
+ return (
+ output_right_context_utterance,
+ output_memory,
+ key[M + R :],
+ value[M + R :],
+ )
+
+
+class EmformerEncoderLayer(nn.Module):
+ """Emformer layer that constitutes Emformer.
+
+ Args:
+ d_model (int):
+ Input dimension.
+ nhead (int):
+ Number of attention heads.
+ dim_feedforward (int):
+ Hidden layer dimension of feedforward network.
+ chunk_length (int):
+ Length of each input segment.
+ dropout (float, optional):
+ Dropout probability. (Default: 0.0)
+ layer_dropout (float, optional):
+ Layer dropout probability. (Default: 0.0)
+ cnn_module_kernel (int):
+ Kernel size of convolution module.
+ left_context_length (int, optional):
+ Length of left context. (Default: 0)
+ right_context_length (int, optional):
+ Length of right context. (Default: 0)
+ memory_size (int, optional):
+ Number of memory elements to use. (Default: 0)
+ tanh_on_mem (bool, optional):
+ If ``True``, applies tanh to memory elements. (Default: ``False``)
+ negative_inf (float, optional):
+ Value to use for negative infinity in attention weights. (Default: -1e8)
+ """
+
+ def __init__(
+ self,
+ d_model: int,
+ nhead: int,
+ dim_feedforward: int,
+ chunk_length: int,
+ dropout: float = 0.1,
+ layer_dropout: float = 0.075,
+ cnn_module_kernel: int = 31,
+ left_context_length: int = 0,
+ right_context_length: int = 0,
+ memory_size: int = 0,
+ tanh_on_mem: bool = False,
+ negative_inf: float = -1e8,
+ ):
+ super().__init__()
+
+ self.attention = EmformerAttention(
+ embed_dim=d_model,
+ nhead=nhead,
+ dropout=dropout,
+ tanh_on_mem=tanh_on_mem,
+ negative_inf=negative_inf,
+ )
+ self.summary_op = nn.AvgPool1d(
+ kernel_size=chunk_length, stride=chunk_length, ceil_mode=True
+ )
+
+ 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.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.conv_module = ConvolutionModule(
+ chunk_length,
+ right_context_length,
+ 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)
+
+ self.layer_dropout = layer_dropout
+ self.left_context_length = left_context_length
+ self.chunk_length = chunk_length
+ self.memory_size = memory_size
+ self.d_model = d_model
+ self.use_memory = memory_size > 0
+
+ def _update_attn_cache(
+ self,
+ next_key: torch.Tensor,
+ next_val: torch.Tensor,
+ memory: torch.Tensor,
+ attn_cache: List[torch.Tensor],
+ ) -> List[torch.Tensor]:
+ """Update cached attention state:
+ 1) output memory of current chunk in the lower layer;
+ 2) attention key and value in current chunk's computation, which would
+ be resued in next chunk's computation.
+ """
+ new_memory = torch.cat([attn_cache[0], memory])
+ new_key = torch.cat([attn_cache[1], next_key])
+ new_val = torch.cat([attn_cache[2], next_val])
+ attn_cache[0] = new_memory[new_memory.size(0) - self.memory_size :]
+ attn_cache[1] = new_key[new_key.size(0) - self.left_context_length :]
+ attn_cache[2] = new_val[new_val.size(0) - self.left_context_length :]
+ return attn_cache
+
+ def _apply_conv_module_forward(
+ self,
+ right_context_utterance: torch.Tensor,
+ R: int,
+ ) -> torch.Tensor:
+ """Apply convolution module in training and validation mode."""
+ utterance = right_context_utterance[R:]
+ right_context = right_context_utterance[:R]
+ utterance, right_context = self.conv_module(utterance, right_context)
+ right_context_utterance = torch.cat([right_context, utterance])
+ return right_context_utterance
+
+ def _apply_conv_module_infer(
+ self,
+ right_context_utterance: torch.Tensor,
+ R: int,
+ conv_cache: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Apply convolution module on utterance in inference mode."""
+ utterance = right_context_utterance[R:]
+ right_context = right_context_utterance[:R]
+ utterance, right_context, conv_cache = self.conv_module.infer(
+ utterance, right_context, conv_cache
+ )
+ right_context_utterance = torch.cat([right_context, utterance])
+ return right_context_utterance, conv_cache
+
+ def _apply_attention_module_forward(
+ self,
+ right_context_utterance: torch.Tensor,
+ R: int,
+ memory: torch.Tensor,
+ attention_mask: torch.Tensor,
+ padding_mask: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Apply attention module in training and validation mode."""
+ utterance = right_context_utterance[R:]
+ right_context = right_context_utterance[:R]
+
+ if self.use_memory:
+ summary = self.summary_op(utterance.permute(1, 2, 0)).permute(
+ 2, 0, 1
+ )
+ else:
+ summary = torch.empty(0).to(
+ dtype=utterance.dtype, device=utterance.device
+ )
+ output_right_context_utterance, output_memory = self.attention(
+ utterance=utterance,
+ right_context=right_context,
+ summary=summary,
+ memory=memory,
+ attention_mask=attention_mask,
+ padding_mask=padding_mask,
+ )
+
+ return output_right_context_utterance, output_memory
+
+ def _apply_attention_module_infer(
+ self,
+ right_context_utterance: torch.Tensor,
+ R: int,
+ memory: torch.Tensor,
+ attn_cache: List[torch.Tensor],
+ padding_mask: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor, List[torch.Tensor]]:
+ """Apply attention module in inference mode.
+ 1) Unpack cached states including:
+ - memory from previous chunks in the lower layer;
+ - attention key and value of left context from preceding
+ chunk's compuation;
+ 2) Apply attention computation;
+ 3) Update cached attention states including:
+ - output memory of current chunk in the lower layer;
+ - attention key and value in current chunk's computation, which would
+ be resued in next chunk's computation.
+ """
+ utterance = right_context_utterance[R:]
+ right_context = right_context_utterance[:R]
+
+ pre_memory = attn_cache[0]
+ left_context_key = attn_cache[1]
+ left_context_val = attn_cache[2]
+
+ if self.use_memory:
+ summary = self.summary_op(utterance.permute(1, 2, 0)).permute(
+ 2, 0, 1
+ )
+ summary = summary[:1]
+ else:
+ summary = torch.empty(0).to(
+ dtype=utterance.dtype, device=utterance.device
+ )
+ (
+ output_right_context_utterance,
+ output_memory,
+ next_key,
+ next_val,
+ ) = self.attention.infer(
+ utterance=utterance,
+ right_context=right_context,
+ summary=summary,
+ memory=pre_memory,
+ left_context_key=left_context_key,
+ left_context_val=left_context_val,
+ padding_mask=padding_mask,
+ )
+ attn_cache = self._update_attn_cache(
+ next_key, next_val, memory, attn_cache
+ )
+ return output_right_context_utterance, output_memory, attn_cache
+
+ def forward(
+ self,
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ memory: torch.Tensor,
+ attention_mask: torch.Tensor,
+ padding_mask: Optional[torch.Tensor] = None,
+ warmup: float = 1.0,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ r"""Forward pass for training and validation mode.
+
+ B: batch size;
+ D: embedding dimension;
+ R: length of hard-copied right contexts;
+ U: length of full utterance;
+ M: length of memory vectors.
+
+ Args:
+ utterance (torch.Tensor):
+ Utterance frames, with shape (U, B, D).
+ right_context (torch.Tensor):
+ Right context frames, with shape (R, B, D).
+ memory (torch.Tensor):
+ Memory elements, with shape (M, B, D).
+ It is an empty tensor without using memory.
+ attention_mask (torch.Tensor):
+ Attention mask for underlying attention module,
+ with shape (Q, KV), where Q = R + U + S, KV = M + R + U.
+ padding_mask (torch.Tensor):
+ Padding mask of ker tensor, with shape (B, KV).
+
+ Returns:
+ A tuple containing 3 tensors:
+ - output utterance, with shape (U, B, D).
+ - output right context, with shape (R, B, D).
+ - output memory, with shape (M, B, D).
+ """
+ R = right_context.size(0)
+ src = torch.cat([right_context, utterance])
+ src_orig = src
+
+ warmup_scale = min(0.1 + 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 0.1
+ )
+ else:
+ alpha = 1.0
+
+ # macaron style feed forward module
+ src = src + self.dropout(self.feed_forward_macaron(src))
+
+ # emformer attention module
+ src_att, output_memory = self._apply_attention_module_forward(
+ src, R, memory, attention_mask, padding_mask=padding_mask
+ )
+ src = src + self.dropout(src_att)
+
+ # convolution module
+ src_conv = self._apply_conv_module_forward(src, R)
+ src = src + self.dropout(src_conv)
+
+ # 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
+
+ output_utterance = src[R:]
+ output_right_context = src[:R]
+ return output_utterance, output_right_context, output_memory
+
+ @torch.jit.export
+ def infer(
+ self,
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ memory: torch.Tensor,
+ attn_cache: List[torch.Tensor],
+ conv_cache: torch.Tensor,
+ padding_mask: Optional[torch.Tensor] = None,
+ ) -> Tuple[
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ List[torch.Tensor],
+ torch.Tensor,
+ ]:
+ """Forward pass for inference.
+
+ B: batch size;
+ D: embedding dimension;
+ R: length of right_context;
+ U: length of utterance;
+ M: length of memory.
+
+ Args:
+ utterance (torch.Tensor):
+ Utterance frames, with shape (U, B, D).
+ right_context (torch.Tensor):
+ Right context frames, with shape (R, B, D).
+ memory (torch.Tensor):
+ Memory elements, with shape (M, B, D).
+ attn_cache (List[torch.Tensor]):
+ Cached attention tensors generated in preceding computation,
+ including memory, key and value of left context.
+ conv_cache (torch.Tensor, optional):
+ Cache tensor of left context for causal convolution.
+ padding_mask (torch.Tensor):
+ Padding mask of ker tensor.
+
+ Returns:
+ (Tensor, Tensor, List[torch.Tensor], Tensor):
+ - output utterance, with shape (U, B, D);
+ - output right_context, with shape (R, B, D);
+ - output memory, with shape (1, B, D) or (0, B, D).
+ - updated attention cache.
+ - updated convolution cache.
+ """
+ R = right_context.size(0)
+ src = torch.cat([right_context, utterance])
+
+ # macaron style feed forward module
+ src = src + self.dropout(self.feed_forward_macaron(src))
+
+ # emformer attention module
+ (
+ src_att,
+ output_memory,
+ attn_cache,
+ ) = self._apply_attention_module_infer(
+ src, R, memory, attn_cache, padding_mask=padding_mask
+ )
+ src = src + self.dropout(src_att)
+
+ # convolution module
+ src_conv, conv_cache = self._apply_conv_module_infer(src, R, conv_cache)
+ src = src + self.dropout(src_conv)
+
+ # feed forward module
+ src = src + self.dropout(self.feed_forward(src))
+
+ src = self.norm_final(self.balancer(src))
+
+ output_utterance = src[R:]
+ output_right_context = src[:R]
+ return (
+ output_utterance,
+ output_right_context,
+ output_memory,
+ attn_cache,
+ conv_cache,
+ )
+
+
+def _gen_attention_mask_block(
+ col_widths: List[int],
+ col_mask: List[bool],
+ num_rows: int,
+ device: torch.device,
+) -> torch.Tensor:
+ assert len(col_widths) == len(
+ col_mask
+ ), "Length of col_widths must match that of col_mask"
+
+ mask_block = [
+ torch.ones(num_rows, col_width, device=device)
+ if is_ones_col
+ else torch.zeros(num_rows, col_width, device=device)
+ for col_width, is_ones_col in zip(col_widths, col_mask)
+ ]
+ return torch.cat(mask_block, dim=1)
+
+
+class EmformerEncoder(nn.Module):
+ """Implements the Emformer architecture introduced in
+ *Emformer: Efficient Memory Transformer Based Acoustic Model for Low Latency
+ Streaming Speech Recognition*
+ [:footcite:`shi2021emformer`].
+
+ Args:
+ d_model (int):
+ Input dimension.
+ nhead (int):
+ Number of attention heads in each emformer layer.
+ dim_feedforward (int):
+ Hidden layer dimension of each emformer layer's feedforward network.
+ num_encoder_layers (int):
+ Number of emformer layers to instantiate.
+ chunk_length (int):
+ Length of each input segment.
+ dropout (float, optional):
+ Dropout probability. (default: 0.0)
+ layer_dropout (float, optional):
+ Layer dropout probability. (default: 0.0)
+ cnn_module_kernel (int):
+ Kernel size of convolution module.
+ left_context_length (int, optional):
+ Length of left context. (default: 0)
+ right_context_length (int, optional):
+ Length of right context. (default: 0)
+ memory_size (int, optional):
+ Number of memory elements to use. (default: 0)
+ tanh_on_mem (bool, optional):
+ If ``true``, applies tanh to memory elements. (default: ``false``)
+ negative_inf (float, optional):
+ Value to use for negative infinity in attention weights. (default: -1e8)
+ """
+
+ def __init__(
+ self,
+ chunk_length: int,
+ d_model: int = 256,
+ nhead: int = 4,
+ dim_feedforward: int = 2048,
+ num_encoder_layers: int = 12,
+ dropout: float = 0.1,
+ layer_dropout: float = 0.075,
+ cnn_module_kernel: int = 31,
+ left_context_length: int = 0,
+ right_context_length: int = 0,
+ memory_size: int = 0,
+ tanh_on_mem: bool = False,
+ negative_inf: float = -1e8,
+ ):
+ super().__init__()
+
+ assert (
+ chunk_length - 1
+ ) & chunk_length == 0, "chunk_length should be a power of 2."
+ self.shift = int(math.log(chunk_length, 2))
+
+ self.use_memory = memory_size > 0
+ self.init_memory_op = nn.AvgPool1d(
+ kernel_size=chunk_length,
+ stride=chunk_length,
+ ceil_mode=True,
+ )
+
+ self.emformer_layers = nn.ModuleList(
+ [
+ EmformerEncoderLayer(
+ d_model=d_model,
+ nhead=nhead,
+ dim_feedforward=dim_feedforward,
+ chunk_length=chunk_length,
+ dropout=dropout,
+ layer_dropout=layer_dropout,
+ cnn_module_kernel=cnn_module_kernel,
+ left_context_length=left_context_length,
+ right_context_length=right_context_length,
+ memory_size=memory_size,
+ tanh_on_mem=tanh_on_mem,
+ negative_inf=negative_inf,
+ )
+ for layer_idx in range(num_encoder_layers)
+ ]
+ )
+
+ self.num_encoder_layers = num_encoder_layers
+ self.d_model = d_model
+ self.left_context_length = left_context_length
+ self.right_context_length = right_context_length
+ self.chunk_length = chunk_length
+ self.memory_size = memory_size
+ self.cnn_module_kernel = cnn_module_kernel
+
+ def _gen_right_context(self, x: torch.Tensor) -> torch.Tensor:
+ """Hard copy each chunk's right context and concat them."""
+ T = x.shape[0]
+ num_chunks = math.ceil(
+ (T - self.right_context_length) / self.chunk_length
+ )
+ # first (num_chunks - 1) right context block
+ intervals = torch.arange(
+ 0, self.chunk_length * (num_chunks - 1), self.chunk_length
+ )
+ first = torch.arange(
+ self.chunk_length, self.chunk_length + self.right_context_length
+ )
+ indexes = intervals.unsqueeze(1) + first.unsqueeze(0)
+ # cat last right context block
+ indexes = torch.cat(
+ [
+ indexes,
+ torch.arange(T - self.right_context_length, T).unsqueeze(0),
+ ]
+ )
+ right_context_blocks = x[indexes.reshape(-1)]
+ return right_context_blocks
+
+ def _gen_attention_mask_col_widths(
+ self, chunk_idx: int, U: int
+ ) -> List[int]:
+ """Calculate column widths (key, value) in attention mask for the
+ chunk_idx chunk."""
+ num_chunks = math.ceil(U / self.chunk_length)
+ rc = self.right_context_length
+ lc = self.left_context_length
+ rc_start = chunk_idx * rc
+ rc_end = rc_start + rc
+ chunk_start = max(chunk_idx * self.chunk_length - lc, 0)
+ chunk_end = min((chunk_idx + 1) * self.chunk_length, U)
+ R = rc * num_chunks
+
+ if self.use_memory:
+ m_start = max(chunk_idx - self.memory_size, 0)
+ M = num_chunks - 1
+ col_widths = [
+ m_start, # before memory
+ chunk_idx - m_start, # memory
+ M - chunk_idx, # after memory
+ rc_start, # before right context
+ rc, # right context
+ R - rc_end, # after right context
+ chunk_start, # before chunk
+ chunk_end - chunk_start, # chunk
+ U - chunk_end, # after chunk
+ ]
+ else:
+ col_widths = [
+ rc_start, # before right context
+ rc, # right context
+ R - rc_end, # after right context
+ chunk_start, # before chunk
+ chunk_end - chunk_start, # chunk
+ U - chunk_end, # after chunk
+ ]
+
+ return col_widths
+
+ def _gen_attention_mask(self, utterance: torch.Tensor) -> torch.Tensor:
+ """Generate attention mask to simulate underlying chunk-wise attention
+ computation, where chunk-wise connections are filled with `False`,
+ and other unnecessary connections beyond chunk are filled with `True`.
+
+ R: length of hard-copied right contexts;
+ U: length of full utterance;
+ S: length of summary vectors;
+ M: length of memory vectors;
+ Q: length of attention query;
+ KV: length of attention key and value.
+
+ The shape of attention mask is (Q, KV).
+ If self.use_memory is `True`:
+ query = [right_context, utterance, summary];
+ key, value = [memory, right_context, utterance];
+ Q = R + U + S, KV = M + R + U.
+ Otherwise:
+ query = [right_context, utterance]
+ key, value = [right_context, utterance]
+ Q = R + U, KV = R + U.
+
+ Suppose:
+ c_i: chunk at index i;
+ r_i: right context that c_i can use;
+ l_i: left context that c_i can use;
+ m_i: past memory vectors from previous layer that c_i can use;
+ s_i: summary vector of c_i.
+ The target chunk-wise attention is:
+ c_i, r_i (in query) -> l_i, c_i, r_i, m_i (in key);
+ s_i (in query) -> l_i, c_i, r_i (in key).
+ """
+ U = utterance.size(0)
+ num_chunks = math.ceil(U / self.chunk_length)
+
+ right_context_mask = []
+ utterance_mask = []
+ summary_mask = []
+
+ if self.use_memory:
+ num_cols = 9
+ # right context and utterance both attend to memory, right context,
+ # utterance
+ right_context_utterance_cols_mask = [
+ idx in [1, 4, 7] for idx in range(num_cols)
+ ]
+ # summary attends to right context, utterance
+ summary_cols_mask = [idx in [4, 7] for idx in range(num_cols)]
+ masks_to_concat = [right_context_mask, utterance_mask, summary_mask]
+ else:
+ num_cols = 6
+ # right context and utterance both attend to right context and
+ # utterance
+ right_context_utterance_cols_mask = [
+ idx in [1, 4] for idx in range(num_cols)
+ ]
+ summary_cols_mask = None
+ masks_to_concat = [right_context_mask, utterance_mask]
+
+ for chunk_idx in range(num_chunks):
+ col_widths = self._gen_attention_mask_col_widths(chunk_idx, U)
+
+ right_context_mask_block = _gen_attention_mask_block(
+ col_widths,
+ right_context_utterance_cols_mask,
+ self.right_context_length,
+ utterance.device,
+ )
+ right_context_mask.append(right_context_mask_block)
+
+ utterance_mask_block = _gen_attention_mask_block(
+ col_widths,
+ right_context_utterance_cols_mask,
+ min(
+ self.chunk_length,
+ U - chunk_idx * self.chunk_length,
+ ),
+ utterance.device,
+ )
+ utterance_mask.append(utterance_mask_block)
+
+ if summary_cols_mask is not None:
+ summary_mask_block = _gen_attention_mask_block(
+ col_widths, summary_cols_mask, 1, utterance.device
+ )
+ summary_mask.append(summary_mask_block)
+
+ attention_mask = (
+ 1 - torch.cat([torch.cat(mask) for mask in masks_to_concat])
+ ).to(torch.bool)
+ return attention_mask
+
+ def forward(
+ self, x: torch.Tensor, lengths: torch.Tensor, warmup: float = 1.0
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Forward pass for training and validation mode.
+
+ B: batch size;
+ D: input dimension;
+ U: length of utterance.
+
+ Args:
+ x (torch.Tensor):
+ Utterance frames right-padded with right context frames,
+ with shape (U + right_context_length, B, D).
+ lengths (torch.Tensor):
+ With shape (B,) and i-th element representing number of valid
+ utterance frames for i-th batch element in x, which contains the
+ right_context at the end.
+
+ Returns:
+ A tuple of 2 tensors:
+ - output utterance frames, with shape (U, B, D).
+ - output_lengths, with shape (B,), without containing the
+ right_context at the end.
+ """
+ U = x.size(0) - self.right_context_length
+
+ right_context = self._gen_right_context(x)
+ utterance = x[:U]
+ output_lengths = torch.clamp(lengths - self.right_context_length, min=0)
+ attention_mask = self._gen_attention_mask(utterance)
+ memory = (
+ self.init_memory_op(utterance.permute(1, 2, 0)).permute(2, 0, 1)[
+ :-1
+ ]
+ if self.use_memory
+ else torch.empty(0).to(dtype=x.dtype, device=x.device)
+ )
+ padding_mask = make_pad_mask(
+ memory.size(0) + right_context.size(0) + output_lengths
+ )
+
+ output = utterance
+ for layer in self.emformer_layers:
+ output, right_context, memory = layer(
+ output,
+ right_context,
+ memory,
+ attention_mask,
+ padding_mask=padding_mask,
+ warmup=warmup,
+ )
+
+ return output, output_lengths
+
+ @torch.jit.export
+ def infer(
+ self,
+ x: torch.Tensor,
+ lengths: torch.Tensor,
+ num_processed_frames: torch.Tensor,
+ states: Tuple[List[List[torch.Tensor]], List[torch.Tensor]],
+ ) -> Tuple[
+ torch.Tensor,
+ torch.Tensor,
+ Tuple[List[List[torch.Tensor]], List[torch.Tensor]],
+ ]:
+ """Forward pass for streaming inference.
+
+ B: batch size;
+ D: input dimension;
+ U: length of utterance.
+
+ Args:
+ x (torch.Tensor):
+ Utterance frames right-padded with right context frames,
+ with shape (U + right_context_length, B, D).
+ lengths (torch.Tensor):
+ With shape (B,) and i-th element representing number of valid
+ utterance frames for i-th batch element in x, which contains the
+ right_context at the end.
+ states (List[torch.Tensor, List[List[torch.Tensor]], List[torch.Tensor]]: # noqa
+ Cached states containing:
+ - attn_caches: attention states from preceding chunk's computation,
+ where each element corresponds to each emformer layer
+ - conv_caches: left context for causal convolution, where each
+ element corresponds to each layer.
+
+ Returns:
+ (Tensor, Tensor, List[List[torch.Tensor]], List[torch.Tensor]):
+ - output utterance frames, with shape (U, B, D).
+ - output lengths, with shape (B,), without containing the
+ right_context at the end.
+ - updated states from current chunk's computation.
+ """
+ assert num_processed_frames.shape == (x.size(1),)
+
+ attn_caches = states[0]
+ assert len(attn_caches) == self.num_encoder_layers, len(attn_caches)
+ for i in range(len(attn_caches)):
+ assert attn_caches[i][0].shape == (
+ self.memory_size,
+ x.size(1),
+ self.d_model,
+ ), attn_caches[i][0].shape
+ assert attn_caches[i][1].shape == (
+ self.left_context_length,
+ x.size(1),
+ self.d_model,
+ ), attn_caches[i][1].shape
+ assert attn_caches[i][2].shape == (
+ self.left_context_length,
+ x.size(1),
+ self.d_model,
+ ), attn_caches[i][2].shape
+
+ conv_caches = states[1]
+ assert len(conv_caches) == self.num_encoder_layers, len(conv_caches)
+ for i in range(len(conv_caches)):
+ assert conv_caches[i].shape == (
+ x.size(1),
+ self.d_model,
+ self.cnn_module_kernel - 1,
+ ), conv_caches[i].shape
+
+ right_context = x[-self.right_context_length :]
+ utterance = x[: -self.right_context_length]
+ output_lengths = torch.clamp(lengths - self.right_context_length, min=0)
+ memory = (
+ self.init_memory_op(utterance.permute(1, 2, 0)).permute(2, 0, 1)
+ if self.use_memory
+ else torch.empty(0).to(dtype=x.dtype, device=x.device)
+ )
+
+ # calcualte padding mask to mask out initial zero caches
+ chunk_mask = make_pad_mask(output_lengths).to(x.device)
+ memory_mask = (
+ (
+ (num_processed_frames >> self.shift).view(x.size(1), 1)
+ <= torch.arange(self.memory_size, device=x.device).expand(
+ x.size(1), self.memory_size
+ )
+ ).flip(1)
+ if self.use_memory
+ else torch.empty(0).to(dtype=torch.bool, device=x.device)
+ )
+ left_context_mask = (
+ num_processed_frames.view(x.size(1), 1)
+ <= torch.arange(self.left_context_length, device=x.device).expand(
+ x.size(1), self.left_context_length
+ )
+ ).flip(1)
+ right_context_mask = torch.zeros(
+ x.size(1),
+ self.right_context_length,
+ dtype=torch.bool,
+ device=x.device,
+ )
+ padding_mask = torch.cat(
+ [memory_mask, right_context_mask, left_context_mask, chunk_mask],
+ dim=1,
+ )
+
+ output = utterance
+ output_attn_caches: List[List[torch.Tensor]] = []
+ output_conv_caches: List[torch.Tensor] = []
+ for layer_idx, layer in enumerate(self.emformer_layers):
+ (
+ output,
+ right_context,
+ memory,
+ output_attn_cache,
+ output_conv_cache,
+ ) = layer.infer(
+ output,
+ right_context,
+ memory,
+ padding_mask=padding_mask,
+ attn_cache=attn_caches[layer_idx],
+ conv_cache=conv_caches[layer_idx],
+ )
+ output_attn_caches.append(output_attn_cache)
+ output_conv_caches.append(output_conv_cache)
+
+ output_states: Tuple[List[List[torch.Tensor]], List[torch.Tensor]] = (
+ output_attn_caches,
+ output_conv_caches,
+ )
+ return output, output_lengths, output_states
+
+ @torch.jit.export
+ def init_states(self, device: torch.device = torch.device("cpu")):
+ """Create initial 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__(
+ self,
+ num_features: int,
+ chunk_length: int,
+ subsampling_factor: int = 4,
+ d_model: int = 256,
+ nhead: int = 4,
+ dim_feedforward: int = 2048,
+ num_encoder_layers: int = 12,
+ dropout: float = 0.1,
+ layer_dropout: float = 0.075,
+ cnn_module_kernel: int = 3,
+ left_context_length: int = 0,
+ right_context_length: int = 0,
+ memory_size: int = 0,
+ tanh_on_mem: bool = False,
+ negative_inf: float = -1e8,
+ ):
+ super().__init__()
+
+ self.subsampling_factor = subsampling_factor
+ self.right_context_length = right_context_length
+ self.chunk_length = chunk_length
+ if subsampling_factor != 4:
+ raise NotImplementedError("Support only 'subsampling_factor=4'.")
+ if chunk_length % subsampling_factor != 0:
+ raise NotImplementedError(
+ "chunk_length must be a mutiple of subsampling_factor."
+ )
+ if (
+ left_context_length != 0
+ and left_context_length % subsampling_factor != 0
+ ):
+ raise NotImplementedError(
+ "left_context_length must be 0 or a mutiple of subsampling_factor." # noqa
+ )
+ if (
+ right_context_length != 0
+ and right_context_length % subsampling_factor != 0
+ ):
+ raise NotImplementedError(
+ "right_context_length must be 0 or a mutiple of subsampling_factor." # noqa
+ )
+
+ # 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 = EmformerEncoder(
+ chunk_length=chunk_length // subsampling_factor,
+ d_model=d_model,
+ nhead=nhead,
+ dim_feedforward=dim_feedforward,
+ num_encoder_layers=num_encoder_layers,
+ dropout=dropout,
+ layer_dropout=layer_dropout,
+ cnn_module_kernel=cnn_module_kernel,
+ left_context_length=left_context_length // subsampling_factor,
+ right_context_length=right_context_length // subsampling_factor,
+ memory_size=memory_size,
+ tanh_on_mem=tanh_on_mem,
+ negative_inf=negative_inf,
+ )
+
+ def forward(
+ self, x: torch.Tensor, x_lens: torch.Tensor, warmup: float = 1.0
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Forward pass for training and non-streaming inference.
+
+ B: batch size;
+ D: feature dimension;
+ T: length of utterance.
+
+ Args:
+ x (torch.Tensor):
+ Utterance frames right-padded with right context frames,
+ with shape (B, T, D).
+ x_lens (torch.Tensor):
+ With shape (B,) and i-th element representing number of valid
+ utterance frames for i-th batch element in x, containing the
+ right_context at the end.
+ 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:
+ (Tensor, Tensor):
+ - output embedding, with shape (B, T', D), where
+ T' = ((T - 1) // 2 - 1) // 2 - self.right_context_length // 4.
+ - output lengths, with shape (B,), without containing the
+ right_context at the end.
+ """
+ x = self.encoder_embed(x)
+ x = x.permute(1, 0, 2) # (N, T, C) -> (T, N, C)
+
+ x_lens = (((x_lens - 1) >> 1) - 1) >> 1
+ assert x.size(0) == x_lens.max().item()
+
+ output, output_lengths = self.encoder(
+ x, x_lens, warmup=warmup
+ ) # (T, N, C)
+
+ output = output.permute(1, 0, 2) # (T, N, C) -> (N, T, C)
+
+ return output, output_lengths
+
+ @torch.jit.export
+ def infer(
+ self,
+ x: torch.Tensor,
+ x_lens: torch.Tensor,
+ num_processed_frames: torch.Tensor,
+ states: Tuple[List[List[torch.Tensor]], List[torch.Tensor]],
+ ) -> Tuple[
+ torch.Tensor,
+ torch.Tensor,
+ Tuple[List[List[torch.Tensor]], List[torch.Tensor]],
+ ]:
+ """Forward pass for streaming inference.
+
+ B: batch size;
+ D: feature dimension;
+ T: length of utterance.
+
+ Args:
+ x (torch.Tensor):
+ Utterance frames right-padded with right context frames,
+ with shape (B, T, D).
+ lengths (torch.Tensor):
+ With shape (B,) and i-th element representing number of valid
+ utterance frames for i-th batch element in x, containing the
+ right_context at the end.
+ states (List[torch.Tensor, List[List[torch.Tensor]], List[torch.Tensor]]: # noqa
+ Cached states containing:
+ - past_lens: number of past frames for each sample in batch
+ - attn_caches: attention states from preceding chunk's computation,
+ where each element corresponds to each emformer layer
+ - conv_caches: left context for causal convolution, where each
+ element corresponds to each layer.
+ Returns:
+ (Tensor, Tensor):
+ - output embedding, with shape (B, T', D), where
+ T' = ((T - 1) // 2 - 1) // 2 - self.right_context_length // 4.
+ - output lengths, with shape (B,), without containing the
+ right_context at the end.
+ - updated states from current chunk's computation.
+ """
+ x = self.encoder_embed(x)
+ # drop the first and last frames
+ x = x[:, 1:-1, :]
+ x = x.permute(1, 0, 2) # (N, T, C) -> (T, N, C)
+
+ # Caution: We assume the subsampling factor is 4!
+ x_lens = (((x_lens - 1) >> 1) - 1) >> 1
+ x_lens -= 2
+ assert x.size(0) == x_lens.max().item()
+
+ num_processed_frames = num_processed_frames >> 2
+
+ output, output_lengths, output_states = self.encoder.infer(
+ x, x_lens, num_processed_frames, states
+ )
+
+ output = output.permute(1, 0, 2) # (T, N, C) -> (N, T, C)
+
+ return output, output_lengths, output_states
+
+ @torch.jit.export
+ def init_states(self, device: torch.device = torch.device("cpu")):
+ """Create initial states."""
+ return self.encoder.init_states(device)
+
+
+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-1)//2 - 1)//2, which approximates T' == T//4
+
+ 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,
+ ) -> 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-1)//2 - 1)//2, out_channels)
+ layer1_channels:
+ Number of channels in layer1
+ layer1_channels:
+ Number of channels in layer2
+ """
+ assert in_channels >= 7
+ super().__init__()
+
+ self.conv = nn.Sequential(
+ ScaledConv2d(
+ in_channels=1,
+ out_channels=layer1_channels,
+ kernel_size=3,
+ padding=1,
+ ),
+ ActivationBalancer(channel_dim=1),
+ DoubleSwish(),
+ ScaledConv2d(
+ in_channels=layer1_channels,
+ out_channels=layer2_channels,
+ kernel_size=3,
+ stride=2,
+ ),
+ ActivationBalancer(channel_dim=1),
+ DoubleSwish(),
+ ScaledConv2d(
+ in_channels=layer2_channels,
+ out_channels=layer3_channels,
+ kernel_size=3,
+ stride=2,
+ ),
+ ActivationBalancer(channel_dim=1),
+ DoubleSwish(),
+ )
+ self.out = ScaledLinear(
+ layer3_channels * (((in_channels - 1) // 2 - 1) // 2), out_channels
+ )
+ # set learn_eps=False because out_norm is preceded by `out`, and `out`
+ # itself has learned scale, so the extra degree of freedom is not
+ # needed.
+ self.out_norm = BasicNorm(out_channels, learn_eps=False)
+ # constrain median of output to be close to zero.
+ self.out_balancer = ActivationBalancer(
+ channel_dim=-1, min_positive=0.45, max_positive=0.55
+ )
+
+ 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-1)//2 - 1)//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-1)//2 - 1)//2, ((idim-1)//2 - 1)//2)
+ b, c, t, f = x.size()
+ x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f))
+ # Now x is of shape (N, ((T-1)//2 - 1))//2, odim)
+ x = self.out_norm(x)
+ x = self.out_balancer(x)
+ return x
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless/encoder_interface.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless/encoder_interface.py
new file mode 120000
index 000000000..b9aa0ae08
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless/encoder_interface.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/encoder_interface.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless/export.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless/export.py
new file mode 100755
index 000000000..4930881ea
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless/export.py
@@ -0,0 +1,287 @@
+#!/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:
+./conv_emformer_transducer_stateless/export.py \
+ --exp-dir ./conv_emformer_transducer_stateless/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 10 \
+ --use-averaged-model=True \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --jit False
+
+It will generate a file exp_dir/pretrained.pt
+
+To use the generated file with `conv_emformer_transducer_stateless/decode.py`,
+you can do:
+
+ cd /path/to/exp_dir
+ ln -s pretrained.pt epoch-9999.pt
+
+ cd /path/to/egs/librispeech/ASR
+ ./conv_emformer_transducer_stateless/decode.py \
+ --exp-dir ./conv_emformer_transducer_stateless/exp \
+ --epoch 9999 \
+ --avg 1 \
+ --max-duration 100 \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --use-averaged-model=False \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32
+"""
+
+import argparse
+import logging
+from pathlib import Path
+
+import sentencepiece as spm
+import torch
+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=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(
+ "--jit",
+ type=str2bool,
+ default=False,
+ help="""True to save a model after applying torch.jit.script.
+ """,
+ )
+
+ 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 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.eval()
+
+ if params.jit:
+ # 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 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/librispeech/ASR/conv_emformer_transducer_stateless/joiner.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless/joiner.py
new file mode 120000
index 000000000..815fd4bb6
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless/joiner.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/joiner.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless/model.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless/model.py
new file mode 120000
index 000000000..ebb6d774d
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless/model.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/model.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless/optim.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless/optim.py
new file mode 120000
index 000000000..e2deb4492
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless/optim.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/optim.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless/scaling.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless/scaling.py
new file mode 120000
index 000000000..09d802cc4
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless/scaling.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/scaling.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless/stream.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless/stream.py
new file mode 100644
index 000000000..9494e1fc1
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless/stream.py
@@ -0,0 +1,156 @@
+# 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.
+
+import math
+from typing import List, Optional, Tuple
+
+import k2
+import torch
+from beam_search import Hypothesis, HypothesisList
+
+from icefall.utils import AttributeDict
+
+
+class Stream(object):
+ def __init__(
+ self,
+ params: AttributeDict,
+ cut_id: str,
+ decoding_graph: Optional[k2.Fsa] = None,
+ device: torch.device = torch.device("cpu"),
+ LOG_EPS: float = math.log(1e-10),
+ ) -> None:
+ """
+ Args:
+ params:
+ It's the return value of :func:`get_params`.
+ decoding_graph:
+ The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+ only when --decoding_method is fast_beam_search.
+ device:
+ The device to run this stream.
+ """
+ self.LOG_EPS = LOG_EPS
+ self.cut_id = cut_id
+
+ # Containing attention caches and convolution caches
+ self.states: Optional[
+ Tuple[List[List[torch.Tensor]], List[torch.Tensor]]
+ ] = None
+
+ # It uses different attributes for different decoding methods.
+ self.context_size = params.context_size
+ self.decoding_method = params.decoding_method
+ 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":
+ # feature_len is needed to get partial results.
+ # The rnnt_decoding_stream for fast_beam_search.
+ self.rnnt_decoding_stream: k2.RnntDecodingStream = (
+ k2.RnntDecodingStream(decoding_graph)
+ )
+ self.hyp: Optional[List[int]] = None
+ else:
+ raise ValueError(
+ f"Unsupported decoding method: {params.decoding_method}"
+ )
+
+ self.ground_truth: str = ""
+
+ self.feature: Optional[torch.Tensor] = None
+ # Make sure all feature frames can be used.
+ # Add 2 here since we will drop the first and last after subsampling.
+ self.chunk_length = params.chunk_length
+ self.pad_length = (
+ params.right_context_length + 2 * params.subsampling_factor + 3
+ )
+ self.num_frames = 0
+ self.num_processed_frames = 0
+
+ # After all feature frames are processed, we set this flag to True
+ self._done = False
+
+ def set_feature(self, feature: torch.Tensor) -> None:
+ assert feature.dim() == 2, feature.dim()
+ self.num_frames = feature.size(0)
+ # tail padding
+ self.feature = torch.nn.functional.pad(
+ feature,
+ (0, 0, 0, self.pad_length),
+ mode="constant",
+ value=self.LOG_EPS,
+ )
+
+ def set_ground_truth(self, ground_truth: str) -> None:
+ self.ground_truth = ground_truth
+
+ def set_states(
+ self, states: Tuple[List[List[torch.Tensor]], List[torch.Tensor]]
+ ) -> None:
+ """Set states."""
+ self.states = states
+
+ def get_feature_chunk(self) -> torch.Tensor:
+ """Get a chunk of feature frames.
+
+ Returns:
+ A tensor of shape (ret_length, feature_dim).
+ """
+ update_length = min(
+ self.num_frames - self.num_processed_frames, self.chunk_length
+ )
+ ret_length = update_length + self.pad_length
+
+ ret_feature = self.feature[
+ self.num_processed_frames : self.num_processed_frames + ret_length
+ ]
+ # Cut off used frames.
+ # self.feature = self.feature[update_length:]
+
+ self.num_processed_frames += update_length
+ if self.num_processed_frames >= self.num_frames:
+ self._done = True
+
+ return ret_feature
+
+ @property
+ def done(self) -> bool:
+ """Return True if all feature frames are processed."""
+ return self._done
+
+ @property
+ def id(self) -> str:
+ return self.cut_id
+
+ def decoding_result(self) -> List[int]:
+ """Obtain current decoding result."""
+ if self.decoding_method == "greedy_search":
+ return self.hyp[self.context_size :]
+ elif self.decoding_method == "modified_beam_search":
+ best_hyp = self.hyps.get_most_probable(length_norm=True)
+ return best_hyp.ys[self.context_size :]
+ else:
+ assert self.decoding_method == "fast_beam_search"
+ return self.hyp
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless/streaming_decode.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless/streaming_decode.py
new file mode 100755
index 000000000..61dbe8658
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless/streaming_decode.py
@@ -0,0 +1,983 @@
+#!/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
+./conv_emformer_transducer_stateless/streaming_decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless/exp \
+ --num-decode-streams 2000 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method greedy_search \
+ --use-averaged-model True
+
+(2) modified beam search
+./conv_emformer_transducer_stateless/streaming_decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless/exp \
+ --num-decode-streams 2000 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method modified_beam_search \
+ --use-averaged-model True \
+ --beam-size 4
+
+(3) fast beam search
+./conv_emformer_transducer_stateless/streaming_decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless/exp \
+ --num-decode-streams 2000 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method fast_beam_search \
+ --use-averaged-model True \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8
+"""
+import argparse
+import logging
+import warnings
+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 beam_search import Hypothesis, HypothesisList, get_hyps_shape
+from emformer import LOG_EPSILON, stack_states, unstack_states
+from kaldifeat import Fbank, FbankOptions
+from lhotse import CutSet
+from stream import Stream
+from torch.nn.utils.rnn import pad_sequence
+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 one_best_decoding
+from icefall.utils import (
+ AttributeDict,
+ get_texts,
+ 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=28,
+ help="It specifies the checkpoint to use for decoding."
+ "Note: Epoch counts from 0.",
+ )
+
+ 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'. ",
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=False,
+ 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="transducer_emformer/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="""Possible values are:
+ - greedy_search
+ - modified_beam_search
+ - fast_beam_search
+ """,
+ )
+
+ 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(
+ "--max-contexts",
+ type=int,
+ default=4,
+ help="""Used only when --decoding-method is
+ fast_beam_search""",
+ )
+
+ parser.add_argument(
+ "--max-states",
+ type=int,
+ default=8,
+ 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(
+ "--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(
+ "--sampling-rate",
+ type=float,
+ default=16000,
+ help="Sample rate of the audio",
+ )
+
+ 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 greedy_search(
+ model: nn.Module,
+ encoder_out: torch.Tensor,
+ streams: List[Stream],
+) -> None:
+ """Greedy search in batch mode. It hardcodes --max-sym-per-frame=1.
+
+ Args:
+ model:
+ The transducer model.
+ encoder_out:
+ Output from the encoder. Its shape is (N, T, C), where N >= 1.
+ streams:
+ A list of Stream objects.
+ """
+ assert len(streams) == encoder_out.size(0)
+ assert encoder_out.ndim == 3
+
+ blank_id = model.decoder.blank_id
+ context_size = model.decoder.context_size
+ device = next(model.parameters()).device
+ T = encoder_out.size(1)
+
+ encoder_out = model.joiner.encoder_proj(encoder_out)
+
+ decoder_input = torch.tensor(
+ [stream.hyp[-context_size:] for stream in streams],
+ device=device,
+ dtype=torch.int64,
+ )
+ # decoder_out is of shape (batch_size, 1, decoder_out_dim)
+ decoder_out = model.decoder(decoder_input, need_pad=False)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+
+ for t in range(T):
+ # current_encoder_out's shape: (batch_size, 1, encoder_out_dim)
+ current_encoder_out = encoder_out[:, t : t + 1, :] # noqa
+
+ logits = model.joiner(
+ current_encoder_out.unsqueeze(2),
+ decoder_out.unsqueeze(1),
+ project_input=False,
+ )
+ # logits'shape (batch_size, vocab_size)
+ logits = logits.squeeze(1).squeeze(1)
+
+ assert logits.ndim == 2, logits.shape
+ y = logits.argmax(dim=1).tolist()
+ emitted = False
+ for i, v in enumerate(y):
+ if v != blank_id:
+ streams[i].hyp.append(v)
+ emitted = True
+ if emitted:
+ # update decoder output
+ decoder_input = torch.tensor(
+ [stream.hyp[-context_size:] for stream in streams],
+ device=device,
+ dtype=torch.int64,
+ )
+ decoder_out = model.decoder(
+ decoder_input,
+ need_pad=False,
+ )
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+
+
+def modified_beam_search(
+ model: nn.Module,
+ encoder_out: torch.Tensor,
+ streams: List[Stream],
+ beam: int = 4,
+):
+ """Beam search in batch mode with --max-sym-per-frame=1 being hardcoded.
+
+ Args:
+ model:
+ The RNN-T model.
+ encoder_out:
+ A 3-D tensor of shape (N, T, encoder_out_dim) containing the output of
+ the encoder model.
+ streams:
+ A list of stream objects.
+ beam:
+ Number of active paths during the beam search.
+ """
+ assert encoder_out.ndim == 3, encoder_out.shape
+ assert len(streams) == encoder_out.size(0)
+
+ blank_id = model.decoder.blank_id
+ context_size = model.decoder.context_size
+ device = next(model.parameters()).device
+ batch_size = len(streams)
+ T = encoder_out.size(1)
+
+ B = [stream.hyps for stream in streams]
+
+ encoder_out = model.joiner.encoder_proj(encoder_out)
+
+ for t in range(T):
+ current_encoder_out = encoder_out[:, t].unsqueeze(1).unsqueeze(1)
+ # current_encoder_out's shape: (batch_size, 1, 1, encoder_out_dim)
+
+ 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.stack(
+ [hyp.log_prob.reshape(1) for hyps in A for hyp in hyps], dim=0
+ ) # (num_hyps, 1)
+
+ 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)
+ # decoder_out is of shape (num_hyps, 1, 1, decoder_output_dim)
+
+ # Note: For torch 1.7.1 and below, it requires a torch.int64 tensor
+ # as index, so we use `to(torch.int64)` below.
+ current_encoder_out = torch.index_select(
+ current_encoder_out,
+ dim=0,
+ index=hyps_shape.row_ids(1).to(torch.int64),
+ ) # (num_hyps, encoder_out_dim)
+
+ logits = model.joiner(
+ current_encoder_out, decoder_out, project_input=False
+ )
+ # logits is of shape (num_hyps, 1, 1, vocab_size)
+
+ logits = logits.squeeze(1).squeeze(1)
+
+ 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 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_ys = hyp.ys[:]
+ new_token = topk_token_indexes[k]
+ if new_token != blank_id:
+ new_ys.append(new_token)
+
+ new_log_prob = topk_log_probs[k]
+ new_hyp = Hypothesis(ys=new_ys, log_prob=new_log_prob)
+ B[i].add(new_hyp)
+
+ for i in range(batch_size):
+ streams[i].hyps = B[i]
+
+
+def fast_beam_search_one_best(
+ model: nn.Module,
+ streams: List[Stream],
+ encoder_out: torch.Tensor,
+ processed_lens: torch.Tensor,
+ beam: float,
+ max_states: int,
+ max_contexts: int,
+) -> None:
+ """It limits the maximum number of symbols per frame to 1.
+
+ A lattice is first obtained using modified beam search, and then
+ the shortest path within the lattice is used as the final output.
+
+ Args:
+ model:
+ An instance of `Transducer`.
+ streams:
+ A list of stream objects.
+ encoder_out:
+ A tensor of shape (N, T, C) from the encoder.
+ processed_lens:
+ A tensor of shape (N,) containing the number of processed frames
+ in `encoder_out` before padding.
+ beam:
+ Beam value, similar to the beam used in Kaldi..
+ max_states:
+ Max states per stream per frame.
+ max_contexts:
+ Max contexts pre stream per frame.
+ """
+ assert encoder_out.ndim == 3
+
+ context_size = model.decoder.context_size
+ vocab_size = model.decoder.vocab_size
+
+ B, T, C = encoder_out.shape
+ assert B == len(streams)
+
+ config = k2.RnntDecodingConfig(
+ vocab_size=vocab_size,
+ decoder_history_len=context_size,
+ beam=beam,
+ max_contexts=max_contexts,
+ max_states=max_states,
+ )
+ individual_streams = []
+ for i in range(B):
+ individual_streams.append(streams[i].rnnt_decoding_stream)
+ decoding_streams = k2.RnntDecodingStreams(individual_streams, config)
+
+ encoder_out = model.joiner.encoder_proj(encoder_out)
+
+ for t in range(T):
+ # shape is a RaggedShape of shape (B, context)
+ # contexts is a Tensor of shape (shape.NumElements(), context_size)
+ shape, contexts = decoding_streams.get_contexts()
+ # `nn.Embedding()` in torch below v1.7.1 supports only torch.int64
+ contexts = contexts.to(torch.int64)
+ # decoder_out is of shape (shape.NumElements(), 1, decoder_out_dim)
+ decoder_out = model.decoder(contexts, need_pad=False)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+ # current_encoder_out is of shape
+ # (shape.NumElements(), 1, joiner_dim)
+ # fmt: off
+ current_encoder_out = torch.index_select(
+ encoder_out[:, t:t + 1, :], 0, shape.row_ids(1).to(torch.int64)
+ )
+ # fmt: on
+ logits = model.joiner(
+ current_encoder_out.unsqueeze(2),
+ decoder_out.unsqueeze(1),
+ project_input=False,
+ )
+ logits = logits.squeeze(1).squeeze(1)
+ log_probs = logits.log_softmax(dim=-1)
+ decoding_streams.advance(log_probs)
+
+ decoding_streams.terminate_and_flush_to_streams()
+
+ lattice = decoding_streams.format_output(processed_lens.tolist())
+
+ best_path = one_best_decoding(lattice)
+ hyps = get_texts(best_path)
+
+ for i in range(B):
+ streams[i].hyp = hyps[i]
+
+
+def decode_one_chunk(
+ model: nn.Module,
+ streams: List[Stream],
+ params: AttributeDict,
+ decoding_graph: Optional[k2.Fsa] = None,
+) -> List[int]:
+ """
+ Args:
+ model:
+ The Transducer model.
+ streams:
+ A list of Stream objects.
+ params:
+ It is returned by :func:`get_params`.
+ decoding_graph:
+ The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+ only when --decoding_method is fast_beam_search.
+
+ Returns:
+ A list of indexes indicating the finished streams.
+ """
+ device = next(model.parameters()).device
+
+ feature_list = []
+ feature_len_list = []
+ state_list = []
+ num_processed_frames_list = []
+
+ for stream in streams:
+ # We should first get `stream.num_processed_frames`
+ # before calling `stream.get_feature_chunk()`
+ # since `stream.num_processed_frames` would be updated
+ num_processed_frames_list.append(stream.num_processed_frames)
+ feature = stream.get_feature_chunk()
+ feature_len = feature.size(0)
+ feature_list.append(feature)
+ feature_len_list.append(feature_len)
+ state_list.append(stream.states)
+
+ features = pad_sequence(
+ feature_list, batch_first=True, padding_value=LOG_EPSILON
+ ).to(device)
+ feature_lens = torch.tensor(feature_len_list, device=device)
+ num_processed_frames = torch.tensor(
+ num_processed_frames_list, device=device
+ )
+
+ # Make sure it has at least 1 frame after subsampling, first-and-last-frame cutting, and right context cutting # noqa
+ tail_length = (
+ 3 * params.subsampling_factor + params.right_context_length + 3
+ )
+ 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_EPSILON,
+ )
+
+ # Stack states of all streams
+ states = stack_states(state_list)
+
+ encoder_out, encoder_out_lens, states = model.encoder.infer(
+ x=features,
+ x_lens=feature_lens,
+ states=states,
+ num_processed_frames=num_processed_frames,
+ )
+
+ if params.decoding_method == "greedy_search":
+ greedy_search(
+ model=model,
+ streams=streams,
+ encoder_out=encoder_out,
+ )
+ elif params.decoding_method == "modified_beam_search":
+ modified_beam_search(
+ model=model,
+ streams=streams,
+ encoder_out=encoder_out,
+ beam=params.beam_size,
+ )
+ elif params.decoding_method == "fast_beam_search":
+ # feature_len is needed to get partial results.
+ # The rnnt_decoding_stream for fast_beam_search.
+ fast_beam_search_one_best(
+ model=model,
+ streams=streams,
+ encoder_out=encoder_out,
+ processed_lens=(num_processed_frames >> 2) + encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ )
+ else:
+ raise ValueError(
+ f"Unsupported decoding method: {params.decoding_method}"
+ )
+
+ # Update cached states of each stream
+ state_list = unstack_states(states)
+ for i, s in enumerate(state_list):
+ streams[i].states = s
+
+ finished_streams = [i for i, stream in enumerate(streams) if stream.done]
+ return finished_streams
+
+
+def create_streaming_feature_extractor() -> Fbank:
+ """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 = 16000
+ opts.mel_opts.num_bins = 80
+ return Fbank(opts)
+
+
+def decode_dataset(
+ cuts: CutSet,
+ model: nn.Module,
+ params: AttributeDict,
+ sp: spm.SentencePieceProcessor,
+ decoding_graph: Optional[k2.Fsa] = None,
+):
+ """Decode dataset.
+
+ Args:
+ cuts:
+ Lhotse Cutset containing the dataset to decode.
+ params:
+ It is returned by :func:`get_params`.
+ model:
+ The Transducer 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 = next(model.parameters()).device
+
+ log_interval = 300
+
+ fbank = create_streaming_feature_extractor()
+
+ decode_results = []
+ streams = []
+ for num, cut in enumerate(cuts):
+ # Each utterance has a Stream.
+ stream = Stream(
+ params=params,
+ cut_id=cut.id,
+ decoding_graph=decoding_graph,
+ device=device,
+ LOG_EPS=LOG_EPSILON,
+ )
+
+ stream.set_states(model.encoder.init_states(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)
+ feature = fbank(samples)
+ stream.set_feature(feature)
+ stream.set_ground_truth(cut.supervisions[0].text)
+
+ streams.append(stream)
+
+ while len(streams) >= params.num_decode_streams:
+ finished_streams = decode_one_chunk(
+ model=model,
+ streams=streams,
+ params=params,
+ decoding_graph=decoding_graph,
+ )
+
+ for i in sorted(finished_streams, reverse=True):
+ decode_results.append(
+ (
+ streams[i].id,
+ streams[i].ground_truth.split(),
+ sp.decode(streams[i].decoding_result()).split(),
+ )
+ )
+ del streams[i]
+
+ if num % log_interval == 0:
+ logging.info(f"Cuts processed until now is {num}.")
+
+ while len(streams) > 0:
+ finished_streams = decode_one_chunk(
+ model=model,
+ streams=streams,
+ params=params,
+ decoding_graph=decoding_graph,
+ )
+
+ for i in sorted(finished_streams, reverse=True):
+ decode_results.append(
+ (
+ streams[i].id,
+ streams[i].ground_truth.split(),
+ sp.decode(streams[i].decoding_result()).split(),
+ )
+ )
+ del 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}"
+ )
+ else:
+ key = f"beam_size_{params.beam_size}"
+
+ 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"
+ )
+ store_transcripts(filename=recog_path, texts=sorted(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",
+ "fast_beam_search",
+ "modified_beam_search",
+ )
+ 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-length-{params.chunk_length}"
+ params.suffix += f"-left-context-length-{params.left_context_length}"
+ params.suffix += f"-right-context-length-{params.right_context_length}"
+ params.suffix += f"-memory-size-{params.memory_size}"
+
+ 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}"
+ 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-streaming-decode")
+ 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()
+
+ params.device = device
+
+ 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.eval()
+
+ if params.decoding_method == "fast_beam_search":
+ 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}")
+
+ 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,
+ model=model,
+ params=params,
+ 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__":
+ torch.manual_seed(20220410)
+ main()
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless/test_emformer.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless/test_emformer.py
new file mode 100644
index 000000000..8cde6205b
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless/test_emformer.py
@@ -0,0 +1,194 @@
+#!/usr/bin/env python3
+#
+# Copyright 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.
+
+
+import torch
+from emformer import ConvolutionModule, Emformer, stack_states, unstack_states
+
+
+def test_convolution_module_forward():
+ B, D = 2, 256
+ chunk_length = 4
+ right_context_length = 2
+ num_chunks = 3
+ U = num_chunks * chunk_length
+ R = num_chunks * right_context_length
+ kernel_size = 31
+ conv_module = ConvolutionModule(
+ chunk_length,
+ right_context_length,
+ D,
+ kernel_size,
+ )
+
+ utterance = torch.randn(U, B, D)
+ right_context = torch.randn(R, B, D)
+
+ utterance, right_context = conv_module(utterance, right_context)
+ assert utterance.shape == (U, B, D), utterance.shape
+ assert right_context.shape == (R, B, D), right_context.shape
+
+
+def test_convolution_module_infer():
+ from emformer import ConvolutionModule
+
+ B, D = 2, 256
+ chunk_length = 4
+ right_context_length = 2
+ num_chunks = 1
+ U = num_chunks * chunk_length
+ R = num_chunks * right_context_length
+ kernel_size = 31
+ conv_module = ConvolutionModule(
+ chunk_length,
+ right_context_length,
+ D,
+ kernel_size,
+ )
+
+ utterance = torch.randn(U, B, D)
+ right_context = torch.randn(R, B, D)
+ cache = torch.randn(B, D, kernel_size - 1)
+
+ utterance, right_context, new_cache = conv_module.infer(
+ utterance, right_context, cache
+ )
+ assert utterance.shape == (U, B, D), utterance.shape
+ assert right_context.shape == (R, B, D), right_context.shape
+ assert new_cache.shape == (B, D, kernel_size - 1), new_cache.shape
+
+
+def test_state_stack_unstack():
+ num_features = 80
+ chunk_length = 32
+ encoder_dim = 512
+ num_encoder_layers = 2
+ kernel_size = 31
+ left_context_length = 32
+ right_context_length = 8
+ memory_size = 32
+
+ model = Emformer(
+ num_features=num_features,
+ chunk_length=chunk_length,
+ subsampling_factor=4,
+ d_model=encoder_dim,
+ num_encoder_layers=num_encoder_layers,
+ cnn_module_kernel=kernel_size,
+ left_context_length=left_context_length,
+ right_context_length=right_context_length,
+ memory_size=memory_size,
+ )
+
+ for batch_size in [1, 2]:
+ attn_caches = [
+ [
+ torch.zeros(memory_size, batch_size, encoder_dim),
+ torch.zeros(left_context_length // 4, batch_size, encoder_dim),
+ torch.zeros(
+ left_context_length // 4,
+ batch_size,
+ encoder_dim,
+ ),
+ ]
+ for _ in range(num_encoder_layers)
+ ]
+ conv_caches = [
+ torch.zeros(batch_size, encoder_dim, kernel_size - 1)
+ for _ in range(num_encoder_layers)
+ ]
+ states = [attn_caches, conv_caches]
+ x = torch.randn(batch_size, 23, num_features)
+ x_lens = torch.full((batch_size,), 23)
+ num_processed_frames = torch.full((batch_size,), 0)
+ y, y_lens, states = model.infer(
+ x, x_lens, num_processed_frames=num_processed_frames, states=states
+ )
+
+ state_list = unstack_states(states)
+ states2 = stack_states(state_list)
+
+ for ss, ss2 in zip(states[0], states2[0]):
+ for s, s2 in zip(ss, ss2):
+ assert torch.allclose(s, s2), f"{s.sum()}, {s2.sum()}"
+
+ for s, s2 in zip(states[1], states2[1]):
+ assert torch.allclose(s, s2), f"{s.sum()}, {s2.sum()}"
+
+
+def test_torchscript_consistency_infer():
+ r"""Verify that scripting Emformer does not change the behavior of method `infer`.""" # noqa
+ num_features = 80
+ chunk_length = 32
+ encoder_dim = 512
+ num_encoder_layers = 2
+ kernel_size = 31
+ left_context_length = 32
+ right_context_length = 8
+ memory_size = 32
+ batch_size = 2
+
+ model = Emformer(
+ num_features=num_features,
+ chunk_length=chunk_length,
+ subsampling_factor=4,
+ d_model=encoder_dim,
+ num_encoder_layers=num_encoder_layers,
+ cnn_module_kernel=kernel_size,
+ left_context_length=left_context_length,
+ right_context_length=right_context_length,
+ memory_size=memory_size,
+ ).eval()
+ attn_caches = [
+ [
+ torch.zeros(memory_size, batch_size, encoder_dim),
+ torch.zeros(left_context_length // 4, batch_size, encoder_dim),
+ torch.zeros(
+ left_context_length // 4,
+ batch_size,
+ encoder_dim,
+ ),
+ ]
+ for _ in range(num_encoder_layers)
+ ]
+ conv_caches = [
+ torch.zeros(batch_size, encoder_dim, kernel_size - 1)
+ for _ in range(num_encoder_layers)
+ ]
+ states = [attn_caches, conv_caches]
+ x = torch.randn(batch_size, 23, num_features)
+ x_lens = torch.full((batch_size,), 23)
+ num_processed_frames = torch.full((batch_size,), 0)
+ y, y_lens, out_states = model.infer(
+ x, x_lens, num_processed_frames=num_processed_frames, states=states
+ )
+
+ sc_model = torch.jit.script(model).eval()
+ sc_y, sc_y_lens, sc_out_states = sc_model.infer(
+ x, x_lens, num_processed_frames=num_processed_frames, states=states
+ )
+
+ assert torch.allclose(y, sc_y)
+
+
+if __name__ == "__main__":
+ test_convolution_module_forward()
+ test_convolution_module_infer()
+ test_state_stack_unstack()
+ test_torchscript_consistency_infer()
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless/train.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless/train.py
new file mode 100755
index 000000000..c07d8f76b
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless/train.py
@@ -0,0 +1,1144 @@
+#!/usr/bin/env python3
+# Copyright 2021 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"
+
+./conv_emformer_transducer_stateless/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --exp-dir conv_emformer_transducer_stateless/exp \
+ --full-libri 1 \
+ --max-duration 280 \
+ --master-port 12321 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32
+
+# For mix precision training:
+./conv_emformer_transducer_stateless/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --use-fp16 1 \
+ --exp-dir conv_emformer_transducer_stateless/exp \
+ --full-libri 1 \
+ --max-duration 300 \
+ --master-port 12321 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32
+"""
+
+
+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 emformer import Emformer
+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, Eve
+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.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.utils import AttributeDict, MetricsTracker, setup_logger, str2bool
+
+LRSchedulerType = Union[
+ torch.optim.lr_scheduler._LRScheduler, optim.LRScheduler
+]
+
+
+def add_model_arguments(parser: argparse.ArgumentParser):
+ parser.add_argument(
+ "--encoder-dim",
+ type=int,
+ default=512,
+ help="Attention dim for the Emformer",
+ )
+
+ parser.add_argument(
+ "--nhead",
+ type=int,
+ default=8,
+ help="Number of attention heads for the Emformer",
+ )
+
+ parser.add_argument(
+ "--dim-feedforward",
+ type=int,
+ default=2048,
+ help="Feed-forward dimension for the Emformer",
+ )
+
+ parser.add_argument(
+ "--num-encoder-layers",
+ type=int,
+ default=12,
+ help="Number of encoder layers for the Emformer",
+ )
+
+ parser.add_argument(
+ "--cnn-module-kernel",
+ type=int,
+ default=31,
+ help="Kernel size for the convolution module.",
+ )
+
+ parser.add_argument(
+ "--left-context-length",
+ type=int,
+ default=32,
+ help="""Number of frames before subsampling for left context
+ in the Emformer.""",
+ )
+
+ parser.add_argument(
+ "--chunk-length",
+ type=int,
+ default=32,
+ help="""Number of frames before subsampling for each chunk
+ in the Emformer.""",
+ )
+
+ parser.add_argument(
+ "--right-context-length",
+ type=int,
+ default=8,
+ help="""Number of frames before subsampling for right context
+ in the Emformer.""",
+ )
+
+ parser.add_argument(
+ "--memory-size",
+ type=int,
+ default=0,
+ help="Number of entries in the memory for the Emformer",
+ )
+
+
+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_stateless2/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(
+ "--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(
+ "--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(
+ "--save-every-n",
+ type=int,
+ default=8000,
+ 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=20,
+ 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.
+
+ - encoder_dim: Hidden dim for multi-head attention model.
+
+ - num_decoder_layers: Number of decoder layer of transformer decoder.
+
+ - 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": 50,
+ "reset_interval": 200,
+ "valid_interval": 3000, # For the 100h subset, use 800
+ # parameters for Emformer
+ "feature_dim": 80,
+ "subsampling_factor": 4,
+ # parameters for decoder
+ "decoder_dim": 512,
+ # parameters for joiner
+ "joiner_dim": 512,
+ # parameters for Noam
+ "model_warm_step": 3000, # arg given to model, not for lrate
+ "env_info": get_env_info(),
+ }
+ )
+
+ return params
+
+
+def get_encoder_model(params: AttributeDict) -> nn.Module:
+ # TODO: We can add an option to switch between Conformer and Transformer
+ encoder = Emformer(
+ num_features=params.feature_dim,
+ chunk_length=params.chunk_length,
+ subsampling_factor=params.subsampling_factor,
+ d_model=params.encoder_dim,
+ nhead=params.nhead,
+ dim_feedforward=params.dim_feedforward,
+ num_encoder_layers=params.num_encoder_layers,
+ cnn_module_kernel=params.cnn_module_kernel,
+ left_context_length=params.left_context_length,
+ right_context_length=params.right_context_length,
+ memory_size=params.memory_size,
+ )
+ 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=params.encoder_dim,
+ 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=params.encoder_dim,
+ 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,
+ warmup: float = 1.0,
+) -> Tuple[Tensor, MetricsTracker]:
+ """
+ Compute RNN-T 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.
+ 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)
+
+ 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,
+ warmup=warmup,
+ )
+ # after the main warmup step, we keep pruned_loss_scale small
+ # for the same amount of time (model_warm_step), to avoid
+ # overwhelming the simple_loss and causing it to diverge,
+ # in case it had not fully learned the alignment yet.
+ pruned_loss_scale = (
+ 0.0
+ if warmup < 1.0
+ else (0.1 if warmup > 1.0 and warmup < 2.0 else 1.0)
+ )
+ loss = (
+ params.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()
+ )
+
+ # `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["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"])
+
+ 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,
+ 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()
+
+ 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 % 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,
+ sp=sp,
+ 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))
+ 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)
+
+ 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 = Eve(model.parameters(), lr=params.initial_lr)
+
+ 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)
+
+ 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
+ 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,
+ sp=sp,
+ 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)
+
+ 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 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,
+ 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,
+ sp=sp,
+ batch=batch,
+ is_training=True,
+ warmup=warmup,
+ )
+ loss.backward()
+ optimizer.step()
+ optimizer.zero_grad()
+ except RuntimeError 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]}) ..."
+ )
+ raise
+
+
+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/conv_emformer_transducer_stateless2/asr_datamodule.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/asr_datamodule.py
new file mode 120000
index 000000000..104eeea5d
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/asr_datamodule.py
@@ -0,0 +1 @@
+../conv_emformer_transducer_stateless/asr_datamodule.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/beam_search.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/beam_search.py
new file mode 120000
index 000000000..8554e44cc
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/beam_search.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/beam_search.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/decode.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/decode.py
new file mode 100755
index 000000000..98b8290b5
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/decode.py
@@ -0,0 +1,661 @@
+#!/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
+./conv_emformer_transducer_stateless2/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --max-duration 300 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method greedy_search \
+ --use-averaged-model True
+
+(2) modified beam search
+./conv_emformer_transducer_stateless2/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --max-duration 300 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method modified_beam_search \
+ --use-averaged-model True \
+ --beam-size 4
+
+(3) fast beam search
+./conv_emformer_transducer_stateless2/decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --max-duration 300 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method fast_beam_search \
+ --use-averaged-model True \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8
+"""
+
+
+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_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.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=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_stateless4/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="""Possible values are:
+ - greedy_search
+ - modified_beam_search
+ - fast_beam_search
+ """,
+ )
+
+ 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=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=8,
+ 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(
+ "--max-sym-per-frame",
+ type=int,
+ default=1,
+ help="""Maximum number of symbols per frame.
+ Used only when --decoding_method is greedy_search""",
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+def decode_one_batch(
+ params: AttributeDict,
+ model: nn.Module,
+ sp: spm.SentencePieceProcessor,
+ 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.
+ sp:
+ The BPE 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 = 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 += params.chunk_length
+ feature = torch.nn.functional.pad(
+ feature,
+ pad=(0, 0, 0, params.chunk_length),
+ 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 == "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 params.decoding_method == "fast_beam_search":
+ return {
+ (
+ f"beam_{params.beam}_"
+ f"max_contexts_{params.max_contexts}_"
+ f"max_states_{params.max_states}"
+ ): 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,
+ 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.
+ 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"]
+ 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,
+ 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",
+ "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}"
+ 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 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 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
+ start = params.epoch - params.avg
+ assert start >= 1
+ 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 params.decoding_method == "fast_beam_search":
+ 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}")
+
+ # 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,
+ 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/conv_emformer_transducer_stateless2/decoder.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/decoder.py
new file mode 120000
index 000000000..1db262df7
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/decoder.py
@@ -0,0 +1 @@
+../conv_emformer_transducer_stateless/decoder.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/emformer.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/emformer.py
new file mode 100644
index 000000000..f16f5acc7
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/emformer.py
@@ -0,0 +1,1841 @@
+# Copyright 2022 Xiaomi Corporation (Author: 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.
+#
+# It is modified based on
+# 1) https://github.com/pytorch/audio/blob/main/torchaudio/models/emformer.py # noqa
+# 2) https://github.com/pytorch/audio/blob/main/torchaudio/prototype/models/conv_emformer.py # noqa
+
+import math
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from encoder_interface import EncoderInterface
+from scaling import (
+ ActivationBalancer,
+ BasicNorm,
+ DoubleSwish,
+ ScaledConv1d,
+ ScaledConv2d,
+ ScaledLinear,
+)
+
+from icefall.utils import make_pad_mask
+
+
+LOG_EPSILON = math.log(1e-10)
+
+
+def unstack_states(
+ states: Tuple[List[List[torch.Tensor]], List[torch.Tensor]]
+) -> List[Tuple[List[List[torch.Tensor]], List[torch.Tensor]]]:
+ """Unstack the emformer 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.
+
+ Args:
+ states:
+ A tuple of 2 elements.
+ ``states[0]`` is the attention caches of a batch of utterance.
+ ``states[1]`` is the convolution caches of a batch of utterance.
+ ``len(states[0])`` and ``len(states[1])`` both eqaul to number of layers. # noqa
+
+ Returns:
+ A list of states.
+ ``states[i]`` is a tuple of 2 elements of i-th utterance.
+ ``states[i][0]`` is the attention caches of i-th utterance.
+ ``states[i][1]`` is the convolution caches of i-th utterance.
+ ``len(states[i][0])`` and ``len(states[i][1])`` both eqaul to number of layers. # noqa
+ """
+
+ attn_caches, conv_caches = states
+ batch_size = conv_caches[0].size(0)
+ num_layers = len(attn_caches)
+
+ list_attn_caches = [None] * batch_size
+ for i in range(batch_size):
+ list_attn_caches[i] = [[] for _ in range(num_layers)]
+ for li, layer in enumerate(attn_caches):
+ for s in layer:
+ s_list = s.unbind(dim=1)
+ for bi, b in enumerate(list_attn_caches):
+ b[li].append(s_list[bi])
+
+ list_conv_caches = [None] * batch_size
+ for i in range(batch_size):
+ list_conv_caches[i] = [None] * num_layers
+ for li, layer in enumerate(conv_caches):
+ c_list = layer.unbind(dim=0)
+ for bi, b in enumerate(list_conv_caches):
+ b[li] = c_list[bi]
+
+ ans = [None] * batch_size
+ for i in range(batch_size):
+ ans[i] = [list_attn_caches[i], list_conv_caches[i]]
+
+ return ans
+
+
+def stack_states(
+ state_list: List[Tuple[List[List[torch.Tensor]], List[torch.Tensor]]]
+) -> Tuple[List[List[torch.Tensor]], List[torch.Tensor]]:
+ """Stack list of emformer states that correspond to separate utterances
+ into a single emformer state so that it can be used as an input for
+ emformer 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 emformer model for a single utterance.
+ ``states[i]`` is a tuple of 2 elements of i-th utterance.
+ ``states[i][0]`` is the attention caches of i-th utterance.
+ ``states[i][1]`` is the convolution caches of i-th utterance.
+ ``len(states[i][0])`` and ``len(states[i][1])`` both eqaul to number of layers. # noqa
+
+ 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)
+
+ attn_caches = []
+ for layer in state_list[0][0]:
+ if batch_size > 1:
+ # Note: We will stack attn_caches[layer][s][] later to get attn_caches[layer][s] # noqa
+ attn_caches.append([[s] for s in layer])
+ else:
+ attn_caches.append([s.unsqueeze(1) for s in layer])
+ for b, states in enumerate(state_list[1:], 1):
+ for li, layer in enumerate(states[0]):
+ for si, s in enumerate(layer):
+ attn_caches[li][si].append(s)
+ if b == batch_size - 1:
+ attn_caches[li][si] = torch.stack(
+ attn_caches[li][si], dim=1
+ )
+
+ conv_caches = []
+ for layer in state_list[0][1]:
+ if batch_size > 1:
+ # Note: We will stack conv_caches[layer][] later to get conv_caches[layer] # noqa
+ conv_caches.append([layer])
+ else:
+ conv_caches.append(layer.unsqueeze(0))
+ for b, states in enumerate(state_list[1:], 1):
+ for li, layer in enumerate(states[1]):
+ conv_caches[li].append(layer)
+ if b == batch_size - 1:
+ conv_caches[li] = torch.stack(conv_caches[li], dim=0)
+
+ return [attn_caches, conv_caches]
+
+
+class ConvolutionModule(nn.Module):
+ """ConvolutionModule.
+
+ Modified from https://github.com/pytorch/audio/blob/main/torchaudio/prototype/models/conv_emformer.py # noqa
+
+ Args:
+ chunk_length (int):
+ Length of each chunk.
+ right_context_length (int):
+ Length of right context.
+ channels (int):
+ The number of input channels and output 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,
+ chunk_length: int,
+ right_context_length: int,
+ channels: int,
+ kernel_size: int,
+ bias: bool = True,
+ ) -> None:
+ """Construct an ConvolutionModule object."""
+ super().__init__()
+ # kernerl_size should be an odd number for 'SAME' padding
+ assert (kernel_size - 1) % 2 == 0, kernel_size
+
+ self.chunk_length = chunk_length
+ self.right_context_length = right_context_length
+ self.channels = channels
+
+ 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
+ )
+
+ # make it causal by padding cached (kernel_size - 1) frames on the left
+ self.cache_size = kernel_size - 1
+ self.depthwise_conv = ScaledConv1d(
+ channels,
+ channels,
+ kernel_size,
+ stride=1,
+ padding=0,
+ 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 _split_right_context(
+ self,
+ pad_utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ ) -> torch.Tensor:
+ """
+ Args:
+ pad_utterance:
+ Its shape is (cache_size + U, B, D).
+ right_context:
+ Its shape is (R, B, D).
+
+ Returns:
+ Right context segments padding with corresponding context.
+ Its shape is (num_segs * B, D, cache_size + right_context_length).
+ """
+ U_, B, D = pad_utterance.size()
+ R = right_context.size(0)
+ assert self.right_context_length != 0
+ assert R % self.right_context_length == 0
+ num_chunks = R // self.right_context_length
+ right_context = right_context.reshape(
+ num_chunks, self.right_context_length, B, D
+ )
+ right_context = right_context.permute(0, 2, 1, 3).reshape(
+ num_chunks * B, self.right_context_length, D
+ )
+
+ intervals = torch.arange(
+ 0, self.chunk_length * (num_chunks - 1), self.chunk_length
+ )
+ first = torch.arange(
+ self.chunk_length, self.chunk_length + self.cache_size
+ )
+ indexes = intervals.unsqueeze(1) + first.unsqueeze(0)
+ indexes = torch.cat(
+ [indexes, torch.arange(U_ - self.cache_size, U_).unsqueeze(0)]
+ )
+ padding = pad_utterance[indexes] # (num_chunks, cache_size, B, D)
+ padding = padding.permute(0, 2, 1, 3).reshape(
+ num_chunks * B, self.cache_size, D
+ )
+
+ pad_right_context = torch.cat([padding, right_context], dim=1)
+ # (num_chunks * B, cache_size + right_context_length, D)
+ return pad_right_context.permute(0, 2, 1)
+
+ def _merge_right_context(
+ self, right_context: torch.Tensor, B: int
+ ) -> torch.Tensor:
+ """
+ Args:
+ right_context:
+ Right context segments.
+ It shape is (num_segs * B, D, right_context_length).
+ B:
+ Batch size.
+
+ Returns:
+ A tensor of shape (B, D, R), where
+ R = num_segs * right_context_length.
+ """
+ right_context = right_context.reshape(
+ -1, B, self.channels, self.right_context_length
+ )
+ right_context = right_context.permute(1, 2, 0, 3)
+ right_context = right_context.reshape(B, self.channels, -1)
+ return right_context
+
+ def forward(
+ self,
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Causal convolution module.
+
+ Args:
+ utterance (torch.Tensor):
+ Utterance tensor of shape (U, B, D).
+ right_context (torch.Tensor):
+ Right context tensor of shape (R, B, D).
+
+ Returns:
+ A tuple of 2 tensors:
+ - output utterance of shape (U, B, D).
+ - output right_context of shape (R, B, D).
+ """
+ U, B, D = utterance.size()
+ R, _, _ = right_context.size()
+
+ # point-wise conv and GLU mechanism
+ x = torch.cat([right_context, utterance], dim=0) # (R + U, B, D)
+ x = x.permute(1, 2, 0) # (B, D, R + U)
+ x = self.pointwise_conv1(x) # (B, 2 * D, R + U)
+ x = self.deriv_balancer1(x)
+ x = nn.functional.glu(x, dim=1) # (B, D, R + U)
+ utterance = x[:, :, R:] # (B, D, U)
+ right_context = x[:, :, :R] # (B, D, R)
+
+ # make causal convolution
+ cache = torch.zeros(
+ B, D, self.cache_size, device=x.device, dtype=x.dtype
+ )
+ pad_utterance = torch.cat(
+ [cache, utterance], dim=2
+ ) # (B, D, cache + U)
+
+ # depth-wise conv on utterance
+ utterance = self.depthwise_conv(pad_utterance) # (B, D, U)
+
+ if self.right_context_length > 0:
+ # depth-wise conv on right_context
+ pad_right_context = self._split_right_context(
+ pad_utterance.permute(2, 0, 1), right_context.permute(2, 0, 1)
+ ) # (num_segs * B, D, cache_size + right_context_length)
+ right_context = self.depthwise_conv(
+ pad_right_context
+ ) # (num_segs * B, D, right_context_length)
+ right_context = self._merge_right_context(
+ right_context, B
+ ) # (B, D, R)
+
+ x = torch.cat([right_context, utterance], dim=2) # (B, D, R + U)
+ x = self.deriv_balancer2(x)
+ x = self.activation(x)
+
+ # point-wise conv
+ x = self.pointwise_conv2(x) # (B, D, R + U)
+
+ right_context = x[:, :, :R] # (B, D, R)
+ utterance = x[:, :, R:] # (B, D, U)
+ return (
+ utterance.permute(2, 0, 1),
+ right_context.permute(2, 0, 1),
+ )
+
+ @torch.jit.export
+ def infer(
+ self,
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ cache: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Causal convolution module applied on both utterance and right_context.
+
+ Args:
+ utterance (torch.Tensor):
+ Utterance tensor of shape (U, B, D).
+ right_context (torch.Tensor):
+ Right context tensor of shape (R, B, D).
+ cache (torch.Tensor, optional):
+ Cached tensor for left padding of shape (B, D, cache_size).
+
+ Returns:
+ A tuple of 3 tensors:
+ - output utterance of shape (U, B, D).
+ - output right_context of shape (R, B, D).
+ - updated cache tensor of shape (B, D, cache_size).
+ """
+ U, B, D = utterance.size()
+ R, _, _ = right_context.size()
+
+ # point-wise conv
+ x = torch.cat([utterance, right_context], dim=0) # (U + R, B, D)
+ x = x.permute(1, 2, 0) # (B, D, U + R)
+ x = self.pointwise_conv1(x) # (B, 2 * D, U + R)
+ x = self.deriv_balancer1(x)
+ x = nn.functional.glu(x, dim=1) # (B, D, U + R)
+
+ # make causal convolution
+ assert cache.shape == (B, D, self.cache_size), cache.shape
+ x = torch.cat([cache, x], dim=2) # (B, D, cache_size + U + R)
+ # update cache
+ new_cache = x[:, :, -R - self.cache_size : -R]
+
+ # 1-D depth-wise conv
+ x = self.depthwise_conv(x) # (B, D, U + R)
+
+ x = self.deriv_balancer2(x)
+ x = self.activation(x)
+
+ # point-wise conv
+ x = self.pointwise_conv2(x) # (B, D, U + R)
+
+ utterance = x[:, :, :U] # (B, D, U)
+ right_context = x[:, :, U:] # (B, D, R)
+ return (
+ utterance.permute(2, 0, 1),
+ right_context.permute(2, 0, 1),
+ new_cache,
+ )
+
+
+class EmformerAttention(nn.Module):
+ r"""Emformer layer attention module.
+
+ Args:
+ embed_dim (int):
+ Embedding dimension.
+ nhead (int):
+ Number of attention heads in each Emformer layer.
+ dropout (float, optional):
+ Dropout probability. (Default: 0.0)
+ tanh_on_mem (bool, optional):
+ If ``True``, applies tanh to memory elements. (Default: ``False``)
+ negative_inf (float, optional):
+ Value to use for negative infinity in attention weights. (Default: -1e8)
+ """
+
+ def __init__(
+ self,
+ embed_dim: int,
+ nhead: int,
+ dropout: float = 0.0,
+ tanh_on_mem: bool = False,
+ negative_inf: float = -1e8,
+ ):
+ super().__init__()
+
+ if embed_dim % nhead != 0:
+ raise ValueError(
+ f"embed_dim ({embed_dim}) is not a multiple of"
+ f"nhead ({nhead})."
+ )
+
+ self.embed_dim = embed_dim
+ self.nhead = nhead
+ self.tanh_on_mem = tanh_on_mem
+ self.negative_inf = negative_inf
+ self.head_dim = embed_dim // nhead
+ self.dropout = dropout
+
+ self.emb_to_key_value = ScaledLinear(
+ embed_dim, 2 * embed_dim, bias=True
+ )
+ self.emb_to_query = ScaledLinear(embed_dim, embed_dim, bias=True)
+ self.out_proj = ScaledLinear(
+ embed_dim, embed_dim, bias=True, initial_scale=0.25
+ )
+
+ def _gen_attention_probs(
+ self,
+ attention_weights: torch.Tensor,
+ attention_mask: torch.Tensor,
+ padding_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """Given the entire attention weights, mask out unecessary connections
+ and optionally with padding positions, to obtain underlying chunk-wise
+ attention probabilities.
+
+ B: batch size;
+ Q: length of query;
+ KV: length of key and value.
+
+ Args:
+ attention_weights (torch.Tensor):
+ Attention weights computed on the entire concatenated tensor
+ with shape (B * nhead, Q, KV).
+ attention_mask (torch.Tensor):
+ Mask tensor where chunk-wise connections are filled with `False`,
+ and other unnecessary connections are filled with `True`,
+ with shape (Q, KV).
+ padding_mask (torch.Tensor, optional):
+ Mask tensor where the padding positions are fill with `True`,
+ and other positions are filled with `False`, with shapa `(B, KV)`.
+
+ Returns:
+ A tensor of shape (B * nhead, Q, KV).
+ """
+ attention_weights_float = attention_weights.float()
+ attention_weights_float = attention_weights_float.masked_fill(
+ attention_mask.unsqueeze(0), self.negative_inf
+ )
+ if padding_mask is not None:
+ Q = attention_weights.size(1)
+ B = attention_weights.size(0) // self.nhead
+ attention_weights_float = attention_weights_float.view(
+ B, self.nhead, Q, -1
+ )
+ attention_weights_float = attention_weights_float.masked_fill(
+ padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool),
+ self.negative_inf,
+ )
+ attention_weights_float = attention_weights_float.view(
+ B * self.nhead, Q, -1
+ )
+
+ attention_probs = nn.functional.softmax(
+ attention_weights_float, dim=-1
+ ).type_as(attention_weights)
+
+ attention_probs = nn.functional.dropout(
+ attention_probs, p=self.dropout, training=self.training
+ )
+ return attention_probs
+
+ def _forward_impl(
+ self,
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ memory: torch.Tensor,
+ attention_mask: torch.Tensor,
+ padding_mask: Optional[torch.Tensor] = None,
+ left_context_key: Optional[torch.Tensor] = None,
+ left_context_val: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Underlying chunk-wise attention implementation."""
+ U, B, _ = utterance.size()
+ R = right_context.size(0)
+ M = memory.size(0)
+ scaling = float(self.head_dim) ** -0.5
+
+ # compute query with [right_context, utterance].
+ query = self.emb_to_query(torch.cat([right_context, utterance]))
+ # compute key and value with [memory, right_context, utterance].
+ key, value = self.emb_to_key_value(
+ torch.cat([memory, right_context, utterance])
+ ).chunk(chunks=2, dim=2)
+
+ if left_context_key is not None and left_context_val is not None:
+ # now compute key and value with
+ # [memory, right context, left context, uttrance]
+ # this is used in inference mode
+ key = torch.cat([key[: M + R], left_context_key, key[M + R :]])
+ value = torch.cat(
+ [value[: M + R], left_context_val, value[M + R :]]
+ )
+ Q = query.size(0)
+ # KV = key.size(0)
+
+ reshaped_query, reshaped_key, reshaped_value = [
+ tensor.contiguous()
+ .view(-1, B * self.nhead, self.head_dim)
+ .transpose(0, 1)
+ for tensor in [query, key, value]
+ ] # (B * nhead, Q or KV, head_dim)
+ attention_weights = torch.bmm(
+ reshaped_query * scaling, reshaped_key.transpose(1, 2)
+ ) # (B * nhead, Q, KV)
+
+ # compute attention probabilities
+ attention_probs = self._gen_attention_probs(
+ attention_weights, attention_mask, padding_mask
+ )
+
+ # compute attention outputs
+ attention = torch.bmm(attention_probs, reshaped_value)
+ assert attention.shape == (B * self.nhead, Q, self.head_dim)
+ attention = (
+ attention.transpose(0, 1).contiguous().view(Q, B, self.embed_dim)
+ )
+
+ # apply output projection
+ output_right_context_utterance = self.out_proj(attention)
+
+ return output_right_context_utterance, key, value
+
+ def forward(
+ self,
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ memory: torch.Tensor,
+ attention_mask: torch.Tensor,
+ padding_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ # TODO: Modify docs.
+ """Forward pass for training and validation mode.
+
+ B: batch size;
+ D: embedding dimension;
+ R: length of the hard-copied right contexts;
+ U: length of full utterance;
+ M: length of memory vectors.
+
+ It computes a `big` attention matrix on full utterance and
+ then utilizes a pre-computed mask to simulate chunk-wise attention.
+
+ It concatenates three blocks: hard-copied right contexts,
+ and full utterance, as a `big` block,
+ to compute the query tensor:
+ query = [right_context, utterance],
+ with length Q = R + U.
+ It concatenates the three blocks: memory vectors,
+ hard-copied right contexts, and full utterance as another `big` block,
+ to compute the key and value tensors:
+ key & value = [memory, right_context, utterance],
+ with length KV = M + R + U.
+ Attention scores is computed with above `big` query and key.
+
+ Then the underlying chunk-wise attention is obtained by applying
+ the attention mask. Suppose
+ c_i: chunk at index i;
+ r_i: right context that c_i can use;
+ l_i: left context that c_i can use;
+ m_i: past memory vectors from previous layer that c_i can use;
+ The target chunk-wise attention is:
+ c_i, r_i (in query) -> l_i, c_i, r_i, m_i (in key)
+
+ Args:
+ utterance (torch.Tensor):
+ Full utterance frames, with shape (U, B, D).
+ right_context (torch.Tensor):
+ Hard-copied right context frames, with shape (R, B, D),
+ where R = num_chunks * right_context_length
+ memory (torch.Tensor):
+ Memory elements, with shape (M, B, D), where M = num_chunks - 1.
+ It is an empty tensor without using memory.
+ attention_mask (torch.Tensor):
+ Pre-computed attention mask to simulate underlying chunk-wise
+ attention, with shape (Q, KV).
+ padding_mask (torch.Tensor):
+ Padding mask of key tensor, with shape (B, KV).
+
+ Returns:
+ Output of right context and utterance, with shape (R + U, B, D).
+ """
+ output_right_context_utterance, _, _ = self._forward_impl(
+ utterance,
+ right_context,
+ memory,
+ attention_mask,
+ padding_mask=padding_mask,
+ )
+ return output_right_context_utterance
+
+ @torch.jit.export
+ def infer(
+ self,
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ memory: torch.Tensor,
+ left_context_key: torch.Tensor,
+ left_context_val: torch.Tensor,
+ padding_mask: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Forward pass for inference.
+
+ B: batch size;
+ D: embedding dimension;
+ R: length of right context;
+ U: length of utterance, i.e., current chunk;
+ L: length of cached left context;
+ M: length of cached memory vectors.
+
+ It concatenates the right context and utterance (i.e., current chunk)
+ of current chunk, to compute the query tensor:
+ query = [right_context, utterance],
+ with length Q = R + U.
+ It concatenates the memory vectors, right context, left context, and
+ current chunk, to compute the key and value tensors:
+ key & value = [memory, right_context, left_context, utterance],
+ with length KV = M + R + L + U.
+
+ The chunk-wise attention is:
+ chunk, right context (in query) ->
+ left context, chunk, right context, memory vectors (in key).
+
+ Args:
+ utterance (torch.Tensor):
+ Current chunk frames, with shape (U, B, D), where U = chunk_length.
+ right_context (torch.Tensor):
+ Right context frames, with shape (R, B, D),
+ where R = right_context_length.
+ memory (torch.Tensor):
+ Memory vectors, with shape (M, B, D), or empty tensor.
+ left_context_key (torch,Tensor):
+ Cached attention key of left context from preceding computation,
+ with shape (L, B, D).
+ left_context_val (torch.Tensor):
+ Cached attention value of left context from preceding computation,
+ with shape (L, B, D).
+ padding_mask (torch.Tensor):
+ Padding mask of key tensor, with shape (B, KV).
+
+ Returns:
+ A tuple containing 4 tensors:
+ - output of right context and utterance, with shape (R + U, B, D).
+ - attention key of left context and utterance, which would be cached
+ for next computation, with shape (L + U, B, D).
+ - attention value of left context and utterance, which would be
+ cached for next computation, with shape (L + U, B, D).
+ """
+ U = utterance.size(0)
+ R = right_context.size(0)
+ L = left_context_key.size(0)
+ M = memory.size(0)
+
+ # query = [right context, utterance]
+ Q = R + U
+ # key, value = [memory, right context, left context, uttrance]
+ KV = M + R + L + U
+ attention_mask = torch.zeros(Q, KV).to(
+ dtype=torch.bool, device=utterance.device
+ )
+
+ output_right_context_utterance, key, value = self._forward_impl(
+ utterance,
+ right_context,
+ memory,
+ attention_mask,
+ padding_mask=padding_mask,
+ left_context_key=left_context_key,
+ left_context_val=left_context_val,
+ )
+ return (
+ output_right_context_utterance,
+ key[M + R :],
+ value[M + R :],
+ )
+
+
+class EmformerEncoderLayer(nn.Module):
+ """Emformer layer that constitutes Emformer.
+
+ Args:
+ d_model (int):
+ Input dimension.
+ nhead (int):
+ Number of attention heads.
+ dim_feedforward (int):
+ Hidden layer dimension of feedforward network.
+ chunk_length (int):
+ Length of each input segment.
+ dropout (float, optional):
+ Dropout probability. (Default: 0.0)
+ layer_dropout (float, optional):
+ Layer dropout probability. (Default: 0.0)
+ cnn_module_kernel (int):
+ Kernel size of convolution module.
+ left_context_length (int, optional):
+ Length of left context. (Default: 0)
+ right_context_length (int, optional):
+ Length of right context. (Default: 0)
+ memory_size (int, optional):
+ Number of memory elements to use. (Default: 0)
+ tanh_on_mem (bool, optional):
+ If ``True``, applies tanh to memory elements. (Default: ``False``)
+ negative_inf (float, optional):
+ Value to use for negative infinity in attention weights. (Default: -1e8)
+ """
+
+ def __init__(
+ self,
+ d_model: int,
+ nhead: int,
+ dim_feedforward: int,
+ chunk_length: int,
+ dropout: float = 0.1,
+ layer_dropout: float = 0.075,
+ cnn_module_kernel: int = 31,
+ left_context_length: int = 0,
+ right_context_length: int = 0,
+ memory_size: int = 0,
+ tanh_on_mem: bool = False,
+ negative_inf: float = -1e8,
+ ):
+ super().__init__()
+
+ self.attention = EmformerAttention(
+ embed_dim=d_model,
+ nhead=nhead,
+ dropout=dropout,
+ tanh_on_mem=tanh_on_mem,
+ negative_inf=negative_inf,
+ )
+ self.summary_op = nn.AvgPool1d(
+ kernel_size=chunk_length, stride=chunk_length, ceil_mode=True
+ )
+
+ 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.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.conv_module = ConvolutionModule(
+ chunk_length,
+ right_context_length,
+ 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)
+
+ self.layer_dropout = layer_dropout
+ self.left_context_length = left_context_length
+ self.chunk_length = chunk_length
+ self.memory_size = memory_size
+ self.d_model = d_model
+ self.use_memory = memory_size > 0
+
+ def _update_attn_cache(
+ self,
+ next_key: torch.Tensor,
+ next_val: torch.Tensor,
+ memory: torch.Tensor,
+ attn_cache: List[torch.Tensor],
+ ) -> List[torch.Tensor]:
+ """Update cached attention state:
+ 1) output memory of current chunk in the lower layer;
+ 2) attention key and value in current chunk's computation, which would
+ be resued in next chunk's computation.
+ """
+ new_memory = torch.cat([attn_cache[0], memory])
+ new_key = torch.cat([attn_cache[1], next_key])
+ new_val = torch.cat([attn_cache[2], next_val])
+ attn_cache[0] = new_memory[new_memory.size(0) - self.memory_size :]
+ attn_cache[1] = new_key[new_key.size(0) - self.left_context_length :]
+ attn_cache[2] = new_val[new_val.size(0) - self.left_context_length :]
+ return attn_cache
+
+ def _apply_conv_module_forward(
+ self,
+ right_context_utterance: torch.Tensor,
+ R: int,
+ ) -> torch.Tensor:
+ """Apply convolution module in training and validation mode."""
+ utterance = right_context_utterance[R:]
+ right_context = right_context_utterance[:R]
+ utterance, right_context = self.conv_module(utterance, right_context)
+ right_context_utterance = torch.cat([right_context, utterance])
+ return right_context_utterance
+
+ def _apply_conv_module_infer(
+ self,
+ right_context_utterance: torch.Tensor,
+ R: int,
+ conv_cache: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Apply convolution module on utterance in inference mode."""
+ utterance = right_context_utterance[R:]
+ right_context = right_context_utterance[:R]
+ utterance, right_context, conv_cache = self.conv_module.infer(
+ utterance, right_context, conv_cache
+ )
+ right_context_utterance = torch.cat([right_context, utterance])
+ return right_context_utterance, conv_cache
+
+ def _apply_attention_module_forward(
+ self,
+ right_context_utterance: torch.Tensor,
+ R: int,
+ attention_mask: torch.Tensor,
+ padding_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """Apply attention module in training and validation mode."""
+ utterance = right_context_utterance[R:]
+ right_context = right_context_utterance[:R]
+
+ if self.use_memory:
+ memory = self.summary_op(utterance.permute(1, 2, 0)).permute(
+ 2, 0, 1
+ )[:-1, :, :]
+ else:
+ memory = torch.empty(0).to(
+ dtype=utterance.dtype, device=utterance.device
+ )
+ output_right_context_utterance = self.attention(
+ utterance=utterance,
+ right_context=right_context,
+ memory=memory,
+ attention_mask=attention_mask,
+ padding_mask=padding_mask,
+ )
+
+ return output_right_context_utterance
+
+ def _apply_attention_module_infer(
+ self,
+ right_context_utterance: torch.Tensor,
+ R: int,
+ attn_cache: List[torch.Tensor],
+ padding_mask: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+ """Apply attention module in inference mode.
+ 1) Unpack cached states including:
+ - memory from previous chunks;
+ - attention key and value of left context from preceding
+ chunk's compuation;
+ 2) Apply attention computation;
+ 3) Update cached attention states including:
+ - memory of current chunk;
+ - attention key and value in current chunk's computation, which would
+ be resued in next chunk's computation.
+ """
+ utterance = right_context_utterance[R:]
+ right_context = right_context_utterance[:R]
+
+ pre_memory = attn_cache[0]
+ left_context_key = attn_cache[1]
+ left_context_val = attn_cache[2]
+
+ if self.use_memory:
+ memory = self.summary_op(utterance.permute(1, 2, 0)).permute(
+ 2, 0, 1
+ )[:1, :, :]
+ else:
+ memory = torch.empty(0).to(
+ dtype=utterance.dtype, device=utterance.device
+ )
+ (
+ output_right_context_utterance,
+ next_key,
+ next_val,
+ ) = self.attention.infer(
+ utterance=utterance,
+ right_context=right_context,
+ memory=pre_memory,
+ left_context_key=left_context_key,
+ left_context_val=left_context_val,
+ padding_mask=padding_mask,
+ )
+ attn_cache = self._update_attn_cache(
+ next_key, next_val, memory, attn_cache
+ )
+ return output_right_context_utterance, attn_cache
+
+ def forward(
+ self,
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ attention_mask: torch.Tensor,
+ padding_mask: Optional[torch.Tensor] = None,
+ warmup: float = 1.0,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ r"""Forward pass for training and validation mode.
+
+ B: batch size;
+ D: embedding dimension;
+ R: length of hard-copied right contexts;
+ U: length of full utterance;
+ M: length of memory vectors.
+
+ Args:
+ utterance (torch.Tensor):
+ Utterance frames, with shape (U, B, D).
+ right_context (torch.Tensor):
+ Right context frames, with shape (R, B, D).
+ attention_mask (torch.Tensor):
+ Attention mask for underlying attention module,
+ with shape (Q, KV), where Q = R + U, KV = M + R + U.
+ padding_mask (torch.Tensor):
+ Padding mask of ker tensor, with shape (B, KV).
+
+ Returns:
+ A tuple containing 2 tensors:
+ - output utterance, with shape (U, B, D).
+ - output right context, with shape (R, B, D).
+ """
+ R = right_context.size(0)
+ src = torch.cat([right_context, utterance])
+ src_orig = src
+
+ warmup_scale = min(0.1 + 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 0.1
+ )
+ else:
+ alpha = 1.0
+
+ # macaron style feed forward module
+ src = src + self.dropout(self.feed_forward_macaron(src))
+
+ # emformer attention module
+ src_att = self._apply_attention_module_forward(
+ src, R, attention_mask, padding_mask=padding_mask
+ )
+ src = src + self.dropout(src_att)
+
+ # convolution module
+ src_conv = self._apply_conv_module_forward(src, R)
+ src = src + self.dropout(src_conv)
+
+ # 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
+
+ output_utterance = src[R:]
+ output_right_context = src[:R]
+ return output_utterance, output_right_context
+
+ @torch.jit.export
+ def infer(
+ self,
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ attn_cache: List[torch.Tensor],
+ conv_cache: torch.Tensor,
+ padding_mask: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor, List[torch.Tensor], torch.Tensor]:
+ """Forward pass for inference.
+
+ B: batch size;
+ D: embedding dimension;
+ R: length of right_context;
+ U: length of utterance;
+ M: length of memory.
+
+ Args:
+ utterance (torch.Tensor):
+ Utterance frames, with shape (U, B, D).
+ right_context (torch.Tensor):
+ Right context frames, with shape (R, B, D).
+ attn_cache (List[torch.Tensor]):
+ Cached attention tensors generated in preceding computation,
+ including memory, key and value of left context.
+ conv_cache (torch.Tensor, optional):
+ Cache tensor of left context for causal convolution.
+ padding_mask (torch.Tensor):
+ Padding mask of ker tensor.
+
+ Returns:
+ (Tensor, Tensor, List[torch.Tensor], Tensor):
+ - output utterance, with shape (U, B, D);
+ - output right_context, with shape (R, B, D);
+ - output attention cache;
+ - output convolution cache.
+ """
+ R = right_context.size(0)
+ src = torch.cat([right_context, utterance])
+
+ # macaron style feed forward module
+ src = src + self.dropout(self.feed_forward_macaron(src))
+
+ # emformer attention module
+ src_att, attn_cache = self._apply_attention_module_infer(
+ src, R, attn_cache, padding_mask=padding_mask
+ )
+ src = src + self.dropout(src_att)
+
+ # convolution module
+ src_conv, conv_cache = self._apply_conv_module_infer(src, R, conv_cache)
+ src = src + self.dropout(src_conv)
+
+ # feed forward module
+ src = src + self.dropout(self.feed_forward(src))
+
+ src = self.norm_final(self.balancer(src))
+
+ output_utterance = src[R:]
+ output_right_context = src[:R]
+ return (
+ output_utterance,
+ output_right_context,
+ attn_cache,
+ conv_cache,
+ )
+
+
+def _gen_attention_mask_block(
+ col_widths: List[int],
+ col_mask: List[bool],
+ num_rows: int,
+ device: torch.device,
+) -> torch.Tensor:
+ assert len(col_widths) == len(
+ col_mask
+ ), "Length of col_widths must match that of col_mask"
+
+ mask_block = [
+ torch.ones(num_rows, col_width, device=device)
+ if is_ones_col
+ else torch.zeros(num_rows, col_width, device=device)
+ for col_width, is_ones_col in zip(col_widths, col_mask)
+ ]
+ return torch.cat(mask_block, dim=1)
+
+
+class EmformerEncoder(nn.Module):
+ """Implements the Emformer architecture introduced in
+ *Emformer: Efficient Memory Transformer Based Acoustic Model for Low Latency
+ Streaming Speech Recognition*
+ [:footcite:`shi2021emformer`].
+
+ In this model, the memory bank computation is simplifed, using the averaged
+ value of each chunk as its memory vector.
+
+ Args:
+ d_model (int):
+ Input dimension.
+ nhead (int):
+ Number of attention heads in each emformer layer.
+ dim_feedforward (int):
+ Hidden layer dimension of each emformer layer's feedforward network.
+ num_encoder_layers (int):
+ Number of emformer layers to instantiate.
+ chunk_length (int):
+ Length of each input segment.
+ dropout (float, optional):
+ Dropout probability. (default: 0.0)
+ layer_dropout (float, optional):
+ Layer dropout probability. (default: 0.0)
+ cnn_module_kernel (int):
+ Kernel size of convolution module.
+ left_context_length (int, optional):
+ Length of left context. (default: 0)
+ right_context_length (int, optional):
+ Length of right context. (default: 0)
+ memory_size (int, optional):
+ Number of memory elements to use. (default: 0)
+ tanh_on_mem (bool, optional):
+ If ``true``, applies tanh to memory elements. (default: ``false``)
+ negative_inf (float, optional):
+ Value to use for negative infinity in attention weights. (default: -1e8)
+ """
+
+ def __init__(
+ self,
+ chunk_length: int,
+ d_model: int = 256,
+ nhead: int = 4,
+ dim_feedforward: int = 2048,
+ num_encoder_layers: int = 12,
+ dropout: float = 0.1,
+ layer_dropout: float = 0.075,
+ cnn_module_kernel: int = 31,
+ left_context_length: int = 0,
+ right_context_length: int = 0,
+ memory_size: int = 0,
+ tanh_on_mem: bool = False,
+ negative_inf: float = -1e8,
+ ):
+ super().__init__()
+
+ assert (
+ chunk_length - 1
+ ) & chunk_length == 0, "chunk_length should be a power of 2."
+ self.shift = int(math.log(chunk_length, 2))
+
+ self.use_memory = memory_size > 0
+
+ self.emformer_layers = nn.ModuleList(
+ [
+ EmformerEncoderLayer(
+ d_model=d_model,
+ nhead=nhead,
+ dim_feedforward=dim_feedforward,
+ chunk_length=chunk_length,
+ dropout=dropout,
+ layer_dropout=layer_dropout,
+ cnn_module_kernel=cnn_module_kernel,
+ left_context_length=left_context_length,
+ right_context_length=right_context_length,
+ memory_size=memory_size,
+ tanh_on_mem=tanh_on_mem,
+ negative_inf=negative_inf,
+ )
+ for layer_idx in range(num_encoder_layers)
+ ]
+ )
+
+ self.num_encoder_layers = num_encoder_layers
+ self.d_model = d_model
+ self.left_context_length = left_context_length
+ self.right_context_length = right_context_length
+ self.chunk_length = chunk_length
+ self.memory_size = memory_size
+ self.cnn_module_kernel = cnn_module_kernel
+
+ def _gen_right_context(self, x: torch.Tensor) -> torch.Tensor:
+ """Hard copy each chunk's right context and concat them."""
+ T = x.shape[0]
+ num_chunks = math.ceil(
+ (T - self.right_context_length) / self.chunk_length
+ )
+ # first (num_chunks - 1) right context block
+ intervals = torch.arange(
+ 0, self.chunk_length * (num_chunks - 1), self.chunk_length
+ )
+ first = torch.arange(
+ self.chunk_length, self.chunk_length + self.right_context_length
+ )
+ indexes = intervals.unsqueeze(1) + first.unsqueeze(0)
+ # cat last right context block
+ indexes = torch.cat(
+ [
+ indexes,
+ torch.arange(T - self.right_context_length, T).unsqueeze(0),
+ ]
+ )
+ right_context_blocks = x[indexes.reshape(-1)]
+ return right_context_blocks
+
+ def _gen_attention_mask_col_widths(
+ self, chunk_idx: int, U: int
+ ) -> List[int]:
+ """Calculate column widths (key, value) in attention mask for the
+ chunk_idx chunk."""
+ num_chunks = math.ceil(U / self.chunk_length)
+ rc = self.right_context_length
+ lc = self.left_context_length
+ rc_start = chunk_idx * rc
+ rc_end = rc_start + rc
+ chunk_start = max(chunk_idx * self.chunk_length - lc, 0)
+ chunk_end = min((chunk_idx + 1) * self.chunk_length, U)
+ R = rc * num_chunks
+
+ if self.use_memory:
+ m_start = max(chunk_idx - self.memory_size, 0)
+ M = num_chunks - 1
+ col_widths = [
+ m_start, # before memory
+ chunk_idx - m_start, # memory
+ M - chunk_idx, # after memory
+ rc_start, # before right context
+ rc, # right context
+ R - rc_end, # after right context
+ chunk_start, # before chunk
+ chunk_end - chunk_start, # chunk
+ U - chunk_end, # after chunk
+ ]
+ else:
+ col_widths = [
+ rc_start, # before right context
+ rc, # right context
+ R - rc_end, # after right context
+ chunk_start, # before chunk
+ chunk_end - chunk_start, # chunk
+ U - chunk_end, # after chunk
+ ]
+
+ return col_widths
+
+ def _gen_attention_mask(self, utterance: torch.Tensor) -> torch.Tensor:
+ """Generate attention mask to simulate underlying chunk-wise attention
+ computation, where chunk-wise connections are filled with `False`,
+ and other unnecessary connections beyond chunk are filled with `True`.
+
+ R: length of hard-copied right contexts;
+ U: length of full utterance;
+ M: length of memory vectors;
+ Q: length of attention query;
+ KV: length of attention key and value.
+
+ The shape of attention mask is (Q, KV).
+ If self.use_memory is `True`:
+ query = [right_context, utterance];
+ key, value = [memory, right_context, utterance];
+ Q = R + U, KV = M + R + U.
+ Otherwise:
+ query = [right_context, utterance]
+ key, value = [right_context, utterance]
+ Q = R + U, KV = R + U.
+
+ Suppose:
+ c_i: chunk at index i;
+ r_i: right context that c_i can use;
+ l_i: left context that c_i can use;
+ m_i: past memory vectors from previous layer that c_i can use;
+ The target chunk-wise attention is:
+ c_i, r_i (in query) -> l_i, c_i, r_i, m_i (in key).
+ """
+ U = utterance.size(0)
+ num_chunks = math.ceil(U / self.chunk_length)
+
+ right_context_mask = []
+ utterance_mask = []
+
+ if self.use_memory:
+ num_cols = 9
+ # right context and utterance both attend to memory, right context,
+ # utterance
+ right_context_utterance_cols_mask = [
+ idx in [1, 4, 7] for idx in range(num_cols)
+ ]
+ else:
+ num_cols = 6
+ # right context and utterance both attend to right context and
+ # utterance
+ right_context_utterance_cols_mask = [
+ idx in [1, 4] for idx in range(num_cols)
+ ]
+ masks_to_concat = [right_context_mask, utterance_mask]
+
+ for chunk_idx in range(num_chunks):
+ col_widths = self._gen_attention_mask_col_widths(chunk_idx, U)
+
+ right_context_mask_block = _gen_attention_mask_block(
+ col_widths,
+ right_context_utterance_cols_mask,
+ self.right_context_length,
+ utterance.device,
+ )
+ right_context_mask.append(right_context_mask_block)
+
+ utterance_mask_block = _gen_attention_mask_block(
+ col_widths,
+ right_context_utterance_cols_mask,
+ min(
+ self.chunk_length,
+ U - chunk_idx * self.chunk_length,
+ ),
+ utterance.device,
+ )
+ utterance_mask.append(utterance_mask_block)
+
+ attention_mask = (
+ 1 - torch.cat([torch.cat(mask) for mask in masks_to_concat])
+ ).to(torch.bool)
+ return attention_mask
+
+ def forward(
+ self, x: torch.Tensor, lengths: torch.Tensor, warmup: float = 1.0
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Forward pass for training and validation mode.
+
+ B: batch size;
+ D: input dimension;
+ U: length of utterance.
+
+ Args:
+ x (torch.Tensor):
+ Utterance frames right-padded with right context frames,
+ with shape (U + right_context_length, B, D).
+ lengths (torch.Tensor):
+ With shape (B,) and i-th element representing number of valid
+ utterance frames for i-th batch element in x, which contains the
+ right_context at the end.
+
+ Returns:
+ A tuple of 2 tensors:
+ - output utterance frames, with shape (U, B, D).
+ - output_lengths, with shape (B,), without containing the
+ right_context at the end.
+ """
+ U = x.size(0) - self.right_context_length
+
+ right_context = self._gen_right_context(x)
+ utterance = x[:U]
+ output_lengths = torch.clamp(lengths - self.right_context_length, min=0)
+ attention_mask = self._gen_attention_mask(utterance)
+
+ M = (
+ right_context.size(0) // self.right_context_length - 1
+ if self.use_memory
+ else 0
+ )
+ padding_mask = make_pad_mask(M + right_context.size(0) + output_lengths)
+
+ output = utterance
+ for layer in self.emformer_layers:
+ output, right_context = layer(
+ output,
+ right_context,
+ attention_mask,
+ padding_mask=padding_mask,
+ warmup=warmup,
+ )
+
+ return output, output_lengths
+
+ @torch.jit.export
+ def infer(
+ self,
+ x: torch.Tensor,
+ lengths: torch.Tensor,
+ num_processed_frames: torch.Tensor,
+ states: Tuple[List[List[torch.Tensor]], List[torch.Tensor]],
+ ) -> Tuple[
+ torch.Tensor,
+ torch.Tensor,
+ Tuple[List[List[torch.Tensor]], List[torch.Tensor]],
+ ]:
+ """Forward pass for streaming inference.
+
+ B: batch size;
+ D: input dimension;
+ U: length of utterance.
+
+ Args:
+ x (torch.Tensor):
+ Utterance frames right-padded with right context frames,
+ with shape (U + right_context_length, B, D).
+ lengths (torch.Tensor):
+ With shape (B,) and i-th element representing number of valid
+ utterance frames for i-th batch element in x, which contains the
+ right_context at the end.
+ states (List[torch.Tensor, List[List[torch.Tensor]], List[torch.Tensor]]: # noqa
+ Cached states containing:
+ - attn_caches: attention states from preceding chunk's computation,
+ where each element corresponds to each emformer layer
+ - conv_caches: left context for causal convolution, where each
+ element corresponds to each layer.
+
+ Returns:
+ (Tensor, Tensor, List[List[torch.Tensor]], List[torch.Tensor]):
+ - output utterance frames, with shape (U, B, D).
+ - output lengths, with shape (B,), without containing the
+ right_context at the end.
+ - updated states from current chunk's computation.
+ """
+ assert num_processed_frames.shape == (x.size(1),)
+
+ attn_caches = states[0]
+ assert len(attn_caches) == self.num_encoder_layers, len(attn_caches)
+ for i in range(len(attn_caches)):
+ assert attn_caches[i][0].shape == (
+ self.memory_size,
+ x.size(1),
+ self.d_model,
+ ), attn_caches[i][0].shape
+ assert attn_caches[i][1].shape == (
+ self.left_context_length,
+ x.size(1),
+ self.d_model,
+ ), attn_caches[i][1].shape
+ assert attn_caches[i][2].shape == (
+ self.left_context_length,
+ x.size(1),
+ self.d_model,
+ ), attn_caches[i][2].shape
+
+ conv_caches = states[1]
+ assert len(conv_caches) == self.num_encoder_layers, len(conv_caches)
+ for i in range(len(conv_caches)):
+ assert conv_caches[i].shape == (
+ x.size(1),
+ self.d_model,
+ self.cnn_module_kernel - 1,
+ ), conv_caches[i].shape
+
+ right_context = x[-self.right_context_length :]
+ utterance = x[: -self.right_context_length]
+ output_lengths = torch.clamp(lengths - self.right_context_length, min=0)
+
+ # calcualte padding mask to mask out initial zero caches
+ chunk_mask = make_pad_mask(output_lengths).to(x.device)
+ memory_mask = (
+ (
+ (num_processed_frames >> self.shift).view(x.size(1), 1)
+ <= torch.arange(self.memory_size, device=x.device).expand(
+ x.size(1), self.memory_size
+ )
+ ).flip(1)
+ if self.use_memory
+ else torch.empty(0).to(dtype=torch.bool, device=x.device)
+ )
+ left_context_mask = (
+ num_processed_frames.view(x.size(1), 1)
+ <= torch.arange(self.left_context_length, device=x.device).expand(
+ x.size(1), self.left_context_length
+ )
+ ).flip(1)
+ right_context_mask = torch.zeros(
+ x.size(1),
+ self.right_context_length,
+ dtype=torch.bool,
+ device=x.device,
+ )
+ padding_mask = torch.cat(
+ [memory_mask, right_context_mask, left_context_mask, chunk_mask],
+ dim=1,
+ )
+
+ output = utterance
+ output_attn_caches: List[List[torch.Tensor]] = []
+ output_conv_caches: List[torch.Tensor] = []
+ for layer_idx, layer in enumerate(self.emformer_layers):
+ (
+ output,
+ right_context,
+ output_attn_cache,
+ output_conv_cache,
+ ) = layer.infer(
+ output,
+ right_context,
+ padding_mask=padding_mask,
+ attn_cache=attn_caches[layer_idx],
+ conv_cache=conv_caches[layer_idx],
+ )
+ output_attn_caches.append(output_attn_cache)
+ output_conv_caches.append(output_conv_cache)
+
+ output_states: Tuple[List[List[torch.Tensor]], List[torch.Tensor]] = (
+ output_attn_caches,
+ output_conv_caches,
+ )
+ return output, output_lengths, output_states
+
+ @torch.jit.export
+ def init_states(self, device: torch.device = torch.device("cpu")):
+ """Create initial 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__(
+ self,
+ num_features: int,
+ chunk_length: int,
+ subsampling_factor: int = 4,
+ d_model: int = 256,
+ nhead: int = 4,
+ dim_feedforward: int = 2048,
+ num_encoder_layers: int = 12,
+ dropout: float = 0.1,
+ layer_dropout: float = 0.075,
+ cnn_module_kernel: int = 3,
+ left_context_length: int = 0,
+ right_context_length: int = 0,
+ memory_size: int = 0,
+ tanh_on_mem: bool = False,
+ negative_inf: float = -1e8,
+ ):
+ super().__init__()
+
+ self.subsampling_factor = subsampling_factor
+ self.right_context_length = right_context_length
+ self.chunk_length = chunk_length
+ if subsampling_factor != 4:
+ raise NotImplementedError("Support only 'subsampling_factor=4'.")
+ if chunk_length % subsampling_factor != 0:
+ raise NotImplementedError(
+ "chunk_length must be a mutiple of subsampling_factor."
+ )
+ if (
+ left_context_length != 0
+ and left_context_length % subsampling_factor != 0
+ ):
+ raise NotImplementedError(
+ "left_context_length must be 0 or a mutiple of subsampling_factor." # noqa
+ )
+ if (
+ right_context_length != 0
+ and right_context_length % subsampling_factor != 0
+ ):
+ raise NotImplementedError(
+ "right_context_length must be 0 or a mutiple of subsampling_factor." # noqa
+ )
+
+ # 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 = EmformerEncoder(
+ chunk_length=chunk_length // subsampling_factor,
+ d_model=d_model,
+ nhead=nhead,
+ dim_feedforward=dim_feedforward,
+ num_encoder_layers=num_encoder_layers,
+ dropout=dropout,
+ layer_dropout=layer_dropout,
+ cnn_module_kernel=cnn_module_kernel,
+ left_context_length=left_context_length // subsampling_factor,
+ right_context_length=right_context_length // subsampling_factor,
+ memory_size=memory_size,
+ tanh_on_mem=tanh_on_mem,
+ negative_inf=negative_inf,
+ )
+
+ def forward(
+ self, x: torch.Tensor, x_lens: torch.Tensor, warmup: float = 1.0
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Forward pass for training and non-streaming inference.
+
+ B: batch size;
+ D: feature dimension;
+ T: length of utterance.
+
+ Args:
+ x (torch.Tensor):
+ Utterance frames right-padded with right context frames,
+ with shape (B, T, D).
+ x_lens (torch.Tensor):
+ With shape (B,) and i-th element representing number of valid
+ utterance frames for i-th batch element in x, containing the
+ right_context at the end.
+ 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:
+ (Tensor, Tensor):
+ - output embedding, with shape (B, T', D), where
+ T' = ((T - 1) // 2 - 1) // 2 - self.right_context_length // 4.
+ - output lengths, with shape (B,), without containing the
+ right_context at the end.
+ """
+ x = self.encoder_embed(x)
+ x = x.permute(1, 0, 2) # (N, T, C) -> (T, N, C)
+
+ x_lens = (((x_lens - 1) >> 1) - 1) >> 1
+ assert x.size(0) == x_lens.max().item()
+
+ output, output_lengths = self.encoder(
+ x, x_lens, warmup=warmup
+ ) # (T, N, C)
+
+ output = output.permute(1, 0, 2) # (T, N, C) -> (N, T, C)
+
+ return output, output_lengths
+
+ @torch.jit.export
+ def infer(
+ self,
+ x: torch.Tensor,
+ x_lens: torch.Tensor,
+ num_processed_frames: torch.Tensor,
+ states: Tuple[List[List[torch.Tensor]], List[torch.Tensor]],
+ ) -> Tuple[
+ torch.Tensor,
+ torch.Tensor,
+ Tuple[List[List[torch.Tensor]], List[torch.Tensor]],
+ ]:
+ """Forward pass for streaming inference.
+
+ B: batch size;
+ D: feature dimension;
+ T: length of utterance.
+
+ Args:
+ x (torch.Tensor):
+ Utterance frames right-padded with right context frames,
+ with shape (B, T, D).
+ lengths (torch.Tensor):
+ With shape (B,) and i-th element representing number of valid
+ utterance frames for i-th batch element in x, containing the
+ right_context at the end.
+ states (List[torch.Tensor, List[List[torch.Tensor]], List[torch.Tensor]]: # noqa
+ Cached states containing:
+ - past_lens: number of past frames for each sample in batch
+ - attn_caches: attention states from preceding chunk's computation,
+ where each element corresponds to each emformer layer
+ - conv_caches: left context for causal convolution, where each
+ element corresponds to each layer.
+ Returns:
+ (Tensor, Tensor):
+ - output embedding, with shape (B, T', D), where
+ T' = ((T - 1) // 2 - 1) // 2 - self.right_context_length // 4.
+ - output lengths, with shape (B,), without containing the
+ right_context at the end.
+ - updated states from current chunk's computation.
+ """
+ x = self.encoder_embed(x)
+ # drop the first and last frames
+ x = x[:, 1:-1, :]
+ x = x.permute(1, 0, 2) # (N, T, C) -> (T, N, C)
+
+ # Caution: We assume the subsampling factor is 4!
+ x_lens = (((x_lens - 1) >> 1) - 1) >> 1
+ x_lens -= 2
+ assert x.size(0) == x_lens.max().item()
+
+ num_processed_frames = num_processed_frames >> 2
+
+ output, output_lengths, output_states = self.encoder.infer(
+ x, x_lens, num_processed_frames, states
+ )
+
+ output = output.permute(1, 0, 2) # (T, N, C) -> (N, T, C)
+
+ return output, output_lengths, output_states
+
+ @torch.jit.export
+ def init_states(self, device: torch.device = torch.device("cpu")):
+ """Create initial states."""
+ return self.encoder.init_states(device)
+
+
+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-1)//2 - 1)//2, which approximates T' == T//4
+
+ 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,
+ ) -> 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-1)//2 - 1)//2, out_channels)
+ layer1_channels:
+ Number of channels in layer1
+ layer1_channels:
+ Number of channels in layer2
+ """
+ assert in_channels >= 7
+ super().__init__()
+
+ self.conv = nn.Sequential(
+ ScaledConv2d(
+ in_channels=1,
+ out_channels=layer1_channels,
+ kernel_size=3,
+ padding=1,
+ ),
+ ActivationBalancer(channel_dim=1),
+ DoubleSwish(),
+ ScaledConv2d(
+ in_channels=layer1_channels,
+ out_channels=layer2_channels,
+ kernel_size=3,
+ stride=2,
+ ),
+ ActivationBalancer(channel_dim=1),
+ DoubleSwish(),
+ ScaledConv2d(
+ in_channels=layer2_channels,
+ out_channels=layer3_channels,
+ kernel_size=3,
+ stride=2,
+ ),
+ ActivationBalancer(channel_dim=1),
+ DoubleSwish(),
+ )
+ self.out = ScaledLinear(
+ layer3_channels * (((in_channels - 1) // 2 - 1) // 2), out_channels
+ )
+ # set learn_eps=False because out_norm is preceded by `out`, and `out`
+ # itself has learned scale, so the extra degree of freedom is not
+ # needed.
+ self.out_norm = BasicNorm(out_channels, learn_eps=False)
+ # constrain median of output to be close to zero.
+ self.out_balancer = ActivationBalancer(
+ channel_dim=-1, min_positive=0.45, max_positive=0.55
+ )
+
+ 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-1)//2 - 1)//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-1)//2 - 1)//2, ((idim-1)//2 - 1)//2)
+ b, c, t, f = x.size()
+ x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f))
+ # Now x is of shape (N, ((T-1)//2 - 1))//2, odim)
+ x = self.out_norm(x)
+ x = self.out_balancer(x)
+ return x
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/encoder_interface.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/encoder_interface.py
new file mode 120000
index 000000000..ee2f09151
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/encoder_interface.py
@@ -0,0 +1 @@
+../conv_emformer_transducer_stateless/encoder_interface.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/export.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/export.py
new file mode 100755
index 000000000..ab15e0241
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/export.py
@@ -0,0 +1,287 @@
+#!/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:
+./conv_emformer_transducer_stateless2/export.py \
+ --exp-dir ./conv_emformer_transducer_stateless2/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 30 \
+ --avg 10 \
+ --use-averaged-model=True \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --jit False
+
+It will generate a file exp_dir/pretrained.pt
+
+To use the generated file with `conv_emformer_transducer_stateless2/decode.py`,
+you can do:
+
+ cd /path/to/exp_dir
+ ln -s pretrained.pt epoch-9999.pt
+
+ cd /path/to/egs/librispeech/ASR
+ ./conv_emformer_transducer_stateless2/decode.py \
+ --exp-dir ./conv_emformer_transducer_stateless2/exp \
+ --epoch 9999 \
+ --avg 1 \
+ --max-duration 100 \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --use-averaged-model=False \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32
+"""
+
+import argparse
+import logging
+from pathlib import Path
+
+import sentencepiece as spm
+import torch
+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=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(
+ "--jit",
+ type=str2bool,
+ default=False,
+ help="""True to save a model after applying torch.jit.script.
+ """,
+ )
+
+ 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 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.eval()
+
+ if params.jit:
+ # 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 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/librispeech/ASR/conv_emformer_transducer_stateless2/joiner.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/joiner.py
new file mode 120000
index 000000000..1eb4dcc83
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/joiner.py
@@ -0,0 +1 @@
+../conv_emformer_transducer_stateless/joiner.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/model.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/model.py
new file mode 120000
index 000000000..322b694e0
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/model.py
@@ -0,0 +1 @@
+../conv_emformer_transducer_stateless/model.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/optim.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/optim.py
new file mode 120000
index 000000000..8f19a99da
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/optim.py
@@ -0,0 +1 @@
+../conv_emformer_transducer_stateless/optim.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/scaling.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/scaling.py
new file mode 120000
index 000000000..12f22cf9c
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/scaling.py
@@ -0,0 +1 @@
+../conv_emformer_transducer_stateless/scaling.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/stream.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/stream.py
new file mode 120000
index 000000000..bf9cbbe2e
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/stream.py
@@ -0,0 +1 @@
+../conv_emformer_transducer_stateless/stream.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/streaming_decode.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/streaming_decode.py
new file mode 100755
index 000000000..71150392d
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/streaming_decode.py
@@ -0,0 +1,983 @@
+#!/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
+./conv_emformer_transducer_stateless2/streaming_decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --num-decode-streams 2000 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method greedy_search \
+ --use-averaged-model True
+
+(2) modified beam search
+./conv_emformer_transducer_stateless2/streaming_decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --num-decode-streams 2000 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method modified_beam_search \
+ --use-averaged-model True \
+ --beam-size 4
+
+(3) fast beam search
+./conv_emformer_transducer_stateless2/streaming_decode.py \
+ --epoch 30 \
+ --avg 10 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --num-decode-streams 2000 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32 \
+ --decoding-method fast_beam_search \
+ --use-averaged-model True \
+ --beam 4 \
+ --max-contexts 4 \
+ --max-states 8
+"""
+import argparse
+import logging
+import warnings
+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 beam_search import Hypothesis, HypothesisList, get_hyps_shape
+from emformer import LOG_EPSILON, stack_states, unstack_states
+from kaldifeat import Fbank, FbankOptions
+from lhotse import CutSet
+from stream import Stream
+from torch.nn.utils.rnn import pad_sequence
+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 one_best_decoding
+from icefall.utils import (
+ AttributeDict,
+ get_texts,
+ 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=28,
+ help="It specifies the checkpoint to use for decoding."
+ "Note: Epoch counts from 0.",
+ )
+
+ 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'. ",
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=False,
+ 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="transducer_emformer/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="""Possible values are:
+ - greedy_search
+ - modified_beam_search
+ - fast_beam_search
+ """,
+ )
+
+ 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(
+ "--max-contexts",
+ type=int,
+ default=4,
+ help="""Used only when --decoding-method is
+ fast_beam_search""",
+ )
+
+ parser.add_argument(
+ "--max-states",
+ type=int,
+ default=8,
+ 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(
+ "--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(
+ "--sampling-rate",
+ type=float,
+ default=16000,
+ help="Sample rate of the audio",
+ )
+
+ 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 greedy_search(
+ model: nn.Module,
+ encoder_out: torch.Tensor,
+ streams: List[Stream],
+) -> None:
+ """Greedy search in batch mode. It hardcodes --max-sym-per-frame=1.
+
+ Args:
+ model:
+ The transducer model.
+ encoder_out:
+ Output from the encoder. Its shape is (N, T, C), where N >= 1.
+ streams:
+ A list of Stream objects.
+ """
+ assert len(streams) == encoder_out.size(0)
+ assert encoder_out.ndim == 3
+
+ blank_id = model.decoder.blank_id
+ context_size = model.decoder.context_size
+ device = next(model.parameters()).device
+ T = encoder_out.size(1)
+
+ encoder_out = model.joiner.encoder_proj(encoder_out)
+
+ decoder_input = torch.tensor(
+ [stream.hyp[-context_size:] for stream in streams],
+ device=device,
+ dtype=torch.int64,
+ )
+ # decoder_out is of shape (batch_size, 1, decoder_out_dim)
+ decoder_out = model.decoder(decoder_input, need_pad=False)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+
+ for t in range(T):
+ # current_encoder_out's shape: (batch_size, 1, encoder_out_dim)
+ current_encoder_out = encoder_out[:, t : t + 1, :] # noqa
+
+ logits = model.joiner(
+ current_encoder_out.unsqueeze(2),
+ decoder_out.unsqueeze(1),
+ project_input=False,
+ )
+ # logits'shape (batch_size, vocab_size)
+ logits = logits.squeeze(1).squeeze(1)
+
+ assert logits.ndim == 2, logits.shape
+ y = logits.argmax(dim=1).tolist()
+ emitted = False
+ for i, v in enumerate(y):
+ if v != blank_id:
+ streams[i].hyp.append(v)
+ emitted = True
+ if emitted:
+ # update decoder output
+ decoder_input = torch.tensor(
+ [stream.hyp[-context_size:] for stream in streams],
+ device=device,
+ dtype=torch.int64,
+ )
+ decoder_out = model.decoder(
+ decoder_input,
+ need_pad=False,
+ )
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+
+
+def modified_beam_search(
+ model: nn.Module,
+ encoder_out: torch.Tensor,
+ streams: List[Stream],
+ beam: int = 4,
+):
+ """Beam search in batch mode with --max-sym-per-frame=1 being hardcoded.
+
+ Args:
+ model:
+ The RNN-T model.
+ encoder_out:
+ A 3-D tensor of shape (N, T, encoder_out_dim) containing the output of
+ the encoder model.
+ streams:
+ A list of stream objects.
+ beam:
+ Number of active paths during the beam search.
+ """
+ assert encoder_out.ndim == 3, encoder_out.shape
+ assert len(streams) == encoder_out.size(0)
+
+ blank_id = model.decoder.blank_id
+ context_size = model.decoder.context_size
+ device = next(model.parameters()).device
+ batch_size = len(streams)
+ T = encoder_out.size(1)
+
+ B = [stream.hyps for stream in streams]
+
+ encoder_out = model.joiner.encoder_proj(encoder_out)
+
+ for t in range(T):
+ current_encoder_out = encoder_out[:, t].unsqueeze(1).unsqueeze(1)
+ # current_encoder_out's shape: (batch_size, 1, 1, encoder_out_dim)
+
+ 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.stack(
+ [hyp.log_prob.reshape(1) for hyps in A for hyp in hyps], dim=0
+ ) # (num_hyps, 1)
+
+ 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)
+ # decoder_out is of shape (num_hyps, 1, 1, decoder_output_dim)
+
+ # Note: For torch 1.7.1 and below, it requires a torch.int64 tensor
+ # as index, so we use `to(torch.int64)` below.
+ current_encoder_out = torch.index_select(
+ current_encoder_out,
+ dim=0,
+ index=hyps_shape.row_ids(1).to(torch.int64),
+ ) # (num_hyps, encoder_out_dim)
+
+ logits = model.joiner(
+ current_encoder_out, decoder_out, project_input=False
+ )
+ # logits is of shape (num_hyps, 1, 1, vocab_size)
+
+ logits = logits.squeeze(1).squeeze(1)
+
+ 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 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_ys = hyp.ys[:]
+ new_token = topk_token_indexes[k]
+ if new_token != blank_id:
+ new_ys.append(new_token)
+
+ new_log_prob = topk_log_probs[k]
+ new_hyp = Hypothesis(ys=new_ys, log_prob=new_log_prob)
+ B[i].add(new_hyp)
+
+ for i in range(batch_size):
+ streams[i].hyps = B[i]
+
+
+def fast_beam_search_one_best(
+ model: nn.Module,
+ streams: List[Stream],
+ encoder_out: torch.Tensor,
+ processed_lens: torch.Tensor,
+ beam: float,
+ max_states: int,
+ max_contexts: int,
+) -> None:
+ """It limits the maximum number of symbols per frame to 1.
+
+ A lattice is first obtained using modified beam search, and then
+ the shortest path within the lattice is used as the final output.
+
+ Args:
+ model:
+ An instance of `Transducer`.
+ streams:
+ A list of stream objects.
+ encoder_out:
+ A tensor of shape (N, T, C) from the encoder.
+ processed_lens:
+ A tensor of shape (N,) containing the number of processed frames
+ in `encoder_out` before padding.
+ beam:
+ Beam value, similar to the beam used in Kaldi..
+ max_states:
+ Max states per stream per frame.
+ max_contexts:
+ Max contexts pre stream per frame.
+ """
+ assert encoder_out.ndim == 3
+
+ context_size = model.decoder.context_size
+ vocab_size = model.decoder.vocab_size
+
+ B, T, C = encoder_out.shape
+ assert B == len(streams)
+
+ config = k2.RnntDecodingConfig(
+ vocab_size=vocab_size,
+ decoder_history_len=context_size,
+ beam=beam,
+ max_contexts=max_contexts,
+ max_states=max_states,
+ )
+ individual_streams = []
+ for i in range(B):
+ individual_streams.append(streams[i].rnnt_decoding_stream)
+ decoding_streams = k2.RnntDecodingStreams(individual_streams, config)
+
+ encoder_out = model.joiner.encoder_proj(encoder_out)
+
+ for t in range(T):
+ # shape is a RaggedShape of shape (B, context)
+ # contexts is a Tensor of shape (shape.NumElements(), context_size)
+ shape, contexts = decoding_streams.get_contexts()
+ # `nn.Embedding()` in torch below v1.7.1 supports only torch.int64
+ contexts = contexts.to(torch.int64)
+ # decoder_out is of shape (shape.NumElements(), 1, decoder_out_dim)
+ decoder_out = model.decoder(contexts, need_pad=False)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+ # current_encoder_out is of shape
+ # (shape.NumElements(), 1, joiner_dim)
+ # fmt: off
+ current_encoder_out = torch.index_select(
+ encoder_out[:, t:t + 1, :], 0, shape.row_ids(1).to(torch.int64)
+ )
+ # fmt: on
+ logits = model.joiner(
+ current_encoder_out.unsqueeze(2),
+ decoder_out.unsqueeze(1),
+ project_input=False,
+ )
+ logits = logits.squeeze(1).squeeze(1)
+ log_probs = logits.log_softmax(dim=-1)
+ decoding_streams.advance(log_probs)
+
+ decoding_streams.terminate_and_flush_to_streams()
+
+ lattice = decoding_streams.format_output(processed_lens.tolist())
+
+ best_path = one_best_decoding(lattice)
+ hyps = get_texts(best_path)
+
+ for i in range(B):
+ streams[i].hyp = hyps[i]
+
+
+def decode_one_chunk(
+ model: nn.Module,
+ streams: List[Stream],
+ params: AttributeDict,
+ decoding_graph: Optional[k2.Fsa] = None,
+) -> List[int]:
+ """
+ Args:
+ model:
+ The Transducer model.
+ streams:
+ A list of Stream objects.
+ params:
+ It is returned by :func:`get_params`.
+ decoding_graph:
+ The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+ only when --decoding_method is fast_beam_search.
+
+ Returns:
+ A list of indexes indicating the finished streams.
+ """
+ device = next(model.parameters()).device
+
+ feature_list = []
+ feature_len_list = []
+ state_list = []
+ num_processed_frames_list = []
+
+ for stream in streams:
+ # We should first get `stream.num_processed_frames`
+ # before calling `stream.get_feature_chunk()`
+ # since `stream.num_processed_frames` would be updated
+ num_processed_frames_list.append(stream.num_processed_frames)
+ feature = stream.get_feature_chunk()
+ feature_len = feature.size(0)
+ feature_list.append(feature)
+ feature_len_list.append(feature_len)
+ state_list.append(stream.states)
+
+ features = pad_sequence(
+ feature_list, batch_first=True, padding_value=LOG_EPSILON
+ ).to(device)
+ feature_lens = torch.tensor(feature_len_list, device=device)
+ num_processed_frames = torch.tensor(
+ num_processed_frames_list, device=device
+ )
+
+ # Make sure it has at least 1 frame after subsampling, first-and-last-frame cutting, and right context cutting # noqa
+ tail_length = (
+ 3 * params.subsampling_factor + params.right_context_length + 3
+ )
+ 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_EPSILON,
+ )
+
+ # Stack states of all streams
+ states = stack_states(state_list)
+
+ encoder_out, encoder_out_lens, states = model.encoder.infer(
+ x=features,
+ x_lens=feature_lens,
+ states=states,
+ num_processed_frames=num_processed_frames,
+ )
+
+ if params.decoding_method == "greedy_search":
+ greedy_search(
+ model=model,
+ streams=streams,
+ encoder_out=encoder_out,
+ )
+ elif params.decoding_method == "modified_beam_search":
+ modified_beam_search(
+ model=model,
+ streams=streams,
+ encoder_out=encoder_out,
+ beam=params.beam_size,
+ )
+ elif params.decoding_method == "fast_beam_search":
+ # feature_len is needed to get partial results.
+ # The rnnt_decoding_stream for fast_beam_search.
+ fast_beam_search_one_best(
+ model=model,
+ streams=streams,
+ encoder_out=encoder_out,
+ processed_lens=(num_processed_frames >> 2) + encoder_out_lens,
+ beam=params.beam,
+ max_contexts=params.max_contexts,
+ max_states=params.max_states,
+ )
+ else:
+ raise ValueError(
+ f"Unsupported decoding method: {params.decoding_method}"
+ )
+
+ # Update cached states of each stream
+ state_list = unstack_states(states)
+ for i, s in enumerate(state_list):
+ streams[i].states = s
+
+ finished_streams = [i for i, stream in enumerate(streams) if stream.done]
+ return finished_streams
+
+
+def create_streaming_feature_extractor() -> Fbank:
+ """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 = 16000
+ opts.mel_opts.num_bins = 80
+ return Fbank(opts)
+
+
+def decode_dataset(
+ cuts: CutSet,
+ model: nn.Module,
+ params: AttributeDict,
+ sp: spm.SentencePieceProcessor,
+ decoding_graph: Optional[k2.Fsa] = None,
+):
+ """Decode dataset.
+
+ Args:
+ cuts:
+ Lhotse Cutset containing the dataset to decode.
+ params:
+ It is returned by :func:`get_params`.
+ model:
+ The Transducer 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 = next(model.parameters()).device
+
+ log_interval = 300
+
+ fbank = create_streaming_feature_extractor()
+
+ decode_results = []
+ streams = []
+ for num, cut in enumerate(cuts):
+ # Each utterance has a Stream.
+ stream = Stream(
+ params=params,
+ cut_id=cut.id,
+ decoding_graph=decoding_graph,
+ device=device,
+ LOG_EPS=LOG_EPSILON,
+ )
+
+ stream.set_states(model.encoder.init_states(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)
+ feature = fbank(samples)
+ stream.set_feature(feature)
+ stream.set_ground_truth(cut.supervisions[0].text)
+
+ streams.append(stream)
+
+ while len(streams) >= params.num_decode_streams:
+ finished_streams = decode_one_chunk(
+ model=model,
+ streams=streams,
+ params=params,
+ decoding_graph=decoding_graph,
+ )
+
+ for i in sorted(finished_streams, reverse=True):
+ decode_results.append(
+ (
+ streams[i].id,
+ streams[i].ground_truth.split(),
+ sp.decode(streams[i].decoding_result()).split(),
+ )
+ )
+ del streams[i]
+
+ if num % log_interval == 0:
+ logging.info(f"Cuts processed until now is {num}.")
+
+ while len(streams) > 0:
+ finished_streams = decode_one_chunk(
+ model=model,
+ streams=streams,
+ params=params,
+ decoding_graph=decoding_graph,
+ )
+
+ for i in sorted(finished_streams, reverse=True):
+ decode_results.append(
+ (
+ streams[i].id,
+ streams[i].ground_truth.split(),
+ sp.decode(streams[i].decoding_result()).split(),
+ )
+ )
+ del 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}"
+ )
+ else:
+ key = f"beam_size_{params.beam_size}"
+
+ 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"
+ )
+ store_transcripts(filename=recog_path, texts=sorted(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",
+ "fast_beam_search",
+ "modified_beam_search",
+ )
+ 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-length-{params.chunk_length}"
+ params.suffix += f"-left-context-length-{params.left_context_length}"
+ params.suffix += f"-right-context-length-{params.right_context_length}"
+ params.suffix += f"-memory-size-{params.memory_size}"
+
+ 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}"
+ 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-streaming-decode")
+ 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()
+
+ params.device = device
+
+ 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.eval()
+
+ if params.decoding_method == "fast_beam_search":
+ 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}")
+
+ 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,
+ model=model,
+ params=params,
+ 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__":
+ torch.manual_seed(20220410)
+ main()
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/test_emformer.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/test_emformer.py
new file mode 100644
index 000000000..8cde6205b
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/test_emformer.py
@@ -0,0 +1,194 @@
+#!/usr/bin/env python3
+#
+# Copyright 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.
+
+
+import torch
+from emformer import ConvolutionModule, Emformer, stack_states, unstack_states
+
+
+def test_convolution_module_forward():
+ B, D = 2, 256
+ chunk_length = 4
+ right_context_length = 2
+ num_chunks = 3
+ U = num_chunks * chunk_length
+ R = num_chunks * right_context_length
+ kernel_size = 31
+ conv_module = ConvolutionModule(
+ chunk_length,
+ right_context_length,
+ D,
+ kernel_size,
+ )
+
+ utterance = torch.randn(U, B, D)
+ right_context = torch.randn(R, B, D)
+
+ utterance, right_context = conv_module(utterance, right_context)
+ assert utterance.shape == (U, B, D), utterance.shape
+ assert right_context.shape == (R, B, D), right_context.shape
+
+
+def test_convolution_module_infer():
+ from emformer import ConvolutionModule
+
+ B, D = 2, 256
+ chunk_length = 4
+ right_context_length = 2
+ num_chunks = 1
+ U = num_chunks * chunk_length
+ R = num_chunks * right_context_length
+ kernel_size = 31
+ conv_module = ConvolutionModule(
+ chunk_length,
+ right_context_length,
+ D,
+ kernel_size,
+ )
+
+ utterance = torch.randn(U, B, D)
+ right_context = torch.randn(R, B, D)
+ cache = torch.randn(B, D, kernel_size - 1)
+
+ utterance, right_context, new_cache = conv_module.infer(
+ utterance, right_context, cache
+ )
+ assert utterance.shape == (U, B, D), utterance.shape
+ assert right_context.shape == (R, B, D), right_context.shape
+ assert new_cache.shape == (B, D, kernel_size - 1), new_cache.shape
+
+
+def test_state_stack_unstack():
+ num_features = 80
+ chunk_length = 32
+ encoder_dim = 512
+ num_encoder_layers = 2
+ kernel_size = 31
+ left_context_length = 32
+ right_context_length = 8
+ memory_size = 32
+
+ model = Emformer(
+ num_features=num_features,
+ chunk_length=chunk_length,
+ subsampling_factor=4,
+ d_model=encoder_dim,
+ num_encoder_layers=num_encoder_layers,
+ cnn_module_kernel=kernel_size,
+ left_context_length=left_context_length,
+ right_context_length=right_context_length,
+ memory_size=memory_size,
+ )
+
+ for batch_size in [1, 2]:
+ attn_caches = [
+ [
+ torch.zeros(memory_size, batch_size, encoder_dim),
+ torch.zeros(left_context_length // 4, batch_size, encoder_dim),
+ torch.zeros(
+ left_context_length // 4,
+ batch_size,
+ encoder_dim,
+ ),
+ ]
+ for _ in range(num_encoder_layers)
+ ]
+ conv_caches = [
+ torch.zeros(batch_size, encoder_dim, kernel_size - 1)
+ for _ in range(num_encoder_layers)
+ ]
+ states = [attn_caches, conv_caches]
+ x = torch.randn(batch_size, 23, num_features)
+ x_lens = torch.full((batch_size,), 23)
+ num_processed_frames = torch.full((batch_size,), 0)
+ y, y_lens, states = model.infer(
+ x, x_lens, num_processed_frames=num_processed_frames, states=states
+ )
+
+ state_list = unstack_states(states)
+ states2 = stack_states(state_list)
+
+ for ss, ss2 in zip(states[0], states2[0]):
+ for s, s2 in zip(ss, ss2):
+ assert torch.allclose(s, s2), f"{s.sum()}, {s2.sum()}"
+
+ for s, s2 in zip(states[1], states2[1]):
+ assert torch.allclose(s, s2), f"{s.sum()}, {s2.sum()}"
+
+
+def test_torchscript_consistency_infer():
+ r"""Verify that scripting Emformer does not change the behavior of method `infer`.""" # noqa
+ num_features = 80
+ chunk_length = 32
+ encoder_dim = 512
+ num_encoder_layers = 2
+ kernel_size = 31
+ left_context_length = 32
+ right_context_length = 8
+ memory_size = 32
+ batch_size = 2
+
+ model = Emformer(
+ num_features=num_features,
+ chunk_length=chunk_length,
+ subsampling_factor=4,
+ d_model=encoder_dim,
+ num_encoder_layers=num_encoder_layers,
+ cnn_module_kernel=kernel_size,
+ left_context_length=left_context_length,
+ right_context_length=right_context_length,
+ memory_size=memory_size,
+ ).eval()
+ attn_caches = [
+ [
+ torch.zeros(memory_size, batch_size, encoder_dim),
+ torch.zeros(left_context_length // 4, batch_size, encoder_dim),
+ torch.zeros(
+ left_context_length // 4,
+ batch_size,
+ encoder_dim,
+ ),
+ ]
+ for _ in range(num_encoder_layers)
+ ]
+ conv_caches = [
+ torch.zeros(batch_size, encoder_dim, kernel_size - 1)
+ for _ in range(num_encoder_layers)
+ ]
+ states = [attn_caches, conv_caches]
+ x = torch.randn(batch_size, 23, num_features)
+ x_lens = torch.full((batch_size,), 23)
+ num_processed_frames = torch.full((batch_size,), 0)
+ y, y_lens, out_states = model.infer(
+ x, x_lens, num_processed_frames=num_processed_frames, states=states
+ )
+
+ sc_model = torch.jit.script(model).eval()
+ sc_y, sc_y_lens, sc_out_states = sc_model.infer(
+ x, x_lens, num_processed_frames=num_processed_frames, states=states
+ )
+
+ assert torch.allclose(y, sc_y)
+
+
+if __name__ == "__main__":
+ test_convolution_module_forward()
+ test_convolution_module_infer()
+ test_state_stack_unstack()
+ test_torchscript_consistency_infer()
diff --git a/egs/librispeech/ASR/conv_emformer_transducer_stateless2/train.py b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/train.py
new file mode 100755
index 000000000..2bbc45d78
--- /dev/null
+++ b/egs/librispeech/ASR/conv_emformer_transducer_stateless2/train.py
@@ -0,0 +1,1145 @@
+#!/usr/bin/env python3
+# Copyright 2021 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"
+
+./conv_emformer_transducer_stateless2/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --full-libri 1 \
+ --max-duration 280 \
+ --master-port 12321 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32
+
+# For mix precision training:
+./conv_emformer_transducer_stateless2/train.py \
+ --world-size 4 \
+ --num-epochs 30 \
+ --start-epoch 1 \
+ --use-fp16 1 \
+ --exp-dir conv_emformer_transducer_stateless2/exp \
+ --full-libri 1 \
+ --max-duration 300 \
+ --master-port 12321 \
+ --num-encoder-layers 12 \
+ --chunk-length 32 \
+ --cnn-module-kernel 31 \
+ --left-context-length 32 \
+ --right-context-length 8 \
+ --memory-size 32
+"""
+
+
+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 emformer import Emformer
+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, Eve
+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.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.utils import AttributeDict, MetricsTracker, setup_logger, str2bool
+
+LRSchedulerType = Union[
+ torch.optim.lr_scheduler._LRScheduler, optim.LRScheduler
+]
+
+
+def add_model_arguments(parser: argparse.ArgumentParser):
+ parser.add_argument(
+ "--encoder-dim",
+ type=int,
+ default=512,
+ help="Attention dim for the Emformer",
+ )
+
+ parser.add_argument(
+ "--nhead",
+ type=int,
+ default=8,
+ help="Number of attention heads for the Emformer",
+ )
+
+ parser.add_argument(
+ "--dim-feedforward",
+ type=int,
+ default=2048,
+ help="Feed-forward dimension for the Emformer",
+ )
+
+ parser.add_argument(
+ "--num-encoder-layers",
+ type=int,
+ default=12,
+ help="Number of encoder layers for the Emformer",
+ )
+
+ parser.add_argument(
+ "--cnn-module-kernel",
+ type=int,
+ default=31,
+ help="Kernel size for the convolution module.",
+ )
+
+ parser.add_argument(
+ "--left-context-length",
+ type=int,
+ default=32,
+ help="""Number of frames before subsampling for left context
+ in the Emformer.""",
+ )
+
+ parser.add_argument(
+ "--chunk-length",
+ type=int,
+ default=32,
+ help="""Number of frames before subsampling for each chunk
+ in the Emformer.""",
+ )
+
+ parser.add_argument(
+ "--right-context-length",
+ type=int,
+ default=8,
+ help="""Number of frames before subsampling for right context
+ in the Emformer.""",
+ )
+
+ parser.add_argument(
+ "--memory-size",
+ type=int,
+ default=0,
+ help="Number of entries in the memory for the Emformer",
+ )
+
+
+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_stateless2/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(
+ "--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(
+ "--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(
+ "--save-every-n",
+ type=int,
+ default=8000,
+ 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=20,
+ 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.
+
+ - encoder_dim: Hidden dim for multi-head attention model.
+
+ - num_decoder_layers: Number of decoder layer of transformer decoder.
+
+ - 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": 50,
+ "reset_interval": 200,
+ "valid_interval": 3000, # For the 100h subset, use 800
+ # parameters for Emformer
+ "feature_dim": 80,
+ "subsampling_factor": 4,
+ # parameters for decoder
+ "decoder_dim": 512,
+ # parameters for joiner
+ "joiner_dim": 512,
+ # parameters for Noam
+ "model_warm_step": 3000, # arg given to model, not for lrate
+ "env_info": get_env_info(),
+ }
+ )
+
+ return params
+
+
+def get_encoder_model(params: AttributeDict) -> nn.Module:
+ # TODO: We can add an option to switch between Conformer and Transformer
+ encoder = Emformer(
+ num_features=params.feature_dim,
+ chunk_length=params.chunk_length,
+ subsampling_factor=params.subsampling_factor,
+ d_model=params.encoder_dim,
+ nhead=params.nhead,
+ dim_feedforward=params.dim_feedforward,
+ num_encoder_layers=params.num_encoder_layers,
+ cnn_module_kernel=params.cnn_module_kernel,
+ left_context_length=params.left_context_length,
+ right_context_length=params.right_context_length,
+ memory_size=params.memory_size,
+ )
+ 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=params.encoder_dim,
+ 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=params.encoder_dim,
+ 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,
+ warmup: float = 1.0,
+) -> Tuple[Tensor, MetricsTracker]:
+ """
+ Compute RNN-T 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.
+ 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)
+
+ 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,
+ warmup=warmup,
+ )
+ # after the main warmup step, we keep pruned_loss_scale small
+ # for the same amount of time (model_warm_step), to avoid
+ # overwhelming the simple_loss and causing it to diverge,
+ # in case it had not fully learned the alignment yet.
+ pruned_loss_scale = (
+ 0.0
+ if warmup < 1.0
+ else (0.1 if warmup > 1.0 and warmup < 2.0 else 1.0)
+ )
+ loss = (
+ params.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()
+ )
+
+ # `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["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"])
+
+ 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,
+ 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()
+
+ 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 % 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,
+ sp=sp,
+ 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))
+ 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)
+
+ 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 = Eve(model.parameters(), lr=params.initial_lr)
+
+ 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)
+
+ 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
+ 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,
+ sp=sp,
+ params=params,
+ )
+
+ 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)
+
+ 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 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:
+ # warmup = 0.0 is so that the derivs for the pruned loss stay zero
+ # (i.e. are not remembered by the decaying-average in adam), because
+ # we want to avoid these params being subject to shrinkage in adam.
+ with torch.cuda.amp.autocast(enabled=params.use_fp16):
+ loss, _ = compute_loss(
+ params=params,
+ model=model,
+ sp=sp,
+ batch=batch,
+ is_training=True,
+ warmup=0.0,
+ )
+ loss.backward()
+ optimizer.step()
+ optimizer.zero_grad()
+ except RuntimeError 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]}) ..."
+ )
+ raise
+
+
+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/distillation_with_hubert.sh b/egs/librispeech/ASR/distillation_with_hubert.sh
old mode 100644
new mode 100755
index e18ba8f55..2a69d3921
--- a/egs/librispeech/ASR/distillation_with_hubert.sh
+++ b/egs/librispeech/ASR/distillation_with_hubert.sh
@@ -1,3 +1,5 @@
+#!/usr/bin/env bash
+#
# A short introduction about distillation framework.
#
# A typical traditional distillation method is
@@ -14,15 +16,15 @@
# teacher embeddings.
# 3. a middle layer 6(1-based) out of total 6 layers is used to extract
# student embeddings.
-
-# This is an example to do distillation with librispeech clean-100 subset.
-# run with command:
-# bash distillation_with_hubert.sh [0|1|2|3|4]
#
-# For example command
-# bash distillation_with_hubert.sh 0
-# will download hubert model.
-stage=$1
+# To directly download the extracted codebook indexes for model distillation, you can
+# set stage=2, stop_stage=4, use_extracted_codebook=True
+#
+# To start from scratch, you can
+# set stage=0, stop_stage=4, use_extracted_codebook=False
+
+stage=0
+stop_stage=4
# Set the GPUs available.
# This script requires at least one GPU.
@@ -33,10 +35,35 @@ stage=$1
# export CUDA_VISIBLE_DEVICES="0"
#
# Suppose GPU 2,3,4,5 are available.
-export CUDA_VISIBLE_DEVICES="2,3,4,5"
+export CUDA_VISIBLE_DEVICES="0,1,2,3"
+exp_dir=./pruned_transducer_stateless6/exp
+mkdir -p $exp_dir
-if [ $stage -eq 0 ]; then
+# 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
+
+# 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
+
+# teacher_model_id can be one of
+# "hubert_xtralarge_ll60k_finetune_ls960" -> fine-tuned model, it is the one we currently use.
+# "hubert_xtralarge_ll60k" -> pretrained model without fintuing
+teacher_model_id=hubert_xtralarge_ll60k_finetune_ls960
+
+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 0 ] && [ $stop_stage -ge 0 ] && [ ! "$use_extracted_codebook" == "True" ]; then
+ log "Stage 0: Download HuBERT model"
# Preparation stage.
# Install fairseq according to:
@@ -45,53 +72,52 @@ if [ $stage -eq 0 ]; then
# commit 806855bf660ea748ed7ffb42fe8dcc881ca3aca0 is used.
has_fairseq=$(python3 -c "import importlib; print(importlib.util.find_spec('fairseq') is not None)")
if [ $has_fairseq == 'False' ]; then
- echo "Please install fairseq before running following stages"
+ log "Please install fairseq before running following stages"
exit 1
fi
# Install quantization toolkit:
- # pip install git+https://github.com/danpovey/quantization.git@master
- # when testing this code:
- # commit c17ffe67aa2e6ca6b6855c50fde812f2eed7870b is used.
+ # pip install git+https://github.com/k2-fsa/multi_quantization.git
+ # or
+ # pip install multi_quantization
- has_quantization=$(python3 -c "import importlib; print(importlib.util.find_spec('quantization') is not None)")
+ has_quantization=$(python3 -c "import importlib; print(importlib.util.find_spec('multi_quantization') is not None)")
if [ $has_quantization == 'False' ]; then
- echo "Please install quantization before running following stages"
+ log "Please install multi_quantization before running following stages"
exit 1
fi
- echo "Download hubert model."
+ log "Download HuBERT model."
# Parameters about model.
- exp_dir=./pruned_transducer_stateless6/exp/
- model_id=hubert_xtralarge_ll60k_finetune_ls960
hubert_model_dir=${exp_dir}/hubert_models
- hubert_model=${hubert_model_dir}/${model_id}.pt
+ hubert_model=${hubert_model_dir}/${teacher_model_id}.pt
mkdir -p ${hubert_model_dir}
# For more models refer to: https://github.com/pytorch/fairseq/tree/main/examples/hubert
if [ -f ${hubert_model} ]; then
- echo "hubert model alread exists."
+ log "HuBERT model alread exists."
else
- wget -c https://dl.fbaipublicfiles.com/hubert/${model_id} -P ${hubert_model}
+ wget -c https://dl.fbaipublicfiles.com/hubert/${teacher_model_id}.pt -P ${hubert_model_dir}
wget -c wget https://dl.fbaipublicfiles.com/fairseq/wav2vec/dict.ltr.txt -P ${hubert_model_dir}
fi
fi
if [ ! -d ./data/fbank ]; then
- echo "This script assumes ./data/fbank is already generated by prepare.sh"
+ log "This script assumes ./data/fbank is already generated by prepare.sh"
exit 1
fi
-if [ $stage -eq 1 ]; then
+if [ $stage -le 1 ] && [ $stop_stage -ge 1 ] && [ ! "$use_extracted_codebook" == "True" ]; then
+ log "Stage 1: Verify that the downloaded HuBERT model is correct."
# This stage is not directly used by codebook indexes extraction.
# It is a method to "prove" that the downloaed hubert model
# is inferenced in an correct way if WERs look like normal.
# Expect WERs:
# [test-clean-ctc_greedy_search] %WER 2.04% [1075 / 52576, 92 ins, 104 del, 879 sub ]
# [test-other-ctc_greedy_search] %WER 3.71% [1942 / 52343, 152 ins, 126 del, 1664 sub ]
- ./pruned_transducer_stateless6/hubert_decode.py
+ ./pruned_transducer_stateless6/hubert_decode.py --exp-dir $exp_dir
fi
-if [ $stage -eq 2 ]; then
+if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
# Analysis of disk usage:
# With num_codebooks==8, each teacher embedding is quantized into
# a sequence of eight 8-bit integers, i.e. only eight bytes are needed.
@@ -113,25 +139,61 @@ if [ $stage -eq 2 ]; then
# During quantizer's training data(teacher embedding) and it's training,
# only the first ONE GPU is used.
# During codebook indexes extraction, ALL GPUs set by CUDA_VISIBLE_DEVICES are used.
+
+ if [ "$use_extracted_codebook" == "True" ]; then
+ if [ ! "$teacher_model_id" == "hubert_xtralarge_ll60k_finetune_ls960" ]; then
+ log "Currently we only uploaded codebook indexes from teacher model hubert_xtralarge_ll60k_finetune_ls960"
+ exit 1
+ fi
+ mkdir -p $exp_dir/vq
+ 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
+ log "$codebook_download_dir exists, you should remove it first."
+ exit 1
+ fi
+ log "Downloading extracted codebook indexes to $codebook_download_dir"
+ # Make sure you have git-lfs installed (https://git-lfs.github.com)
+ git lfs install
+ git clone https://huggingface.co/Zengwei/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/
+ mkdir -p $codebook_dir/splits4
+ mv $codebook_download_dir/*.h5 $codebook_dir/splits4/
+ log "Remove $codebook_download_dir"
+ rm -rf $codebook_download_dir
+ fi
+
./pruned_transducer_stateless6/extract_codebook_index.py \
- --full-libri False
+ --full-libri $full_libri \
+ --exp-dir $exp_dir \
+ --embedding-layer 36 \
+ --num-utts 1000 \
+ --num-codebooks 8 \
+ --max-duration 100 \
+ --teacher-model-id $teacher_model_id \
+ --use-extracted-codebook $use_extracted_codebook
fi
-if [ $stage -eq 3 ]; then
+if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
# Example training script.
# Note: it's better to set spec-aug-time-warpi-factor=-1
WORLD_SIZE=$(echo ${CUDA_VISIBLE_DEVICES} | awk '{n=split($1, _, ","); print n}')
./pruned_transducer_stateless6/train.py \
--manifest-dir ./data/vq_fbank \
--master-port 12359 \
- --full-libri False \
+ --full-libri $full_libri \
--spec-aug-time-warp-factor -1 \
--max-duration 300 \
--world-size ${WORLD_SIZE} \
- --num-epochs 20
+ --num-epochs 20 \
+ --exp-dir $exp_dir \
+ --enable-distillation True
fi
-if [ $stage -eq 4 ]; then
+if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
# Results should be similar to:
# errs-test-clean-beam_size_4-epoch-20-avg-10-beam-4.txt:%WER = 5.67
# errs-test-other-beam_size_4-epoch-20-avg-10-beam-4.txt:%WER = 15.60
@@ -140,5 +202,6 @@ if [ $stage -eq 4 ]; then
--epoch 20 \
--avg 10 \
--max-duration 200 \
- --exp-dir ./pruned_transducer_stateless6/exp
+ --exp-dir $exp_dir \
+ --enable-distillation True
fi
diff --git a/egs/librispeech/ASR/generate-lm.sh b/egs/librispeech/ASR/generate-lm.sh
new file mode 100755
index 000000000..6baccd381
--- /dev/null
+++ b/egs/librispeech/ASR/generate-lm.sh
@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+
+lang_dir=data/lang_bpe_500
+
+for ngram in 2 3 5; do
+ if [ ! -f $lang_dir/${ngram}gram.arpa ]; then
+ ./shared/make_kn_lm.py \
+ -ngram-order ${ngram} \
+ -text $lang_dir/transcript_tokens.txt \
+ -lm $lang_dir/${ngram}gram.arpa
+ fi
+
+ if [ ! -f $lang_dir/${ngram}gram.fst.txt ]; then
+ python3 -m kaldilm \
+ --read-symbol-table="$lang_dir/tokens.txt" \
+ --disambig-symbol='#0' \
+ --max-order=${ngram} \
+ $lang_dir/${ngram}gram.arpa > $lang_dir/${ngram}gram.fst.txt
+ fi
+done
diff --git a/egs/librispeech/ASR/local/compute_fbank_gigaspeech_dev_test.py b/egs/librispeech/ASR/local/compute_fbank_gigaspeech_dev_test.py
index 9f1039893..c0c7ef8c5 100644
--- a/egs/librispeech/ASR/local/compute_fbank_gigaspeech_dev_test.py
+++ b/egs/librispeech/ASR/local/compute_fbank_gigaspeech_dev_test.py
@@ -20,11 +20,7 @@ import logging
from pathlib import Path
import torch
-from lhotse import (
- CutSet,
- KaldifeatFbank,
- KaldifeatFbankConfig,
-)
+from lhotse import CutSet, KaldifeatFbank, KaldifeatFbankConfig
# Torch's multithreaded behavior needs to be disabled or
# it wastes a lot of CPU and slow things down.
@@ -51,13 +47,16 @@ def compute_fbank_gigaspeech_dev_test():
logging.info(f"device: {device}")
+ prefix = "gigaspeech"
+ suffix = "jsonl.gz"
+
for partition in subsets:
- cuts_path = in_out_dir / f"cuts_{partition}.jsonl.gz"
+ cuts_path = in_out_dir / f"{prefix}_cuts_{partition}.{suffix}"
if cuts_path.is_file():
logging.info(f"{cuts_path} exists - skipping")
continue
- raw_cuts_path = in_out_dir / f"cuts_{partition}_raw.jsonl.gz"
+ raw_cuts_path = in_out_dir / f"{prefix}_cuts_{partition}_raw.{suffix}"
logging.info(f"Loading {raw_cuts_path}")
cut_set = CutSet.from_file(raw_cuts_path)
@@ -66,9 +65,10 @@ def compute_fbank_gigaspeech_dev_test():
cut_set = cut_set.compute_and_store_features_batch(
extractor=extractor,
- storage_path=f"{in_out_dir}/feats_{partition}",
+ storage_path=f"{in_out_dir}/{prefix}_feats_{partition}",
num_workers=num_workers,
batch_duration=batch_duration,
+ overwrite=True,
)
cut_set = cut_set.trim_to_supervisions(
keep_overlapping=False, min_duration=None
diff --git a/egs/librispeech/ASR/local/compute_fbank_gigaspeech_splits.py b/egs/librispeech/ASR/local/compute_fbank_gigaspeech_splits.py
index a7ed2467d..5587106e5 100644
--- a/egs/librispeech/ASR/local/compute_fbank_gigaspeech_splits.py
+++ b/egs/librispeech/ASR/local/compute_fbank_gigaspeech_splits.py
@@ -77,7 +77,7 @@ def get_parser():
def compute_fbank_gigaspeech_splits(args):
num_splits = args.num_splits
- output_dir = f"data/fbank/XL_split_{num_splits}"
+ output_dir = f"data/fbank/gigaspeech_XL_split_{num_splits}"
output_dir = Path(output_dir)
assert output_dir.exists(), f"{output_dir} does not exist!"
@@ -96,17 +96,19 @@ def compute_fbank_gigaspeech_splits(args):
extractor = KaldifeatFbank(KaldifeatFbankConfig(device=device))
logging.info(f"device: {device}")
+ prefix = "gigaspeech"
+
num_digits = 8 # num_digits is fixed by lhotse split-lazy
for i in range(start, stop):
idx = f"{i + 1}".zfill(num_digits)
logging.info(f"Processing {idx}/{num_splits}")
- cuts_path = output_dir / f"cuts_XL.{idx}.jsonl.gz"
+ cuts_path = output_dir / f"{prefix}_cuts_XL.{idx}.jsonl.gz"
if cuts_path.is_file():
logging.info(f"{cuts_path} exists - skipping")
continue
- raw_cuts_path = output_dir / f"cuts_XL_raw.{idx}.jsonl.gz"
+ raw_cuts_path = output_dir / f"{prefix}_cuts_XL_raw.{idx}.jsonl.gz"
if not raw_cuts_path.is_file():
logging.info(f"{raw_cuts_path} does not exist - skipping it")
continue
@@ -115,15 +117,16 @@ def compute_fbank_gigaspeech_splits(args):
cut_set = CutSet.from_file(raw_cuts_path)
logging.info("Computing features")
- if (output_dir / f"feats_XL_{idx}.lca").exists():
- logging.info(f"Removing {output_dir}/feats_XL_{idx}.lca")
- os.remove(output_dir / f"feats_XL_{idx}.lca")
+ if (output_dir / f"{prefix}_feats_XL_{idx}.lca").exists():
+ logging.info(f"Removing {output_dir}/{prefix}_feats_XL_{idx}.lca")
+ os.remove(output_dir / f"{prefix}_feats_XL_{idx}.lca")
cut_set = cut_set.compute_and_store_features_batch(
extractor=extractor,
- storage_path=f"{output_dir}/feats_XL_{idx}",
+ storage_path=f"{output_dir}/{prefix}_feats_XL_{idx}",
num_workers=args.num_workers,
batch_duration=args.batch_duration,
+ overwrite=True,
)
logging.info("About to split cuts into smaller chunks.")
diff --git a/egs/librispeech/ASR/local/compute_fbank_librispeech.py b/egs/librispeech/ASR/local/compute_fbank_librispeech.py
index 92f4f6ab7..f3e15e039 100755
--- a/egs/librispeech/ASR/local/compute_fbank_librispeech.py
+++ b/egs/librispeech/ASR/local/compute_fbank_librispeech.py
@@ -28,7 +28,7 @@ import os
from pathlib import Path
import torch
-from lhotse import ChunkedLilcomHdf5Writer, CutSet, Fbank, FbankConfig
+from lhotse import CutSet, Fbank, FbankConfig, LilcomChunkyWriter
from lhotse.recipes.utils import read_manifests_if_cached
from icefall.utils import get_executor
@@ -56,16 +56,29 @@ def compute_fbank_librispeech():
"train-clean-360",
"train-other-500",
)
+ prefix = "librispeech"
+ suffix = "jsonl.gz"
manifests = read_manifests_if_cached(
- prefix="librispeech", dataset_parts=dataset_parts, output_dir=src_dir
+ dataset_parts=dataset_parts,
+ output_dir=src_dir,
+ prefix=prefix,
+ suffix=suffix,
)
assert manifests is not None
+ assert len(manifests) == len(dataset_parts), (
+ len(manifests),
+ len(dataset_parts),
+ list(manifests.keys()),
+ dataset_parts,
+ )
+
extractor = Fbank(FbankConfig(num_mel_bins=num_mel_bins))
with get_executor() as ex: # Initialize the executor only once.
for partition, m in manifests.items():
- if (output_dir / f"cuts_{partition}.json.gz").is_file():
+ cuts_filename = f"{prefix}_cuts_{partition}.{suffix}"
+ if (output_dir / cuts_filename).is_file():
logging.info(f"{partition} already exists - skipping.")
continue
logging.info(f"Processing {partition}")
@@ -81,13 +94,13 @@ def compute_fbank_librispeech():
)
cut_set = cut_set.compute_and_store_features(
extractor=extractor,
- storage_path=f"{output_dir}/feats_{partition}",
+ storage_path=f"{output_dir}/{prefix}_feats_{partition}",
# when an executor is specified, make more partitions
num_jobs=num_jobs if ex is None else 80,
executor=ex,
- storage_type=ChunkedLilcomHdf5Writer,
+ storage_type=LilcomChunkyWriter,
)
- cut_set.to_json(output_dir / f"cuts_{partition}.json.gz")
+ cut_set.to_file(output_dir / cuts_filename)
if __name__ == "__main__":
diff --git a/egs/librispeech/ASR/local/compute_fbank_musan.py b/egs/librispeech/ASR/local/compute_fbank_musan.py
index 368bea4e8..056da29e5 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 ChunkedLilcomHdf5Writer, CutSet, Fbank, FbankConfig, combine
+from lhotse import CutSet, Fbank, FbankConfig, LilcomChunkyWriter, combine
from lhotse.recipes.utils import read_manifests_if_cached
from icefall.utils import get_executor
@@ -52,12 +52,24 @@ def compute_fbank_musan():
"speech",
"noise",
)
+ prefix = "musan"
+ suffix = "jsonl.gz"
manifests = read_manifests_if_cached(
- prefix="musan", dataset_parts=dataset_parts, output_dir=src_dir
+ dataset_parts=dataset_parts,
+ output_dir=src_dir,
+ prefix=prefix,
+ suffix=suffix,
)
assert manifests is not None
- musan_cuts_path = output_dir / "cuts_musan.json.gz"
+ assert len(manifests) == len(dataset_parts), (
+ len(manifests),
+ len(dataset_parts),
+ list(manifests.keys()),
+ dataset_parts,
+ )
+
+ musan_cuts_path = output_dir / "musan_cuts.jsonl.gz"
if musan_cuts_path.is_file():
logging.info(f"{musan_cuts_path} already exists - skipping")
@@ -79,13 +91,13 @@ def compute_fbank_musan():
.filter(lambda c: c.duration > 5)
.compute_and_store_features(
extractor=extractor,
- storage_path=f"{output_dir}/feats_musan",
+ storage_path=f"{output_dir}/musan_feats",
num_jobs=num_jobs if ex is None else 80,
executor=ex,
- storage_type=ChunkedLilcomHdf5Writer,
+ storage_type=LilcomChunkyWriter,
)
)
- musan_cuts.to_json(musan_cuts_path)
+ musan_cuts.to_file(musan_cuts_path)
if __name__ == "__main__":
diff --git a/egs/librispeech/ASR/local/display_manifest_statistics.py b/egs/librispeech/ASR/local/display_manifest_statistics.py
index 15bd206fa..c3c684235 100755
--- a/egs/librispeech/ASR/local/display_manifest_statistics.py
+++ b/egs/librispeech/ASR/local/display_manifest_statistics.py
@@ -25,19 +25,19 @@ for usage.
"""
-from lhotse import load_manifest
+from lhotse import load_manifest_lazy
def main():
- path = "./data/fbank/cuts_train-clean-100.json.gz"
- path = "./data/fbank/cuts_train-clean-360.json.gz"
- path = "./data/fbank/cuts_train-other-500.json.gz"
- path = "./data/fbank/cuts_dev-clean.json.gz"
- path = "./data/fbank/cuts_dev-other.json.gz"
- path = "./data/fbank/cuts_test-clean.json.gz"
- path = "./data/fbank/cuts_test-other.json.gz"
+ # path = "./data/fbank/librispeech_cuts_train-clean-100.jsonl.gz"
+ # path = "./data/fbank/librispeech_cuts_train-clean-360.jsonl.gz"
+ # path = "./data/fbank/librispeech_cuts_train-other-500.jsonl.gz"
+ # path = "./data/fbank/librispeech_cuts_dev-clean.jsonl.gz"
+ # path = "./data/fbank/librispeech_cuts_dev-other.jsonl.gz"
+ # path = "./data/fbank/librispeech_cuts_test-clean.jsonl.gz"
+ path = "./data/fbank/librispeech_cuts_test-other.jsonl.gz"
- cuts = load_manifest(path)
+ cuts = load_manifest_lazy(path)
cuts.describe()
diff --git a/egs/librispeech/ASR/local/download_lm.py b/egs/librispeech/ASR/local/download_lm.py
index 94d23afed..030122aa7 100755
--- a/egs/librispeech/ASR/local/download_lm.py
+++ b/egs/librispeech/ASR/local/download_lm.py
@@ -23,6 +23,7 @@ This file downloads the following LibriSpeech LM files:
- 4-gram.arpa.gz
- librispeech-vocab.txt
- librispeech-lexicon.txt
+ - librispeech-lm-norm.txt.gz
from http://www.openslr.org/resources/11
and save them in the user provided directory.
@@ -61,6 +62,7 @@ def main(out_dir: str):
"4-gram.arpa.gz",
"librispeech-vocab.txt",
"librispeech-lexicon.txt",
+ "librispeech-lm-norm.txt.gz",
)
for f in tqdm(files_to_download, desc="Downloading LibriSpeech LM files"):
diff --git a/egs/librispeech/ASR/local/prepare_lm_training_data.py b/egs/librispeech/ASR/local/prepare_lm_training_data.py
new file mode 100755
index 000000000..5070341f1
--- /dev/null
+++ b/egs/librispeech/ASR/local/prepare_lm_training_data.py
@@ -0,0 +1,172 @@
+#!/usr/bin/env python3
+
+# Copyright (c) 2021 Xiaomi Corporation (authors: Daniel Povey
+# 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 takes a `bpe.model` and a text file such as
+./download/lm/librispeech-lm-norm.txt
+and outputs the LM training data to a supplied directory such
+as data/lm_training_bpe_500. The format is as follows:
+
+It creates a PyTorch archive (.pt file), say data/lm_training.pt, which is a
+representation of a dict with the following format:
+
+ 'words' -> a k2.RaggedTensor of two axes [word][token] with dtype torch.int32
+ containing the BPE representations of each word, indexed by
+ integer word ID. (These integer word IDS are present in
+ 'lm_data'). The sentencepiece object can be used to turn the
+ words and BPE units into string form.
+ 'sentences' -> a k2.RaggedTensor of two axes [sentence][word] with dtype
+ torch.int32 containing all the sentences, as word-ids (we don't
+ output the string form of this directly but it can be worked out
+ together with 'words' and the bpe.model).
+ 'sentence_lengths' -> a 1-D torch.Tensor of dtype torch.int32, containing
+ number of BPE tokens of each sentence.
+"""
+
+import argparse
+import logging
+from pathlib import Path
+
+import k2
+import sentencepiece as spm
+import torch
+
+
+def get_args():
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ help="Input BPE model, e.g. data/bpe_500/bpe.model",
+ )
+ parser.add_argument(
+ "--lm-data",
+ type=str,
+ help="""Input LM training data as text, e.g.
+ download/pb.train.txt""",
+ )
+ parser.add_argument(
+ "--lm-archive",
+ type=str,
+ help="""Path to output archive, e.g. data/bpe_500/lm_data.pt;
+ look at the source of this script to see the format.""",
+ )
+
+ return parser.parse_args()
+
+
+def main():
+ args = get_args()
+
+ if Path(args.lm_archive).exists():
+ logging.warning(f"{args.lm_archive} exists - skipping")
+ return
+
+ sp = spm.SentencePieceProcessor()
+ sp.load(args.bpe_model)
+
+ # word2index is a dictionary from words to integer ids. No need to reserve
+ # space for epsilon, etc.; the words are just used as a convenient way to
+ # compress the sequences of BPE pieces.
+ word2index = dict()
+
+ word2bpe = [] # Will be a list-of-list-of-int, representing BPE pieces.
+ sentences = [] # Will be a list-of-list-of-int, representing word-ids.
+
+ if "librispeech-lm-norm" in args.lm_data:
+ num_lines_in_total = 40418261.0
+ step = 5000000
+ elif "valid" in args.lm_data:
+ num_lines_in_total = 5567.0
+ step = 3000
+ elif "test" in args.lm_data:
+ num_lines_in_total = 5559.0
+ step = 3000
+ else:
+ num_lines_in_total = None
+ step = None
+
+ processed = 0
+
+ with open(args.lm_data) as f:
+ while True:
+ line = f.readline()
+ if line == "":
+ break
+
+ if step and processed % step == 0:
+ logging.info(
+ f"Processed number of lines: {processed} "
+ f"({processed/num_lines_in_total*100: .3f}%)"
+ )
+ processed += 1
+
+ line_words = line.split()
+ for w in line_words:
+ if w not in word2index:
+ w_bpe = sp.encode(w)
+ word2index[w] = len(word2bpe)
+ word2bpe.append(w_bpe)
+ sentences.append([word2index[w] for w in line_words])
+
+ logging.info("Constructing ragged tensors")
+ words = k2.ragged.RaggedTensor(word2bpe)
+ sentences = k2.ragged.RaggedTensor(sentences)
+
+ output = dict(words=words, sentences=sentences)
+
+ num_sentences = sentences.dim0
+ logging.info(f"Computing sentence lengths, num_sentences: {num_sentences}")
+ sentence_lengths = [0] * num_sentences
+ for i in range(num_sentences):
+ if step and i % step == 0:
+ logging.info(
+ f"Processed number of lines: {i} "
+ f"({i/num_sentences*100: .3f}%)"
+ )
+
+ word_ids = sentences[i]
+
+ # NOTE: If word_ids is a tensor with only 1 entry,
+ # token_ids is a torch.Tensor
+ token_ids = words[word_ids]
+ if isinstance(token_ids, k2.RaggedTensor):
+ token_ids = token_ids.values
+
+ # token_ids is a 1-D tensor containing the BPE tokens
+ # of the current sentence
+
+ sentence_lengths[i] = token_ids.numel()
+
+ output["sentence_lengths"] = torch.tensor(
+ sentence_lengths, dtype=torch.int32
+ )
+
+ torch.save(output, args.lm_archive)
+ logging.info(f"Saved to {args.lm_archive}")
+
+
+if __name__ == "__main__":
+ formatter = (
+ "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
+ )
+
+ logging.basicConfig(format=formatter, level=logging.INFO)
+
+ main()
diff --git a/egs/librispeech/ASR/local/preprocess_gigaspeech.py b/egs/librispeech/ASR/local/preprocess_gigaspeech.py
index cd1345904..077f23039 100644
--- a/egs/librispeech/ASR/local/preprocess_gigaspeech.py
+++ b/egs/librispeech/ASR/local/preprocess_gigaspeech.py
@@ -58,17 +58,26 @@ def preprocess_giga_speech():
)
logging.info("Loading manifest (may take 4 minutes)")
+ prefix = "gigaspeech"
+ suffix = "jsonl.gz"
manifests = read_manifests_if_cached(
dataset_parts=dataset_parts,
output_dir=src_dir,
- prefix="gigaspeech",
- suffix="jsonl.gz",
+ prefix=prefix,
+ suffix=suffix,
)
assert manifests is not None
+ assert len(manifests) == len(dataset_parts), (
+ len(manifests),
+ len(dataset_parts),
+ list(manifests.keys()),
+ dataset_parts,
+ )
+
for partition, m in manifests.items():
logging.info(f"Processing {partition}")
- raw_cuts_path = output_dir / f"cuts_{partition}_raw.jsonl.gz"
+ raw_cuts_path = output_dir / f"{prefix}_cuts_{partition}_raw.{suffix}"
if raw_cuts_path.is_file():
logging.info(f"{partition} already exists - skipping")
continue
diff --git a/egs/librispeech/ASR/local/sort_lm_training_data.py b/egs/librispeech/ASR/local/sort_lm_training_data.py
new file mode 120000
index 000000000..bb26b5f5c
--- /dev/null
+++ b/egs/librispeech/ASR/local/sort_lm_training_data.py
@@ -0,0 +1 @@
+../../../ptb/LM/local/sort_lm_training_data.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/local/train_bpe_model.py b/egs/librispeech/ASR/local/train_bpe_model.py
index bc5812810..42aba9572 100755
--- a/egs/librispeech/ASR/local/train_bpe_model.py
+++ b/egs/librispeech/ASR/local/train_bpe_model.py
@@ -38,7 +38,6 @@ def get_args():
"--lang-dir",
type=str,
help="""Input and output directory.
- It should contain the training corpus: transcript_words.txt.
The generated bpe.model is saved to this directory.
""",
)
diff --git a/egs/librispeech/ASR/local/validate_manifest.py b/egs/librispeech/ASR/local/validate_manifest.py
index 8d3d4c7ce..7c57d629a 100755
--- a/egs/librispeech/ASR/local/validate_manifest.py
+++ b/egs/librispeech/ASR/local/validate_manifest.py
@@ -25,7 +25,7 @@ We will add more checks later if needed.
Usage example:
python3 ./local/validate_manifest.py \
- ./data/fbank/cuts_train-clean-100.json.gz
+ ./data/fbank/librispeech_cuts_train-clean-100.jsonl.gz
"""
@@ -33,7 +33,7 @@ import argparse
import logging
from pathlib import Path
-from lhotse import load_manifest, CutSet
+from lhotse import CutSet, load_manifest_lazy
from lhotse.cut import Cut
@@ -76,7 +76,7 @@ def main():
logging.info(f"Validating {manifest}")
assert manifest.is_file(), f"{manifest} does not exist"
- cut_set = load_manifest(manifest)
+ cut_set = load_manifest_lazy(manifest)
assert isinstance(cut_set, CutSet)
for c in cut_set:
diff --git a/egs/librispeech/ASR/lstm_transducer_stateless/__init__.py b/egs/librispeech/ASR/lstm_transducer_stateless/__init__.py
new file mode 120000
index 000000000..b24e5e357
--- /dev/null
+++ b/egs/librispeech/ASR/lstm_transducer_stateless/__init__.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/__init__.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/lstm_transducer_stateless/asr_datamodule.py b/egs/librispeech/ASR/lstm_transducer_stateless/asr_datamodule.py
new file mode 120000
index 000000000..a074d6085
--- /dev/null
+++ b/egs/librispeech/ASR/lstm_transducer_stateless/asr_datamodule.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/asr_datamodule.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/lstm_transducer_stateless/beam_search.py b/egs/librispeech/ASR/lstm_transducer_stateless/beam_search.py
new file mode 120000
index 000000000..8554e44cc
--- /dev/null
+++ b/egs/librispeech/ASR/lstm_transducer_stateless/beam_search.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/beam_search.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/lstm_transducer_stateless/decode.py b/egs/librispeech/ASR/lstm_transducer_stateless/decode.py
new file mode 100755
index 000000000..27414d717
--- /dev/null
+++ b/egs/librispeech/ASR/lstm_transducer_stateless/decode.py
@@ -0,0 +1,818 @@
+#!/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
+./lstm_transducer_stateless/decode.py \
+ --epoch 35 \
+ --avg 15 \
+ --exp-dir ./lstm_transducer_stateless/exp \
+ --max-duration 600 \
+ --decoding-method greedy_search
+
+(2) beam search (not recommended)
+./lstm_transducer_stateless/decode.py \
+ --epoch 35 \
+ --avg 15 \
+ --exp-dir ./lstm_transducer_stateless/exp \
+ --max-duration 600 \
+ --decoding-method beam_search \
+ --beam-size 4
+
+(3) modified beam search
+./lstm_transducer_stateless/decode.py \
+ --epoch 35 \
+ --avg 15 \
+ --exp-dir ./lstm_transducer_stateless/exp \
+ --max-duration 600 \
+ --decoding-method modified_beam_search \
+ --beam-size 4
+
+(4) fast beam search (one best)
+./lstm_transducer_stateless/decode.py \
+ --epoch 35 \
+ --avg 15 \
+ --exp-dir ./lstm_transducer_stateless/exp \
+ --max-duration 600 \
+ --decoding-method fast_beam_search \
+ --beam 20.0 \
+ --max-contexts 8 \
+ --max-states 64
+
+(5) fast beam search (nbest)
+./lstm_transducer_stateless/decode.py \
+ --epoch 30 \
+ --avg 15 \
+ --exp-dir ./pruned_transducer_stateless3/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)
+./lstm_transducer_stateless/decode.py \
+ --epoch 35 \
+ --avg 15 \
+ --exp-dir ./lstm_transducer_stateless/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)
+./lstm_transducer_stateless/decode.py \
+ --epoch 35 \
+ --avg 15 \
+ --exp-dir ./lstm_transducer_stateless/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=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="lstm_transducer_stateless/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 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.
+ 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)
+
+ # tail padding here to alleviate the tail deletion problem
+ num_tail_padded_frames = 35
+ feature = torch.nn.functional.pad(
+ feature,
+ (0, 0, 0, num_tail_padded_frames),
+ mode="constant",
+ value=LOG_EPS,
+ )
+ feature_lens += num_tail_padded_frames
+
+ 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 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.
+ 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 "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)
+
+ 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/lstm_transducer_stateless/decoder.py b/egs/librispeech/ASR/lstm_transducer_stateless/decoder.py
new file mode 120000
index 000000000..0793c5709
--- /dev/null
+++ b/egs/librispeech/ASR/lstm_transducer_stateless/decoder.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/decoder.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/lstm_transducer_stateless/encoder_interface.py b/egs/librispeech/ASR/lstm_transducer_stateless/encoder_interface.py
new file mode 120000
index 000000000..aa5d0217a
--- /dev/null
+++ b/egs/librispeech/ASR/lstm_transducer_stateless/encoder_interface.py
@@ -0,0 +1 @@
+../transducer_stateless/encoder_interface.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/lstm_transducer_stateless/export.py b/egs/librispeech/ASR/lstm_transducer_stateless/export.py
new file mode 100755
index 000000000..13dac6009
--- /dev/null
+++ b/egs/librispeech/ASR/lstm_transducer_stateless/export.py
@@ -0,0 +1,388 @@
+#!/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.
+
+# This script converts several saved checkpoints
+# to a single one using model averaging.
+"""
+
+Usage:
+
+(1) Export to torchscript model using torch.jit.trace()
+
+./lstm_transducer_stateless/export.py \
+ --exp-dir ./lstm_transducer_stateless/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 35 \
+ --avg 10 \
+ --jit-trace 1
+
+It will generate 3 files: `encoder_jit_trace.pt`,
+`decoder_jit_trace.pt`, and `joiner_jit_trace.pt`.
+
+(2) Export `model.state_dict()`
+
+./lstm_transducer_stateless/export.py \
+ --exp-dir ./lstm_transducer_stateless/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 35 \
+ --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 `lstm_transducer_stateless/decode.py`,
+you can do:
+
+ cd /path/to/exp_dir
+ ln -s pretrained.pt epoch-9999.pt
+
+ cd /path/to/egs/librispeech/ASR
+ ./lstm_transducer_stateless/decode.py \
+ --exp-dir ./lstm_transducer_stateless/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-lstm-transducer-stateless-2022-08-18
+
+with the following commands:
+
+ sudo apt-get install git-lfs
+ git lfs install
+ git clone https://huggingface.co/Zengwei/icefall-asr-librispeech-lstm-transducer-stateless-2022-08-18
+ # You will find the pre-trained model in icefall-asr-librispeech-lstm-transducer-stateless-2022-08-18/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=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(
+ "--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_stateless3/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-trace",
+ type=str2bool,
+ default=False,
+ help="""True to save a model after applying torch.jit.trace.
+ It will generate 3 files:
+ - encoder_jit_trace.pt
+ - decoder_jit_trace.pt
+ - joiner_jit_trace.pt
+
+ Check ./jit_pretrained.py for how to use them.
+ """,
+ )
+
+ 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
+
+
+def export_encoder_model_jit_trace(
+ encoder_model: nn.Module,
+ encoder_filename: str,
+) -> 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.
+ """
+ x = torch.zeros(1, 100, 80, dtype=torch.float32)
+ x_lens = torch.tensor([100], dtype=torch.int64)
+ states = encoder_model.get_init_states()
+
+ 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: 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: 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")
+ 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)
+
+ 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_trace is True:
+ convert_scaled_to_non_scaled(model, inplace=True)
+ logging.info("Using torch.jit.trace()")
+ encoder_filename = params.exp_dir / "encoder_jit_trace.pt"
+ export_encoder_model_jit_trace(model.encoder, encoder_filename)
+
+ decoder_filename = params.exp_dir / "decoder_jit_trace.pt"
+ export_decoder_model_jit_trace(model.decoder, decoder_filename)
+
+ joiner_filename = params.exp_dir / "joiner_jit_trace.pt"
+ export_joiner_model_jit_trace(model.joiner, joiner_filename)
+ else:
+ logging.info("Not using torchscript")
+ # 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/lstm_transducer_stateless/jit_pretrained.py b/egs/librispeech/ASR/lstm_transducer_stateless/jit_pretrained.py
new file mode 100755
index 000000000..594c33e4f
--- /dev/null
+++ b/egs/librispeech/ASR/lstm_transducer_stateless/jit_pretrained.py
@@ -0,0 +1,322 @@
+#!/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, either exported by `torch.jit.trace()`
+or by `torch.jit.script()`, and uses them to decode waves.
+You can use the following command to get the exported models:
+
+./lstm_transducer_stateless/export.py \
+ --exp-dir ./lstm_transducer_stateless/exp \
+ --bpe-model data/lang_bpe_500/bpe.model \
+ --epoch 20 \
+ --avg 10 \
+ --jit-trace 1
+
+Usage of this script:
+
+./lstm_transducer_stateless/jit_pretrained.py \
+ --encoder-model-filename ./lstm_transducer_stateless/exp/encoder_jit_trace.pt \
+ --decoder-model-filename ./lstm_transducer_stateless/exp/decoder_jit_trace.pt \
+ --joiner-model-filename ./lstm_transducer_stateless/exp/joiner_jit_trace.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(
+ "--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(
+ "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}. "
+ f"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,
+ encoder_out_lens: torch.Tensor,
+ 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.
+ 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.
+ """
+ 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)
+
+ hyps = [[blank_id] * context_size for _ in range(N)]
+
+ decoder_input = torch.tensor(
+ hyps,
+ device=device,
+ dtype=torch.int64,
+ ) # (N, context_size)
+
+ decoder_out = 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 = 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 = 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}")
+
+ 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")
+ opts = kaldifeat.FbankOptions()
+ opts.device = device
+ 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,
+ )
+ 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)
+
+ states = encoder.get_init_states(batch_size=features.size(0), device=device)
+
+ encoder_out, encoder_out_lens, _ = encoder(
+ x=features,
+ x_lens=feature_lengths,
+ states=states,
+ )
+
+ hyps = greedy_search(
+ decoder=decoder,
+ joiner=joiner,
+ 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/lstm_transducer_stateless/joiner.py b/egs/librispeech/ASR/lstm_transducer_stateless/joiner.py
new file mode 120000
index 000000000..815fd4bb6
--- /dev/null
+++ b/egs/librispeech/ASR/lstm_transducer_stateless/joiner.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/joiner.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/lstm_transducer_stateless/lstm.py b/egs/librispeech/ASR/lstm_transducer_stateless/lstm.py
new file mode 100644
index 000000000..c54a4c478
--- /dev/null
+++ b/egs/librispeech/ASR/lstm_transducer_stateless/lstm.py
@@ -0,0 +1,871 @@
+# 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.
+
+import copy
+import math
+from typing import List, Optional, Tuple
+
+import torch
+from encoder_interface import EncoderInterface
+from scaling import (
+ ActivationBalancer,
+ BasicNorm,
+ DoubleSwish,
+ ScaledConv2d,
+ ScaledLinear,
+ ScaledLSTM,
+)
+from torch import nn
+
+LOG_EPSILON = math.log(1e-10)
+
+
+def unstack_states(
+ states: Tuple[torch.Tensor, torch.Tensor]
+) -> List[Tuple[torch.Tensor, torch.Tensor]]:
+ """
+ Unstack the lstm states corresponding to a batch of utterances into a list
+ of states, where the i-th entry is the state from the i-th utterance.
+
+ Args:
+ states:
+ A tuple of 2 elements.
+ ``states[0]`` is the lstm hidden states, of a batch of utterance.
+ ``states[1]`` is the lstm cell states, of a batch of utterances.
+
+ Returns:
+ A list of states.
+ ``states[i]`` is a tuple of 2 elememts of i-th utterance.
+ ``states[i][0]`` is the lstm hidden states of i-th utterance.
+ ``states[i][1]`` is the lstm cell states of i-th utterance.
+ """
+ hidden_states, cell_states = states
+
+ list_hidden_states = hidden_states.unbind(dim=1)
+ list_cell_states = cell_states.unbind(dim=1)
+
+ ans = [
+ (h.unsqueeze(1), c.unsqueeze(1))
+ for (h, c) in zip(list_hidden_states, list_cell_states)
+ ]
+ return ans
+
+
+def stack_states(
+ states_list: List[Tuple[torch.Tensor, torch.Tensor]]
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Stack list of lstm states corresponding to separate utterances into a single
+ lstm state so that it can be used as an input for lstm when those utterances
+ are formed into a batch.
+
+ Args:
+ state_list:
+ Each element in state_list corresponds to the lstm state for a single
+ utterance.
+ ``states[i]`` is a tuple of 2 elememts of i-th utterance.
+ ``states[i][0]`` is the lstm hidden states of i-th utterance.
+ ``states[i][1]`` is the lstm cell states of i-th utterance.
+
+
+ Returns:
+ A new state corresponding to a batch of utterances.
+ It is a tuple of 2 elements.
+ ``states[0]`` is the lstm hidden states, of a batch of utterance.
+ ``states[1]`` is the lstm cell states, of a batch of utterances.
+ """
+ hidden_states = torch.cat([s[0] for s in states_list], dim=1)
+ cell_states = torch.cat([s[1] for s in states_list], dim=1)
+ ans = (hidden_states, cell_states)
+ return ans
+
+
+class RNN(EncoderInterface):
+ """
+ Args:
+ num_features (int):
+ Number of input features.
+ subsampling_factor (int):
+ Subsampling factor of encoder (convolution layers before lstm layers) (default=4). # noqa
+ d_model (int):
+ Output dimension (default=512).
+ dim_feedforward (int):
+ Feedforward dimension (default=2048).
+ rnn_hidden_size (int):
+ Hidden dimension for lstm layers (default=1024).
+ num_encoder_layers (int):
+ Number of encoder layers (default=12).
+ dropout (float):
+ Dropout rate (default=0.1).
+ layer_dropout (float):
+ Dropout value for model-level warmup (default=0.075).
+ aux_layer_period (int):
+ Period of auxiliary layers used for random combiner during training.
+ If set to 0, will not use the random combiner (Default).
+ You can set a positive integer to use the random combiner, e.g., 3.
+ is_pnnx:
+ True to make this class exportable via PNNX.
+ """
+
+ def __init__(
+ self,
+ num_features: int,
+ subsampling_factor: int = 4,
+ d_model: int = 512,
+ dim_feedforward: int = 2048,
+ rnn_hidden_size: int = 1024,
+ num_encoder_layers: int = 12,
+ dropout: float = 0.1,
+ layer_dropout: float = 0.075,
+ aux_layer_period: int = 0,
+ is_pnnx: bool = False,
+ ) -> None:
+ super(RNN, self).__init__()
+
+ 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,
+ is_pnnx=is_pnnx,
+ )
+
+ self.is_pnnx = is_pnnx
+
+ self.num_encoder_layers = num_encoder_layers
+ self.d_model = d_model
+ self.rnn_hidden_size = rnn_hidden_size
+
+ encoder_layer = RNNEncoderLayer(
+ d_model=d_model,
+ dim_feedforward=dim_feedforward,
+ rnn_hidden_size=rnn_hidden_size,
+ dropout=dropout,
+ layer_dropout=layer_dropout,
+ )
+ self.encoder = RNNEncoder(
+ encoder_layer,
+ num_encoder_layers,
+ aux_layers=list(
+ range(
+ num_encoder_layers // 3,
+ num_encoder_layers - 1,
+ aux_layer_period,
+ )
+ )
+ if aux_layer_period > 0
+ else None,
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ x_lens: torch.Tensor,
+ states: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ warmup: float = 1.0,
+ ) -> Tuple[torch.Tensor, torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+ """
+ Args:
+ x:
+ The input tensor. Its shape is (N, T, C), where N is the batch size,
+ T is the sequence length, C is the feature dimension.
+ x_lens:
+ A tensor of shape (N,), containing the number of frames in `x`
+ before padding.
+ states:
+ A tuple of 2 tensors (optional). It is for streaming inference.
+ states[0] is the hidden states of all layers,
+ with shape of (num_layers, N, d_model);
+ states[1] is the cell states of all layers,
+ with shape of (num_layers, N, rnn_hidden_size).
+ 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 tuple of 3 tensors:
+ - embeddings: its shape is (N, T', d_model), where T' is the output
+ sequence lengths.
+ - lengths: a tensor of shape (batch_size,) containing the number of
+ frames in `embeddings` before padding.
+ - updated states, whose shape is the same as the input states.
+ """
+ x = self.encoder_embed(x)
+ x = x.permute(1, 0, 2) # (N, T, C) -> (T, N, C)
+
+ # lengths = ((x_lens - 3) // 2 - 1) // 2 # issue an warning
+ #
+ # Note: rounding_mode in torch.div() is available only in torch >= 1.8.0
+ if not self.is_pnnx:
+ lengths = (((x_lens - 3) >> 1) - 1) >> 1
+ else:
+ lengths1 = torch.floor((x_lens - 3) / 2)
+ lengths = torch.floor((lengths1 - 1) / 2)
+ lengths = lengths.to(x_lens)
+
+ if not torch.jit.is_tracing():
+ assert x.size(0) == lengths.max().item()
+
+ if states is None:
+ x = self.encoder(x, warmup=warmup)[0]
+ # torch.jit.trace requires returned types to be the same as annotated # noqa
+ new_states = (torch.empty(0), torch.empty(0))
+ else:
+ assert not self.training
+ assert len(states) == 2
+ if not torch.jit.is_tracing():
+ # for hidden state
+ assert states[0].shape == (
+ self.num_encoder_layers,
+ x.size(1),
+ self.d_model,
+ )
+ # for cell state
+ assert states[1].shape == (
+ self.num_encoder_layers,
+ x.size(1),
+ self.rnn_hidden_size,
+ )
+ x, new_states = self.encoder(x, states)
+
+ x = x.permute(1, 0, 2) # (T, N, C) -> (N, T, C)
+ return x, lengths, new_states
+
+ @torch.jit.export
+ def get_init_states(
+ self, batch_size: int = 1, device: torch.device = torch.device("cpu")
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Get model initial states."""
+ # for rnn hidden states
+ hidden_states = torch.zeros(
+ (self.num_encoder_layers, batch_size, self.d_model), device=device
+ )
+ cell_states = torch.zeros(
+ (self.num_encoder_layers, batch_size, self.rnn_hidden_size),
+ device=device,
+ )
+ return (hidden_states, cell_states)
+
+
+class RNNEncoderLayer(nn.Module):
+ """
+ RNNEncoderLayer is made up of lstm and feedforward networks.
+
+ Args:
+ d_model:
+ The number of expected features in the input (required).
+ dim_feedforward:
+ The dimension of feedforward network model (default=2048).
+ rnn_hidden_size:
+ The hidden dimension of rnn layer.
+ dropout:
+ The dropout value (default=0.1).
+ layer_dropout:
+ The dropout value for model-level warmup (default=0.075).
+ """
+
+ def __init__(
+ self,
+ d_model: int,
+ dim_feedforward: int,
+ rnn_hidden_size: int,
+ dropout: float = 0.1,
+ layer_dropout: float = 0.075,
+ ) -> None:
+ super(RNNEncoderLayer, self).__init__()
+ self.layer_dropout = layer_dropout
+ self.d_model = d_model
+ self.rnn_hidden_size = rnn_hidden_size
+
+ assert rnn_hidden_size >= d_model, (rnn_hidden_size, d_model)
+ self.lstm = ScaledLSTM(
+ input_size=d_model,
+ hidden_size=rnn_hidden_size,
+ proj_size=d_model if rnn_hidden_size > d_model else 0,
+ num_layers=1,
+ 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.norm_final = BasicNorm(d_model)
+
+ # try to ensure the output is close to zero-mean (or at least, zero-median). # noqa
+ 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,
+ states: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ warmup: float = 1.0,
+ ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+ """
+ Pass the input through the encoder layer.
+
+ Args:
+ src:
+ The sequence to the encoder layer (required).
+ Its shape is (S, N, E), where S is the sequence length,
+ N is the batch size, and E is the feature number.
+ states:
+ A tuple of 2 tensors (optional). It is for streaming inference.
+ states[0] is the hidden states of all layers,
+ with shape of (1, N, d_model);
+ states[1] is the cell states of all layers,
+ with shape of (1, N, rnn_hidden_size).
+ warmup:
+ It controls selective bypass of of layers; if < 1.0, we will
+ bypass layers more frequently.
+ """
+ src_orig = src
+
+ warmup_scale = min(0.1 + 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 0.1
+ )
+ else:
+ alpha = 1.0
+
+ # lstm module
+ if states is None:
+ src_lstm = self.lstm(src)[0]
+ # torch.jit.trace requires returned types be the same as annotated
+ new_states = (torch.empty(0), torch.empty(0))
+ else:
+ assert not self.training
+ assert len(states) == 2
+ if not torch.jit.is_tracing():
+ # for hidden state
+ assert states[0].shape == (1, src.size(1), self.d_model)
+ # for cell state
+ assert states[1].shape == (1, src.size(1), self.rnn_hidden_size)
+ src_lstm, new_states = self.lstm(src, states)
+ src = self.dropout(src_lstm) + 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, new_states
+
+
+class RNNEncoder(nn.Module):
+ """
+ RNNEncoder is a stack of N encoder layers.
+
+ Args:
+ encoder_layer:
+ An instance of the RNNEncoderLayer() class (required).
+ num_layers:
+ The number of sub-encoder-layers in the encoder (required).
+ """
+
+ def __init__(
+ self,
+ encoder_layer: nn.Module,
+ num_layers: int,
+ aux_layers: Optional[List[int]] = None,
+ ) -> None:
+ super(RNNEncoder, self).__init__()
+ 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.rnn_hidden_size = encoder_layer.rnn_hidden_size
+
+ self.aux_layers: List[int] = []
+ self.combiner: Optional[nn.Module] = None
+ if aux_layers is not None:
+ 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,
+ states: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ warmup: float = 1.0,
+ ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+ """
+ Pass the input through the encoder layer in turn.
+
+ Args:
+ src:
+ The sequence to the encoder layer (required).
+ Its shape is (S, N, E), where S is the sequence length,
+ N is the batch size, and E is the feature number.
+ states:
+ A tuple of 2 tensors (optional). It is for streaming inference.
+ states[0] is the hidden states of all layers,
+ with shape of (num_layers, N, d_model);
+ states[1] is the cell states of all layers,
+ with shape of (num_layers, N, rnn_hidden_size).
+ warmup:
+ It controls selective bypass of of layers; if < 1.0, we will
+ bypass layers more frequently.
+ """
+ if states is not None:
+ assert not self.training
+ assert len(states) == 2
+ if not torch.jit.is_tracing():
+ # for hidden state
+ assert states[0].shape == (
+ self.num_layers,
+ src.size(1),
+ self.d_model,
+ )
+ # for cell state
+ assert states[1].shape == (
+ self.num_layers,
+ src.size(1),
+ self.rnn_hidden_size,
+ )
+
+ output = src
+
+ outputs = []
+
+ new_hidden_states = []
+ new_cell_states = []
+
+ for i, mod in enumerate(self.layers):
+ if states is None:
+ output = mod(output, warmup=warmup)[0]
+ else:
+ layer_state = (
+ states[0][i : i + 1, :, :], # h: (1, N, d_model)
+ states[1][i : i + 1, :, :], # c: (1, N, rnn_hidden_size)
+ )
+ output, (h, c) = mod(output, layer_state)
+ new_hidden_states.append(h)
+ new_cell_states.append(c)
+
+ if self.combiner is not None and i in self.aux_layers:
+ outputs.append(output)
+
+ if self.combiner is not None:
+ output = self.combiner(outputs)
+
+ if states is None:
+ new_states = (torch.empty(0), torch.empty(0))
+ else:
+ new_states = (
+ torch.cat(new_hidden_states, dim=0),
+ torch.cat(new_cell_states, dim=0),
+ )
+
+ return output, new_states
+
+
+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-1)//2, which approximates T' == T//4
+
+ 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,
+ is_pnnx: bool = False,
+ ) -> None:
+ """
+ Args:
+ in_channels:
+ Number of channels in. The input shape is (N, T, in_channels).
+ Caution: It requires: T >= 9, in_channels >= 9.
+ out_channels
+ Output dim. The output shape is (N, ((T-3)//2-1)//2, out_channels)
+ layer1_channels:
+ Number of channels in layer1
+ layer1_channels:
+ Number of channels in layer2
+ is_pnnx:
+ True if we are converting the model to PNNX format.
+ False otherwise.
+ """
+ assert in_channels >= 9
+ super().__init__()
+
+ self.conv = nn.Sequential(
+ ScaledConv2d(
+ in_channels=1,
+ out_channels=layer1_channels,
+ kernel_size=3,
+ padding=0,
+ ),
+ ActivationBalancer(channel_dim=1),
+ DoubleSwish(),
+ ScaledConv2d(
+ in_channels=layer1_channels,
+ out_channels=layer2_channels,
+ kernel_size=3,
+ stride=2,
+ ),
+ ActivationBalancer(channel_dim=1),
+ DoubleSwish(),
+ ScaledConv2d(
+ in_channels=layer2_channels,
+ out_channels=layer3_channels,
+ kernel_size=3,
+ stride=2,
+ ),
+ ActivationBalancer(channel_dim=1),
+ DoubleSwish(),
+ )
+ self.out = ScaledLinear(
+ layer3_channels * (((in_channels - 3) // 2 - 1) // 2), out_channels
+ )
+ # set learn_eps=False because out_norm is preceded by `out`, and `out`
+ # itself has learned scale, so the extra degree of freedom is not
+ # needed.
+ self.out_norm = BasicNorm(out_channels, learn_eps=False)
+ # constrain median of output to be close to zero.
+ self.out_balancer = ActivationBalancer(
+ channel_dim=-1, min_positive=0.45, max_positive=0.55
+ )
+
+ # ncnn supports only batch size == 1
+ self.is_pnnx = is_pnnx
+ self.conv_out_dim = self.out.weight.shape[1]
+
+ 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-3)//2-1)//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)
+
+ if torch.jit.is_tracing() and self.is_pnnx:
+ x = x.permute(0, 2, 1, 3).reshape(1, -1, self.conv_out_dim)
+ x = self.out(x)
+ else:
+ # Now x is of shape (N, odim, ((T-3)//2-1)//2, ((idim-3)//2-1)//2)
+ b, c, t, f = x.size()
+ x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f))
+
+ # Now x is of shape (N, ((T-3)//2-1))//2, odim)
+ x = self.out_norm(x)
+ x = self.out_balancer(x)
+ return x
+
+
+class RandomCombine(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
+
+ 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
+ if not self.training or torch.jit.is_scripting():
+ 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,))
+
+ # The following if causes errors for torch script in torch 1.6.0
+ # if __name__ == "__main__":
+ # # for testing only...
+ # print("Weights = ", weights.reshape(num_frames, num_inputs))
+ 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
+ ):
+ """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. # noqa
+ 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
+ ):
+ """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 # noqa
+ )
+ logprobs[:, -1] += self.final_log_weight
+ return logprobs.softmax(dim=1)
+
+
+def _test_random_combine(final_weight: float, pure_prob: float, stddev: float):
+ print(
+ f"_test_random_combine: final_weight={final_weight}, pure_prob={pure_prob}, stddev={stddev}" # noqa
+ )
+ 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():
+ _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)
+
+ feature_dim = 50
+ c = RNN(num_features=feature_dim, d_model=128)
+ batch_size = 5
+ seq_len = 20
+ # 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),
+ )
+ f # to remove flake8 warnings
+
+
+if __name__ == "__main__":
+ feature_dim = 80
+ m = RNN(
+ num_features=feature_dim,
+ d_model=512,
+ rnn_hidden_size=1024,
+ dim_feedforward=2048,
+ num_encoder_layers=12,
+ )
+ batch_size = 5
+ seq_len = 20
+ # Just make sure the forward pass runs.
+ f = m(
+ torch.randn(batch_size, seq_len, feature_dim),
+ torch.full((batch_size,), seq_len, dtype=torch.int64),
+ warmup=0.5,
+ )
+ num_param = sum([p.numel() for p in m.parameters()])
+ print(f"Number of model parameters: {num_param}")
+
+ _test_random_combine_main()
diff --git a/egs/librispeech/ASR/lstm_transducer_stateless/lstmp.py b/egs/librispeech/ASR/lstm_transducer_stateless/lstmp.py
new file mode 120000
index 000000000..4f377cd01
--- /dev/null
+++ b/egs/librispeech/ASR/lstm_transducer_stateless/lstmp.py
@@ -0,0 +1 @@
+../lstm_transducer_stateless2/lstmp.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/lstm_transducer_stateless/model.py b/egs/librispeech/ASR/lstm_transducer_stateless/model.py
new file mode 100644
index 000000000..efbc88a55
--- /dev/null
+++ b/egs/librispeech/ASR/lstm_transducer_stateless/model.py
@@ -0,0 +1,202 @@
+# 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 scaling import ScaledLinear
+
+from icefall.utils import add_sos
+
+
+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,
+ 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.simple_am_proj = ScaledLinear(
+ encoder_dim, vocab_size, initial_speed=0.5
+ )
+ self.simple_lm_proj = ScaledLinear(decoder_dim, vocab_size)
+
+ 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,
+ reduction: str = "sum",
+ ) -> 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.
+ 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 value warmup >= 0 that determines which modules are active, values
+ warmup > 1 "are fully warmed up" and all modules will be active.
+ reduction:
+ "sum" to sum the losses over all utterances in the batch.
+ "none" to return the loss in a 1-D tensor for each utterance
+ in the batch.
+ Returns:
+ Return the transducer loss.
+
+ 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 reduction in ("sum", "none"), reduction
+ 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, warmup=warmup)
+ assert torch.all(x_lens > 0)
+
+ # Now for the decoder, i.e., the prediction network
+ row_splits = y.shape.row_splits(1)
+ y_lens = row_splits[1:] - row_splits[:-1]
+
+ blank_id = self.decoder.blank_id
+ 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
+
+ lm = self.simple_lm_proj(decoder_out)
+ am = self.simple_am_proj(encoder_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=reduction,
+ 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),
+ 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=reduction,
+ )
+
+ return (simple_loss, pruned_loss)
diff --git a/egs/librispeech/ASR/lstm_transducer_stateless/optim.py b/egs/librispeech/ASR/lstm_transducer_stateless/optim.py
new file mode 120000
index 000000000..e2deb4492
--- /dev/null
+++ b/egs/librispeech/ASR/lstm_transducer_stateless/optim.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless2/optim.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/lstm_transducer_stateless/pretrained.py b/egs/librispeech/ASR/lstm_transducer_stateless/pretrained.py
new file mode 100755
index 000000000..2a6e2adc6
--- /dev/null
+++ b/egs/librispeech/ASR/lstm_transducer_stateless/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.
+"""
+Usage:
+
+(1) greedy search
+./lstm_transducer_stateless/pretrained.py \
+ --checkpoint ./lstm_transducer_stateless/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
+./lstm_transducer_stateless/pretrained.py \
+ --checkpoint ./lstm_transducer_stateless/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
+./lstm_transducer_stateless/pretrained.py \
+ --checkpoint ./lstm_transducer_stateless/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
+./lstm_transducer_stateless/pretrained.py \
+ --checkpoint ./lstm_transducer_stateless/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 `./lstm_transducer_stateless/exp/epoch-xx.pt`.
+
+Note: ./lstm_transducer_stateless/exp/pretrained.pt is generated by
+./lstm_transducer_stateless/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)
+
+ #
+## Introduction
+
+icefall contains ASR recipes for various datasets
+using