跳转至

SepFormer

1 概述

1.1 背景介绍

语音分离是指从混合的音频信号中提取出独立的语音源, 例如, 在多人会议录音中, 分离出每个说话人的声音。 本次开源的语音分离算法来源于SpeechBrain发布的Pytorch开源工具包, 该工具包集成了多种先进的音频处理技术, 详情可参考Speechbrain官方说明:

https://github.com/speechbrain/speechbrain

本次采用的开源模型是基于Libri3Mix数据集训练而成的, 因此该模型可分离出三人混声音频, 开源模型地址如下:

https://huggingface.co/speechbrain/sepformer-libri3mix/tree/main

测试音频下载地址:

https://huggingface.co/speechbrain/sepformer-wsj03mix/tree/main

1.2 使用说明

Linux SDK-alkaid中默认带了已经预先转换好的离线模型及板端示例, 相关文件路径如下:

  • 板端示例程序路径

    Linux_SDK/sdk/verify/opendla/source/separation
    
  • 板端离线模型路径

    Linux_SDK/project/board/${chip}/dla_file/ipu_open_models/separation/separation_lib3mix_sim.img
    
  • 板端测试音频路径

    Linux_SDK/sdk/verify/opendla/source/resource/item0_mix.wav
    

如果用户不需要转换模型可直接跳转至第3章节。

2 模型转换

2.1 onnx模型转换

  • python环境搭建

    $conda create -n speechbrain python==3.9
    $conda activate speechbrain
    $git clone https://github.com/speechbrain/speechbrain.git
    $cd speechbrain
    $pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
    $pip install --editable .
    

    注意:这里提供的python环境搭建, 仅作为参考示例, 具体搭建过程请参考官方源码运行教程:

    https://github.com/speechbrain/speechbrain/blob/develop/README.md
    
  • 模型测试

    • 编写模型测试脚本inference.py

      from speechbrain.pretrained import SepformerSeparation as separator
      import torchaudio
      
      model = separator.from_hparams(source="speechbrain/sepformer-libri3mix", savedir='speechbrain/sepformer-libri3mix')
      
      est_sources = model.separate_file(path='speechbrain/sepformer-wsj03mix/test_mixture_3spks.wav')
      
      torchaudio.save("source1hat.wav", est_sources[:, :, 0].detach().cpu(), 8000)
      torchaudio.save("source2hat.wav", est_sources[:, :, 1].detach().cpu(), 8000)
      torchaudio.save("source3hat.wav", est_sources[:, :, 2].detach().cpu(), 8000)
      
    • 运行模型测试脚本, 确保speechbrain环境配置正确。

      $python inference.py
      
  • 模型导出

    • 安装依赖库

      $pip install onnx -i https://pypi.tuna.tsinghua.edu.cn/simple
      $pip install onnx-simplifier -i https://pypi.tuna.tsinghua.edu.cn/simple
      
    • 编写模型转换脚本export_onnx.py

      import os
      import sys
      sys.path.append(os.getcwd())
      import torch
      import  numpy as np
      
      from speechbrain.inference.separation import SepformerSeparation as separator
      from speechbrain.utils.fetching import fetch
      from speechbrain.utils.data_utils import split_path
      import torchaudio
      
      import onnx
      import onnxsim
      import onnxruntime as ort
      
      model = separator.from_hparams(source="speechbrain/sepformer-libri3mix", savedir='speechbrain/sepformer-libri3mix')
      
      class SeparationModel(torch.nn.Module):
          def __init__(self, encoder, masknet, decoder):
          super().__init__()
          self.encoder = encoder
          self.masknet = masknet
          self.decoder = decoder
      
          def forward(self, wav):
          source = self.encoder(wav)
          print(source.size())
          source_mask = self.masknet(source)
          print(source_mask.size())
          source = torch.stack([source] * 3)
          sep_h = source * source_mask
      
          # Decoding
          est_source = torch.cat(
                  [
                      self.decoder(sep_h[i]).unsqueeze(-1)
                      for i in range(3)
                  ],
                  dim=-1,
          )
          print(est_source.size())
      
          T_origin = wav.size(1)
          T_est = est_source.size(1)
          if T_origin > T_est:
                  est_source = F.pad(est_source, (0, 0, 0, T_origin - T_est))
          else:
                  est_source = est_source[:, :T_origin, :]
      
          return est_source
      
      if __name__ == '__main__':
          source, fl = split_path('speechbrain/sepformer-wsj03mix/test_mixture_3spks.wav')
          path = fetch(fl, source="speechbrain/sepformer-libri3mix", savedir='speechbrain/sepformer-libri3mix')
      
          batch, fs_file = torchaudio.load(path)
          fs_model = 8000
      
          # resample the data if needed
          if fs_file != fs_model:
          print(
                  "Resampling the audio from {} Hz to {} Hz".format(
                      fs_file, fs_model
                  )
          )
          tf = torchaudio.transforms.Resample(
                  orig_freq=fs_file, new_freq=fs_model
          )
          batch = batch.mean(dim=0, keepdim=True)
          batch = tf(batch)
      
          mod = batch.shape[1] % 8000
          batch = torch.cat([batch, torch.zeros((1, 8000 - mod))],dim=1)
          div = batch.shape[1] // 8000
          sub_model = SeparationModel(model.mods.encoder, model.mods.masknet, model.mods.decoder)
      
          result_list = []
          onnx_path = './opendla/separation_lib3mix.onnx'
          for i in range(div):
      
          if i == div-1:
                  input_tensor = batch[:, i*8000:]
          else:
                  input_tensor = batch[:, i*8000:(i+1)*8000]
      
          est_sources = sub_model(input_tensor)
      
          result_list.append(est_sources)
      
          torch.onnx.export(
                  sub_model,
                  (input_tensor),
                  onnx_path,
                  export_params=True,
                  opset_version=13,
                  do_constant_folding=True,
                  input_names=["audio"],
                  output_names=["probs"],
                  verbose=False,
          )
      
          model_onnx = onnx.load(onnx_path)  # load onnx model
          onnx.checker.check_model(model_onnx)  # check onnx model
          model_onnx, check = onnxsim.simplify(model_onnx)
          onnx.save(model_onnx, onnx_path.replace('separation','separation_sim'))
          exit(1)
      
    • 运行模型转换脚本export_onnx.py

      python export_onnx.py
      

2.2 离线模型转换

2.2.1 预&后处理说明

  • 预处理 语音输入到模型之前, 通常需要通过torchaudio.load接口将音频数据转换为适合模型输入的张量。然后, 为了固定输入长度, 需要对返回的张量数据做截取或者补零操作。最后再对张量数据做归一化处理。模型输入信息如下所示:
  • 后处理 该语音分离模型没有后处理操作, 获取到模型输出数据后, 可直接根据输出的通道数量分别保存不同说话人的音频。模型输出信息如下所示:

2.2.2 offline模型转换流程

注意:1)OpenDLAModel对应的是压缩包image-dev_model_convert.tar解压之后的smodel文件。2)转换命令需要在docker环境下运行, 请先根据Docker开发环境教程, 加载SGS Docker环境。

  • 拷贝onnx模型到转换代码目录

    $cp speechbrain/sepformer-libri3mix/separation_sim_lib3mix.onnx OpenDLAModel/separation/sepformer/onnx
    
  • 转换命令

    $cd IPU_SDK_Release/docker
    $bash run_docker.sh
    #进入到docker环境下的OpenDLAModel目录
    $cd /work/SGS_XXX/OpenDLAModel
    $bash convert.sh -a separation/sepformer -c config/separation_sepformer.cfg -p SGS_IPU_Toolchain(绝对路径) -s false
    
  • 最终生成的模型地址

    output/${chip}_${时间}/separation_sim_lib3mix.img
    output/${chip}_${时间}/separation_sim_lib3mix_fixed.sim
    output/${chip}_${时间}/separation_sim_lib3mix_float.sim
    

2.2.3 关键脚本参数解析

-   input_config.ini

        [INPUT_CONFIG]
        inputs=audio;                               #onnx 输入节点名称, 如果有多个需以“,”隔开;
        training_input_formats=RAWDATA_F32_NHWC;    #模型训练数据的格式
        input_formats=RAWDATA_F32_NHWC;             #板端输入格式, 可以根据onnx的输入格式选择, 例如float:RAWDATA_F32_NHWC, int32:RAWDATA_S16_NHWC;
        quantizations=TRUE;                         #打开输入量化, 不需要修改;

        [OUTPUT_CONFIG]
        outputs=probs;                              #onnx 输出节点名称, 如果有多个需以“,”隔开;
        dequantizations=FALSE;                      #是否开启反量化, 根据实际需求填写, 建议为TRUE。设为False, 输出为int16; 设为True, 输出为float32

-   separation_sepfofmer.cfg

        [SEPFORMER]
        CHIP_LIST=pcupid                            #平台名称, 必须和板端平台一致, 否则模型无法运行
        Model_LIST=separation_sim_lib3mix           #输入onnx模型名称
        INPUT_SIZE_LIST=0                           #模型输入分辨率
        INPUT_INI_LIST=input_config.ini             #配置文件
        CLASS_NUM_LIST=0                            #填0即可
        SAVE_NAME_LIST=separation_sim_lib3mix.img   #输出模型名称
        QUANT_DATA_PATH=image_lists.txt             #量化数据路径

2.3 模型仿真

  • 获取float/fixed/offline模型输出

    $bash convert.sh -a separation/sepformer -c config/separation_sepformer.cfg -p SGS_IPU_Toolchain(绝对路径) -s true
    

    执行上述命令后, 会默认将float模型的输出tensor保存到separation/sepformer/log/output路径下的txt文件中。此外, 在separation/sepformer/convert.sh脚本中也提供了fixedoffline的仿真示例, 用户在运行时可以通过打开注释代码块, 分别获取fixedoffline模型输出。

  • 模型精度对比

    在保证输入和上述模型相同的情况下, 进入2.1章节搭建好的环境, 在inference.py文件中第8行处添加打印:

    print(est_sources)
    

    获取pytorch模型对应节点的输出tensor, 进而和float、fixed、offline模型进行对比。此外需要特别注意的是, 原始模型的输出格式是NCHW, 而float/fixed/offline模型输出的格式是NHWC

3 板端部署

3.1 程序编译

示例程序编译之前需要先根据板子(nand/nor/emmc, ddr型号等)选择deconfig进行sdk整包编译, 具体可以参考alkaid sdk sigdoc《开发环境搭建》文档。

  • 编译板端sepformer示例。

    $cd sdk/verify/opendla
    $make clean && make source/separation/sepformer -j8
    
  • 最终生成的可执行文件地址

    sdk/verify/opendla/out/${AARCH}/app/prog_separation_sepformer
    

3.2 运行文件

运行程序时, 需要先将以下几个文件拷贝到板端

  • prog_separation_sepformer
  • item0_mix.wav
  • separation_sim_lib3mix.img

3.3 运行说明

  • Usage: ./prog_separation_sepformer wav model(执行文件使用命令)

  • Required Input:

    • wav: 音频文件
    • model: 模型文件
  • Typical Output:

    ./prog_separation_sepformer item0_mix.wav models/separation_sim_lib3mix.img
    
        client [830] connected, module:ipu
        invoke time: 1064.185000 ms
        invoke time: 2128.087000 ms
        invoke time: 3192.150000 ms
        invoke time: 4255.845000 ms
        invoke time: 5320.448000 ms
        all invoke time: 5320.547000 ms
        WAV file 'spk_0.wav' has been written
        WAV file 'spk_1.wav' has been written
        WAV file 'spk_2.wav' has been written
        ------shutdown IPU0------
        client [830] disconnected, module:ipu