本次开始,对WeNet的测试代码,进行逐行学习。

运行脚本

if[${stage}-le8]&&[${stop_stage}-ge8];then
  # Test model, please specify the model you want to test by --checkpoint
  cmvn_opts=
  $cmvn && cmvn_opts="--cmvn data/${train_set}/global_cmvn"
  mkdir -p $dir/test
  if [ ${average_checkpoint} == true ]; then
    decode_checkpoint=$dir/avg_${average_num}.pt
    echo "do model average and final checkpoint is $decode_checkpoint"
    python -m ipdb wenet/bin/average_model.py \
      --dst_model $decode_checkpoint \
      --src_path $dir  \
      --num ${average_num} \
      --val_best
  fi
  # Specify decoding_chunk_size if it's a unified dynamic chunk trained model
  # -1 for full chunk
  decoding_chunk_size=-1
  ctc_weight=0.5
  # Polling GPU id begin with index 0
  #num_gpus=$(echo $CUDA_VISIBLE_DEVICES | awk -F "," '{print NF}')
  num_gpus=1 #$(echo $CUDA_VISIBLE_DEVICES | awk -F "," '{print NF}')
  idx=0
  for test in $recog_set; do
    for mode in ${decode_modes}; do
    {
      {
        test_dir=$dir/${test}_${mode}
        mkdir -p $test_dir
        gpu_id=$(echo $CUDA_VISIBLE_DEVICES | cut -d',' -f$[$idx+1])
        python wenet/bin/recognize.py --gpu $gpu_id \
          --mode $mode \
          --config $dir/train.yaml \
          --data_type raw \
          --test_data $wave_data/$test/data.list \
          --checkpoint $decode_checkpoint \
          --beam_size 10 \
          --batch_size 1 \
          --penalty 0.0 \
          --dict $dict \
          --result_file $test_dir/text_bpe \
          --ctc_weight $ctc_weight \
          $cmvn_opts \
          ${decoding_chunk_size:+--decoding_chunk_size $decoding_chunk_size}

        cut -f2- -d " " $test_dir/text_bpe > $test_dir/text_bpe_value_tmp
        cut -f1 -d " " $test_dir/text_bpe > $test_dir/text_bpe_key_tmp
        tools/spm_decode --model=${bpemodel}.model --input_format=piece \
          < $test_dir/text_bpe_value_tmp | sed -e "s/▁/ /g" > $test_dir/text_value_tmp
        paste -d " " $test_dir/text_bpe_key_tmp $test_dir/text_value_tmp > $test_dir/text

        python tools/compute-wer.py --char=1 --v=1 \
          $wave_data/$test/text $test_dir/text > $test_dir/wer
      } #&

      ((idx+=1))
      if [ $idx -eq $num_gpus ]; then
        idx=0
      fi
    }
    done
  done
  #wait
fi

上面的脚本,涉及到两个部分,一个是对多个checkpoint进行平均,一个是循环多个decode模式来解码。

decode_modes="attention_rescoring ctc_greedy_search ctc_prefix_beam_search attention"

一共考察学习上面四种decode模式。

有个小坑

--config $dir/train.yaml \

这个一行,是在.pt保存路径上的那个配置文件!!

我的版本,如果指定用下面的【训练的时候用的配置文件】:

train_config=conf/train_conformer.yaml

会出错。。。因为train_config里面没有配置"cmvn_file",而recognize.py又不接受cmvn_opts。。。


checkpoint平均

>/workspace/asr/wenet/examples/csj/s0/wenet/bin/average_model.py(3)()

这个的逻辑相对容易理解一些:


接受的参数,以及基本的准备工作

然后就是读取目前的所有的checkpoints,按照cv_loss(在validation set上的loss),从小到大排序,然后就是取前面的10个(或者其他指定的数量)checkpoints,读取它们,并平均,最后保存:


多个checkpoint平均的整体流程

上面的流程还是比较容易理解的,使用torch.load, torch.true_devide,以及torch.save来搞定。


test

[/workspace/asr/wenet/examples/csj/s0/wenet/bin/recognize.py]

开始测试:


测试inference相关的代码

args

>/workspace/asr/wenet/examples/csj/s0/wenet/bin/recognize.py(108)main()
    107     args = get_args()
--> 108     logging.basicConfig(level=logging.DEBUG,
    109                         format='%(asctime)s %(levelname)s %(message)s')

ipdb> p args
Namespace(batch_size=1, beam_size=10, bpe_model=None, 
checkpoint='exp/sp_spec_aug_conformer_bidecoder_large/avg_10.pt', 
config='exp/sp_spec_aug_conformer_bidecoder_large/train.yaml', 
ctc_weight=0.5, data_type='raw', decoding_chunk_size=-1, 
dict='data/lang_char/train_bpe4096_units.txt', gpu=0, 
mode='attention_rescoring', non_lang_syms=None, 
num_decoding_left_chunks=-1, override_config=[], 
penalty=0.0, 
result_file='exp/sp_spec_aug_conformer_bidecoder_large/test1_attention_rescoring/text_bpe', 
reverse_weight=0.0, simulate_streaming=False, test_data='data/test1/data.list')

注意,输入的命令行参数里面,有beam_size=10

configs

自train.yaml读取。

ipdb>configs
{
'accum_grad': 4,
 'cmvn_file': 'data/train/global_cmvn', 
'dataset_conf': {
	'batch_conf': {
		'batch_size': 12,
		 'batch_type': 'static'},
	 'fbank_conf': {
		'dither': 1.0,
		 'frame_length': 25,
		 'frame_shift': 10,
		 'num_mel_bins': 80},
	 'filter_conf': {
		'max_length': 2000,
		 'max_output_input_ratio': 10.0,
		 'min_length': 50,
		 'min_output_input_ratio': 0.05,
		 'token_max_length': 400,
		 'token_min_length': 1},
	 'resample_conf': {
		'resample_rate': 16000},
 	'shuffle': True,
	 'shuffle_conf': {
		'shuffle_size': 1500},
 	'sort': True,
	 'sort_conf': {
		'sort_size': 500},
 	'spec_aug': True,
	 'spec_aug_conf': {
		'max_f': 10,
		 'max_t': 50,
		 'num_f_mask': 2,
		 'num_t_mask': 3},
 	'speed_perturb': True},


 'decoder': 'bitransformer',
 'decoder_conf': {
	'attention_heads': 8,
	 'dropout_rate': 0.1,
	 'linear_units': 2048,
	 'num_blocks': 3,
	 'positional_dropout_rate': 0.1,
	 'r_num_blocks': 3,
	 'self_attention_dropout_rate': 0.1,
	 'src_attention_dropout_rate': 0.1},

 'encoder': 'conformer',
 'encoder_conf': {
	'activation_type': 'swish',
	 'attention_dropout_rate': 0.1,
	 'attention_heads': 8,
	 'cnn_module_kernel': 31,
	 'cnn_module_norm': 'layer_norm',
	 'dropout_rate': 0.1,
	 'input_layer': 'conv2d',
	 'linear_units': 2048,
	 'normalize_before': True,
	 'num_blocks': 12,
	 'output_size': 512,
	 'pos_enc_layer_type': 'rel_pos',
	 'positional_dropout_rate': 0.1,
	 'selfattention_layer_type': 'rel_selfattn',
	 'use_cnn_module': True},

 'grad_clip': 5,
 'input_dim': 80,
 'is_json_cmvn': True,
 'log_interval': 100,
 'max_epoch': 200,

 'model_conf': {
	'ctc_weight': 0.3,
	 'length_normalized_loss': False,
	 'lsm_weight': 0.1,
	 'reverse_weight': 0.3},

 'optim': 'adam',

 'optim_conf': {
	'lr': 0.002},
	 'output_dim': 5502,
	 'scheduler': 'warmuplr',
	 'scheduler_conf': {
		'warmup_steps': 50000}
}

test_conf

ipdb>test_conf
{'batch_conf': {'batch_size': 1, 'batch_type': 'static'}, 
'fbank_conf': {'dither': 0.0, 'frame_length': 25, 'frame_shift': 10, 'num_mel_bins': 80}, 
'filter_conf': {'max_length': 102400, 'max_output_input_ratio': 102400, 
  'min_length': 0, 'min_output_input_ratio': 0, 'token_max_length': 102400, 'token_min_length': 0}, 
'resample_conf': {'resample_rate': 16000}, 
'shuffle': False, 'shuffle_conf': {'shuffle_size': 1500}, 
'sort': False, 'sort_conf': {'sort_size': 500}, 
'spec_aug': False, 

'spec_aug_conf': {'max_f': 10, 'max_t': 50, 'num_f_mask': 2, 'num_t_mask': 3}, 
'speed_perturb': False}

读取test data,并且构造data loader的代码,和train的时候 train data相通,这里不再冗述。


目前的位置,读取一个batch,并且开始解码!

one batch test

这里为了简单,以及使用beam search,设定batch size = 1

>/workspace/asr/wenet/examples/csj/s0/wenet/bin/recognize.py(173)main()
    172             keys, feats, target, feats_lengths, target_lengths = batch
--> 173             feats = feats.to(device)
    174             target = target.to(device)

ipdb> keys.shape
*** AttributeError: 'list' object has no attribute 'shape'
ipdb> keys
['A03M0156_00000.612_00002.674']
ipdb> feats.shape
torch.Size([1, 204, 80])
ipdb> target.shape
torch.Size([1, 12])
ipdb> feats_lengths
tensor([204], dtype=torch.int32)
ipdb> target_lengths
tensor([12], dtype=torch.int32)
ipdb>

keys是当前的wav相关的id识别码。【输入数据会准备它】

feats是经过kaldi.fbank处理之后的wav的表示,这里是1=batch, 204=frame number, 80=fbank维度(类似于一个frame被表示成一个80维度的向量)。

target是id化的目标文本序列;

feats_lengths=[204],输入 wav batch的长度列表(本batch);

target_lengths=[12],输出文本序列的长度列表(本batch)。

attention_rescoring

这次关注的是attention rescoring解码算法。

思想:

  1. 基于ctc prefix beam search来获取top-10(beam size=10)文本候选+得分 score1;
  2. 每个候选和source wav都扔给decoder,打分score2,然后打分和ctc的得分加权相加:w1 * score1 + w2 *score2。
  3. 输出top-1结果。


attention_rescoring算法相关的逻辑

关于model.attention_rescoring()算法的细节脑图为:


model.attention_rescoring函数的细节脑图

上面函数中,重要的有三个部分:

  1. self._ctc_prefix_beam_search;
  2. self.decoder
  3. for i, hyp部分的,ctc解码得分,和decoder解码得分,加权相加。

准备部分


直接来自args.beam_size是None,不过来自输入的命令行的beam_size=10.

注意,输入的命令行参数里面,有beam_size=10。

不过reverse_weight=0.0这个,后续的确需要尝试其他值。

1 _ctc_prefix_beam_search


ctc prefix beam search的脑图

上面这段脑图,直接爽死。

主要是,对wav,使用model encoder编码,然后调用ctc_lo来把512维度改成5502维度,从而每个frame都有自己的候选词列表以及概率了!

然后,搞成真正的softmax -> log,即对数概率的形式;

之后就是迷死个人的beam search了,这个需要结合ctc解码,以及ctc论文来综合解读。

1.1 _forward_encoder

这块其实和train的时候,对于输入wav frames的基于ConformerEncoderLayers的表示处理,相同:


上面的脑图,也就是绘制了一部分,感觉足够了。

返回的结果:

  1. encoder_out: [1, 50, 512],是对50个frames的512维度的向量化。
  2. encoder_mask,[1, 1, 50],因为当前batch.size=1,所以这个mask其实没啥用,都是True。

1.2 ctc.log_softmax


这里有个线性层,负责把一个frame的512维度的向量,映射成5502-词表大小。从而当前frame就有自己的对应的候选词,以及候选词的概率了!!!

1.3 beam search

涉及到的代码为:


有嚼头的ctc decoding方法,beam size=10

上面这块的细节,特别是区分当前字符的三种情况(空,和前一个相同,其他),可以参考原始的ctc论文。

1.3.1 t=0..49

我们跟着走两步,走两步,走两步:


这是对50个frames循环,然后根据beam size=10来构造候选


对于当前frame,我们看它的10个得分最高的候选token

1.3.2 s=0

进去看看s=0的情况:


根据s=0来扩展当前的cur_hyps中的所有候选

这个是严格按照ctc论文中定义的更新方法来的。

1.3.3 s=1554


根据s=1来扩展当前的cur_hyps中的所有候选

因为s!=0,而且s和prefix的当前的最后一个token不同的时候的更新。

1.3.4 s=1396


根据s=2来扩展当前的cur_hyps中的所有候选

以此类推。就可以搞定了。

输出

>/workspace/asr/wenet/wenet/transformer/asr_model.py(413)_ctc_prefix_beam_search()
    412             cur_hyps = next_hyps[:beam_size]
1-> 413         hyps = [(y[0], log_add([y[1][0], y[1][1]])) for y in cur_hyps]
    414         return hyps, encoder_out

ipdb> cur_hyps
[((1396, 1396, 1396, 1396, 1396, 1396, 1885, 1670), (-15.889783398806042, -21.83005307378142)), 
((1396, 1396, 1396, 1396, 1396, 1396, 1885), (-15.996046160775117, -inf)), 
((1396, 1396, 1396, 1396, 1396, 1396, 1396, 1670), (-16.10800202701202, -22.039127286713356)), 
((1396, 1396, 1396, 1396, 1396, 1396, 1396, 1885, 1670), 
(-16.112852794392175, -22.05312246936756)), 

((1396, 1396, 1396, 1396, 1396, 1885, 1670), (-16.12371970252195, -22.06398937749733)), 
((1396, 1396, 1396, 1396, 1396, 1396, 1741, 1670), (-16.169015181751966, -inf)), 
((1396, 1396, 1396, 1396, 1396, 1396, 1670), (-16.19584450996598, -inf)), 
((1396, 1396, 1396, 1396, 1396, 1396, 1396), (-16.20512037370705, -inf)), 
((1396, 1396, 1396, 1396, 1396, 1396, 1396, 1885), (-16.219115556361253, -inf)), 
((1396, 1396, 1396, 1396, 1396, 1885), (-16.229982464491023, -inf))]

上面,是执行了maxlen=50次大的循环之后,得到的hyps的内容。

然后,进一步把一个候选序列的两个对数概率相加(上面Line 413里面的,log_add!),就得到了一个候选序列的完整的概率:

[((1396,1396,1396,1396,1396,1396,1885,1670),-15.887155534713223),
((1396, 1396, 1396, 1396, 1396, 1396, 1885), -15.996046160775117), 
((1396, 1396, 1396, 1396, 1396, 1396, 1396, 1670), -16.105350054413357), 
((1396, 1396, 1396, 1396, 1396, 1396, 1396, 1885, 1670), -16.110224930299356), 
((1396, 1396, 1396, 1396, 1396, 1885, 1670), -16.12109183842913), 
((1396, 1396, 1396, 1396, 1396, 1396, 1741, 1670), -16.169015181751966), 
((1396, 1396, 1396, 1396, 1396, 1396, 1670), -16.19584450996598), 
((1396, 1396, 1396, 1396, 1396, 1396, 1396), -16.20512037370705), 
((1396, 1396, 1396, 1396, 1396, 1396, 1396, 1885), -16.219115556361253), 
((1396, 1396, 1396, 1396, 1396, 1885), -16.229982464491023)]

上面这个列表,就是_ctc_prefix_beam_search的第一个输出。

第二个输出为:encoder_out.shape = torch.Size([1, 50, 512])

准备decoder


pad刚才ctc的输出,并准备decoder的调用

pad_sequence(tgt)

pad的是-1。这是beam=10的例子:

>/workspace/asr/wenet/wenet/transformer/asr_model.py(499)attention_rescoring()
    498         ], True, self.ignore_id)  # (beam_size, max_hyps_len)
3-> 499         ori_hyps_pad = hyps_pad
    500         hyps_lens = torch.tensor([len(hyp[0]) for hyp in hyps],

ipdb> hyps_pad
tensor([[1396, 1396, 1396, 1396, 1396, 1396, 1885, 1670,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1885,   -1,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1396, 1670,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1396, 1885, 1670],
        [1396, 1396, 1396, 1396, 1396, 1885, 1670,   -1,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1741, 1670,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1670,   -1,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1396,   -1,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1396, 1885,   -1],
        [1396, 1396, 1396, 1396, 1396, 1885,   -1,   -1,   -1]],
       device='cuda:0')

add_sos_eos

ipdb>hyps_pad# shape=(10, 9)
tensor([[1396, 1396, 1396, 1396, 1396, 1396, 1885, 1670,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1885,   -1,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1396, 1670,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1396, 1885, 1670],
        [1396, 1396, 1396, 1396, 1396, 1885, 1670,   -1,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1741, 1670,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1670,   -1,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1396,   -1,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1396, 1885,   -1],
        [1396, 1396, 1396, 1396, 1396, 1885,   -1,   -1,   -1]],
       device='cuda:0')
ipdb> n
> /workspace/asr/wenet/wenet/transformer/asr_model.py(504)attention_rescoring()
    503         hyps_pad, _ = add_sos_eos(hyps_pad, self.sos, self.eos, self.ignore_id)
--> 504         hyps_lens = hyps_lens + 1  # Add  at begining
    505         encoder_out = encoder_out.repeat(beam_size, 1, 1)

ipdb> hyps_pad # shape=(10,10)
tensor([[5501, 1396, 1396, 1396, 1396, 1396, 1396, 1885, 1670, 5501],
        [5501, 1396, 1396, 1396, 1396, 1396, 1396, 1885, 5501, 5501],
        [5501, 1396, 1396, 1396, 1396, 1396, 1396, 1396, 1670, 5501],
        [5501, 1396, 1396, 1396, 1396, 1396, 1396, 1396, 1885, 1670],
        [5501, 1396, 1396, 1396, 1396, 1396, 1885, 1670, 5501, 5501],
        [5501, 1396, 1396, 1396, 1396, 1396, 1396, 1741, 1670, 5501],
        [5501, 1396, 1396, 1396, 1396, 1396, 1396, 1670, 5501, 5501],
        [5501, 1396, 1396, 1396, 1396, 1396, 1396, 1396, 5501, 5501],
        [5501, 1396, 1396, 1396, 1396, 1396, 1396, 1396, 1885, 5501],
        [5501, 1396, 1396, 1396, 1396, 1396, 1885, 5501, 5501, 5501]],
       device='cuda:0')

reverse_pad_list

逆序搞一次pad

>/workspace/asr/wenet/wenet/transformer/asr_model.py(513)attention_rescoring()
    512         r_hyps_pad = reverse_pad_list(ori_hyps_pad, hyps_lens, self.ignore_id)
--> 513         r_hyps_pad, _ = add_sos_eos(r_hyps_pad, self.sos, self.eos,
    514                                     self.ignore_id)

ipdb> ori_hyps_pad
tensor([[1396, 1396, 1396, 1396, 1396, 1396, 1885, 1670,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1885,   -1,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1396, 1670,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1396, 1885, 1670],
        [1396, 1396, 1396, 1396, 1396, 1885, 1670,   -1,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1741, 1670,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1670,   -1,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1396,   -1,   -1],
        [1396, 1396, 1396, 1396, 1396, 1396, 1396, 1885,   -1],
        [1396, 1396, 1396, 1396, 1396, 1885,   -1,   -1,   -1]],
       device='cuda:0')
ipdb> r_hyps_pad
tensor([[  -1, 1670, 1885, 1396, 1396, 1396, 1396, 1396, 1396],
        [  -1, 1885, 1396, 1396, 1396, 1396, 1396, 1396,   -1],
        [  -1, 1670, 1396, 1396, 1396, 1396, 1396, 1396, 1396],
        [1670, 1885, 1396, 1396, 1396, 1396, 1396, 1396, 1396],
        [  -1, 1670, 1885, 1396, 1396, 1396, 1396, 1396,   -1],
        [  -1, 1670, 1741, 1396, 1396, 1396, 1396, 1396, 1396],
        [  -1, 1670, 1396, 1396, 1396, 1396, 1396, 1396,   -1],
        [  -1, 1396, 1396, 1396, 1396, 1396, 1396, 1396,   -1],
        [  -1, 1885, 1396, 1396, 1396, 1396, 1396, 1396, 1396],
        [  -1, 1885, 1396, 1396, 1396, 1396, 1396,   -1,   -1]],
       device='cuda:0', dtype=torch.int32)
ipdb>

add_sos_eos

ipdb>r_hyps_pad
tensor([[5501, 1670, 1885, 1396, 1396, 1396, 1396, 1396, 1396, 5501],
        [5501, 1885, 1396, 1396, 1396, 1396, 1396, 1396, 5501, 5501],
        [5501, 1670, 1396, 1396, 1396, 1396, 1396, 1396, 1396, 5501],
        [5501, 1670, 1885, 1396, 1396, 1396, 1396, 1396, 1396, 1396],
        [5501, 1670, 1885, 1396, 1396, 1396, 1396, 1396, 5501, 5501],
        [5501, 1670, 1741, 1396, 1396, 1396, 1396, 1396, 1396, 5501],
        [5501, 1670, 1396, 1396, 1396, 1396, 1396, 1396, 5501, 5501],
        [5501, 1396, 1396, 1396, 1396, 1396, 1396, 1396, 5501, 5501],
        [5501, 1885, 1396, 1396, 1396, 1396, 1396, 1396, 1396, 5501],
        [5501, 1885, 1396, 1396, 1396, 1396, 1396, 5501, 5501, 5501]],
       device='cuda:0')
ipdb>

如此,就可以调用decoder了。

2 decoder


这个和train里面的decoder相同。使用的decoder input是来自ctc prefix beam search得到的!

这个decoder的输出是:

  1. decoder_out, [10, 10, 5502];
  2. r_decoder_out = tensor(0.)
  3. olens output lengths

其中olens为:

olens-没有用

在left_decoder和right_decoder之后:【不过它没啥用】

tensor([[10,9,8,7,6,5,4,3,2,0],
        [10,  9,  8,  7,  6,  5,  4,  3,  0,  0],
        [10,  9,  8,  7,  6,  5,  4,  3,  2,  0],
        [10,  9,  8,  7,  6,  5,  4,  3,  2,  1],
        [10,  9,  8,  7,  6,  5,  4,  3,  0,  0],
        [10,  9,  8,  7,  6,  5,  4,  3,  2,  0],
        [10,  9,  8,  7,  6,  5,  4,  3,  0,  0],
        [10,  9,  8,  7,  6,  5,  4,  3,  0,  0],
        [10,  9,  8,  7,  6,  5,  4,  3,  2,  0],
        [10,  9,  8,  7,  6,  5,  4,  0,  0,  0]], device='cuda:0')

3 得分加权相加


这是基于decoder的输出概率和ctc beam search的概率,加权叠加的情况

最终结果收集

如此,就可以最终确定候选文本了。


again,写入结果文件。