使用 torchaudio 的 MVDR

作者 Zhaoheng Ni

概述

这是一个关于如何使用 torchaudio.transforms.MVDR()应用MVDR波束成形的教程。

步骤

• 理想比率掩码（IRM）是通过将干净信号/噪声的幅度除以混合信号的幅度而生成的。

• 我们测试了 torchaudio 的 MVDR 模块的所有三种解决方案 (ref_channel, stv_evd, stv_power)。

• 我们针对 MVDR 波束成形测试了单通道和多通道掩码。在分别计算语音和噪声的协方差矩阵时，多通道掩码会沿通道维度进行平均。

准备

首先，我们导入必要的包并获取数据。

多通道音频示例选自 ConferencingSpeech 数据集。

原始文件名为

SSB07200001\#noise-sound-bible-0038\#7.86_6.16_3.00_3.14_4.84_134.5285_191.7899_0.4735\#15217\#25.16333303751458\#0.2101221178590021.wav

由以下生成；

• SSB07200001.wav 来自 AISHELL-3 (Apache License v.2.0)

• noise-sound-bible-0038.wav 来自 MUSAN（署名 4.0 国际 — CC BY 4.0）# noqa: E501

import os

import IPython.display as ipd
import requests
import torch
import torchaudio

torch.random.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print(torch.__version__)
print(torchaudio.__version__)
print(device)

filenames = [
    "mix.wav",
    "reverb_clean.wav",
    "clean.wav",
]
base_url = "https://download.pytorch.org/torchaudio/tutorial-assets/mvdr"

for filename in filenames:
    os.makedirs("_assets", exist_ok=True)
    if not os.path.exists(filename):
        with open(f"_assets/{filename}", "wb") as file:
            file.write(requests.get(f"{base_url}/{filename}").content)

Out:

1.11.0+cpu
0.11.0+cpu
cpu

生成理想比率掩码 (IRM)

加载音频数据

mix, sr = torchaudio.load("_assets/mix.wav")
reverb_clean, sr2 = torchaudio.load("_assets/reverb_clean.wav")
clean, sr3 = torchaudio.load("_assets/clean.wav")
assert sr == sr2

noise = mix - reverb_clean

注意

MVDR 模块需要 torch.cdouble dtype 用于含噪 STFT。 我们需要将波形的 dtype 转换为 torch.double

mix = mix.to(torch.double)
noise = noise.to(torch.double)
clean = clean.to(torch.double)
reverb_clean = reverb_clean.to(torch.double)

计算 STFT

stft = torchaudio.transforms.Spectrogram(
    n_fft=1024,
    hop_length=256,
    power=None,
)
istft = torchaudio.transforms.InverseSpectrogram(n_fft=1024, hop_length=256)

spec_mix = stft(mix)
spec_clean = stft(clean)
spec_reverb_clean = stft(reverb_clean)
spec_noise = stft(noise)

生成理想比率掩码 (IRM)

注意

我们发现直接使用掩码的效果优于使用其平方根。这与 IRM 的定义略有不同。

def get_irms(spec_clean, spec_noise):
    mag_clean = spec_clean.abs() ** 2
    mag_noise = spec_noise.abs() ** 2
    irm_speech = mag_clean / (mag_clean + mag_noise)
    irm_noise = mag_noise / (mag_clean + mag_noise)

    return irm_speech, irm_noise

注意

此处我们将混响清晰语音作为目标，您也可以将其设置为干声清晰语音。

irm_speech, irm_noise = get_irms(spec_reverb_clean, spec_noise)

应用 MVDR

使用多通道掩码应用 MVDR 波束成形

results_multi = {}
for solution in ["ref_channel", "stv_evd", "stv_power"]:
    mvdr = torchaudio.transforms.MVDR(ref_channel=0, solution=solution, multi_mask=True)
    stft_est = mvdr(spec_mix, irm_speech, irm_noise)
    est = istft(stft_est, length=mix.shape[-1])
    results_multi[solution] = est

使用单通道掩码应用 MVDR 波束成形

我们以第 1 个通道为例。 通道选择可能取决于麦克风阵列的设计。

results_single = {}
for solution in ["ref_channel", "stv_evd", "stv_power"]:
    mvdr = torchaudio.transforms.MVDR(ref_channel=0, solution=solution, multi_mask=False)
    stft_est = mvdr(spec_mix, irm_speech[0], irm_noise[0])
    est = istft(stft_est, length=mix.shape[-1])
    results_single[solution] = est

计算 Si-SDR 分数

def si_sdr(estimate, reference, epsilon=1e-8):
    estimate = estimate - estimate.mean()
    reference = reference - reference.mean()
    reference_pow = reference.pow(2).mean(axis=1, keepdim=True)
    mix_pow = (estimate * reference).mean(axis=1, keepdim=True)
    scale = mix_pow / (reference_pow + epsilon)

    reference = scale * reference
    error = estimate - reference

    reference_pow = reference.pow(2)
    error_pow = error.pow(2)

    reference_pow = reference_pow.mean(axis=1)
    error_pow = error_pow.mean(axis=1)

    sisdr = 10 * torch.log10(reference_pow) - 10 * torch.log10(error_pow)
    return sisdr.item()

结果

单通道掩码结果

for solution in results_single:
    print(solution + ": ", si_sdr(results_single[solution][None, ...], reverb_clean[0:1]))

Out:

ref_channel:  15.035907456985868
stv_evd:  16.56373467383255
stv_power:  17.820481909929903

多通道掩码结果

for solution in results_multi:
    print(solution + ": ", si_sdr(results_multi[solution][None, ...], reverb_clean[0:1]))

Out:

ref_channel:  13.177373866143256
stv_evd:  12.433610809532858
stv_power:  12.897505397104673

原始音频

混合语音

ipd.Audio(mix[0], rate=16000)

噪声

ipd.Audio(noise[0], rate=16000)

清晰的语音

ipd.Audio(clean[0], rate=16000)

增强音频

多通道掩码，ref_channel 解决方案

ipd.Audio(results_multi["ref_channel"], rate=16000)

多通道掩码，stv_evd 解决方案

ipd.Audio(results_multi["stv_evd"], rate=16000)

多通道掩码，stv_power 解决方案

ipd.Audio(results_multi["stv_power"], rate=16000)

单通道掩码，ref_channel 解决方案

ipd.Audio(results_single["ref_channel"], rate=16000)

单通道掩码，stv_evd 解决方案

ipd.Audio(results_single["stv_evd"], rate=16000)

单通道掩码，stv_power 解决方案

ipd.Audio(results_single["stv_power"], rate=16000)