From 18b1449d1ff06c469e54190508c4f6be05c73d85 Mon Sep 17 00:00:00 2001
From: 夜雨飘零 <yeyupiaoling@foxmail.com>
Date: 星期二, 05 十二月 2023 22:04:14 +0800
Subject: [PATCH] 分角色语音识别支持更多的模型
---
funasr/utils/speaker_utils.py | 650 ----------------------------------------------------------
1 files changed, 6 insertions(+), 644 deletions(-)
diff --git a/funasr/utils/speaker_utils.py b/funasr/utils/speaker_utils.py
index df3eca7..a3eebf9 100644
--- a/funasr/utils/speaker_utils.py
+++ b/funasr/utils/speaker_utils.py
@@ -1,25 +1,18 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
""" Some implementations are adapted from https://github.com/yuyq96/D-TDNN
"""
-import math
-
-import torch
-import torch.nn.functional as F
-import torch.utils.checkpoint as cp
-from torch import nn
import io
-import os
-from typing import Any, Dict, List, Union
+from typing import Union
-import numpy as np
import librosa as sf
+import numpy as np
import torch
-import torchaudio
-import logging
-from funasr.utils.modelscope_file import File
-from collections import OrderedDict
+import torch.nn.functional as F
import torchaudio.compliance.kaldi as Kaldi
+from torch import nn
+
+from funasr.utils.modelscope_file import File
def check_audio_list(audio: list):
@@ -104,230 +97,6 @@
return segs
-class BasicResBlock(nn.Module):
- expansion = 1
-
- def __init__(self, in_planes, planes, stride=1):
- super(BasicResBlock, self).__init__()
- self.conv1 = nn.Conv2d(
- in_planes,
- planes,
- kernel_size=3,
- stride=(stride, 1),
- padding=1,
- bias=False)
- self.bn1 = nn.BatchNorm2d(planes)
- self.conv2 = nn.Conv2d(
- planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
- self.bn2 = nn.BatchNorm2d(planes)
-
- self.shortcut = nn.Sequential()
- if stride != 1 or in_planes != self.expansion * planes:
- self.shortcut = nn.Sequential(
- nn.Conv2d(
- in_planes,
- self.expansion * planes,
- kernel_size=1,
- stride=(stride, 1),
- bias=False), nn.BatchNorm2d(self.expansion * planes))
-
- def forward(self, x):
- out = F.relu(self.bn1(self.conv1(x)))
- out = self.bn2(self.conv2(out))
- out += self.shortcut(x)
- out = F.relu(out)
- return out
-
-
-class FCM(nn.Module):
-
- def __init__(self,
- block=BasicResBlock,
- num_blocks=[2, 2],
- m_channels=32,
- feat_dim=80):
- super(FCM, self).__init__()
- self.in_planes = m_channels
- self.conv1 = nn.Conv2d(
- 1, m_channels, kernel_size=3, stride=1, padding=1, bias=False)
- self.bn1 = nn.BatchNorm2d(m_channels)
-
- self.layer1 = self._make_layer(
- block, m_channels, num_blocks[0], stride=2)
- self.layer2 = self._make_layer(
- block, m_channels, num_blocks[0], stride=2)
-
- self.conv2 = nn.Conv2d(
- m_channels,
- m_channels,
- kernel_size=3,
- stride=(2, 1),
- padding=1,
- bias=False)
- self.bn2 = nn.BatchNorm2d(m_channels)
- self.out_channels = m_channels * (feat_dim // 8)
-
- def _make_layer(self, block, planes, num_blocks, stride):
- strides = [stride] + [1] * (num_blocks - 1)
- layers = []
- for stride in strides:
- layers.append(block(self.in_planes, planes, stride))
- self.in_planes = planes * block.expansion
- return nn.Sequential(*layers)
-
- def forward(self, x):
- x = x.unsqueeze(1)
- out = F.relu(self.bn1(self.conv1(x)))
- out = self.layer1(out)
- out = self.layer2(out)
- out = F.relu(self.bn2(self.conv2(out)))
-
- shape = out.shape
- out = out.reshape(shape[0], shape[1] * shape[2], shape[3])
- return out
-
-
-class CAMPPlus(nn.Module):
-
- def __init__(self,
- feat_dim=80,
- embedding_size=192,
- growth_rate=32,
- bn_size=4,
- init_channels=128,
- config_str='batchnorm-relu',
- memory_efficient=True,
- output_level='segment'):
- super(CAMPPlus, self).__init__()
-
- self.head = FCM(feat_dim=feat_dim)
- channels = self.head.out_channels
- self.output_level = output_level
-
- self.xvector = nn.Sequential(
- OrderedDict([
- ('tdnn',
- TDNNLayer(
- channels,
- init_channels,
- 5,
- stride=2,
- dilation=1,
- padding=-1,
- config_str=config_str)),
- ]))
- channels = init_channels
- for i, (num_layers, kernel_size, dilation) in enumerate(
- zip((12, 24, 16), (3, 3, 3), (1, 2, 2))):
- block = CAMDenseTDNNBlock(
- num_layers=num_layers,
- in_channels=channels,
- out_channels=growth_rate,
- bn_channels=bn_size * growth_rate,
- kernel_size=kernel_size,
- dilation=dilation,
- config_str=config_str,
- memory_efficient=memory_efficient)
- self.xvector.add_module('block%d' % (i + 1), block)
- channels = channels + num_layers * growth_rate
- self.xvector.add_module(
- 'transit%d' % (i + 1),
- TransitLayer(
- channels, channels // 2, bias=False,
- config_str=config_str))
- channels //= 2
-
- self.xvector.add_module('out_nonlinear',
- get_nonlinear(config_str, channels))
-
- if self.output_level == 'segment':
- self.xvector.add_module('stats', StatsPool())
- self.xvector.add_module(
- 'dense',
- DenseLayer(
- channels * 2, embedding_size, config_str='batchnorm_'))
- else:
- assert self.output_level == 'frame', '`output_level` should be set to \'segment\' or \'frame\'. '
-
- for m in self.modules():
- if isinstance(m, (nn.Conv1d, nn.Linear)):
- nn.init.kaiming_normal_(m.weight.data)
- if m.bias is not None:
- nn.init.zeros_(m.bias)
-
- def forward(self, x):
- x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
- x = self.head(x)
- x = self.xvector(x)
- if self.output_level == 'frame':
- x = x.transpose(1, 2)
- return x
-
-
-def get_nonlinear(config_str, channels):
- nonlinear = nn.Sequential()
- for name in config_str.split('-'):
- if name == 'relu':
- nonlinear.add_module('relu', nn.ReLU(inplace=True))
- elif name == 'prelu':
- nonlinear.add_module('prelu', nn.PReLU(channels))
- elif name == 'batchnorm':
- nonlinear.add_module('batchnorm', nn.BatchNorm1d(channels))
- elif name == 'batchnorm_':
- nonlinear.add_module('batchnorm',
- nn.BatchNorm1d(channels, affine=False))
- else:
- raise ValueError('Unexpected module ({}).'.format(name))
- return nonlinear
-
-
-def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, eps=1e-2):
- mean = x.mean(dim=dim)
- std = x.std(dim=dim, unbiased=unbiased)
- stats = torch.cat([mean, std], dim=-1)
- if keepdim:
- stats = stats.unsqueeze(dim=dim)
- return stats
-
-
-class StatsPool(nn.Module):
-
- def forward(self, x):
- return statistics_pooling(x)
-
-
-class TDNNLayer(nn.Module):
-
- def __init__(self,
- in_channels,
- out_channels,
- kernel_size,
- stride=1,
- padding=0,
- dilation=1,
- bias=False,
- config_str='batchnorm-relu'):
- super(TDNNLayer, self).__init__()
- if padding < 0:
- assert kernel_size % 2 == 1, 'Expect equal paddings, but got even kernel size ({})'.format(
- kernel_size)
- padding = (kernel_size - 1) // 2 * dilation
- self.linear = nn.Conv1d(
- in_channels,
- out_channels,
- kernel_size,
- stride=stride,
- padding=padding,
- dilation=dilation,
- bias=bias)
- self.nonlinear = get_nonlinear(config_str, out_channels)
-
- def forward(self, x):
- x = self.linear(x)
- x = self.nonlinear(x)
- return x
-
-
def extract_feature(audio):
features = []
for au in audio:
@@ -386,116 +155,6 @@
seg = seg[..., :x.shape[-1]]
return seg
-
-class CAMDenseTDNNLayer(nn.Module):
-
- def __init__(self,
- in_channels,
- out_channels,
- bn_channels,
- kernel_size,
- stride=1,
- dilation=1,
- bias=False,
- config_str='batchnorm-relu',
- memory_efficient=False):
- super(CAMDenseTDNNLayer, self).__init__()
- assert kernel_size % 2 == 1, 'Expect equal paddings, but got even kernel size ({})'.format(
- kernel_size)
- padding = (kernel_size - 1) // 2 * dilation
- self.memory_efficient = memory_efficient
- self.nonlinear1 = get_nonlinear(config_str, in_channels)
- self.linear1 = nn.Conv1d(in_channels, bn_channels, 1, bias=False)
- self.nonlinear2 = get_nonlinear(config_str, bn_channels)
- self.cam_layer = CAMLayer(
- bn_channels,
- out_channels,
- kernel_size,
- stride=stride,
- padding=padding,
- dilation=dilation,
- bias=bias)
-
- def bn_function(self, x):
- return self.linear1(self.nonlinear1(x))
-
- def forward(self, x):
- if self.training and self.memory_efficient:
- x = cp.checkpoint(self.bn_function, x)
- else:
- x = self.bn_function(x)
- x = self.cam_layer(self.nonlinear2(x))
- return x
-
-
-class CAMDenseTDNNBlock(nn.ModuleList):
-
- def __init__(self,
- num_layers,
- in_channels,
- out_channels,
- bn_channels,
- kernel_size,
- stride=1,
- dilation=1,
- bias=False,
- config_str='batchnorm-relu',
- memory_efficient=False):
- super(CAMDenseTDNNBlock, self).__init__()
- for i in range(num_layers):
- layer = CAMDenseTDNNLayer(
- in_channels=in_channels + i * out_channels,
- out_channels=out_channels,
- bn_channels=bn_channels,
- kernel_size=kernel_size,
- stride=stride,
- dilation=dilation,
- bias=bias,
- config_str=config_str,
- memory_efficient=memory_efficient)
- self.add_module('tdnnd%d' % (i + 1), layer)
-
- def forward(self, x):
- for layer in self:
- x = torch.cat([x, layer(x)], dim=1)
- return x
-
-
-class TransitLayer(nn.Module):
-
- def __init__(self,
- in_channels,
- out_channels,
- bias=True,
- config_str='batchnorm-relu'):
- super(TransitLayer, self).__init__()
- self.nonlinear = get_nonlinear(config_str, in_channels)
- self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias)
-
- def forward(self, x):
- x = self.nonlinear(x)
- x = self.linear(x)
- return x
-
-
-class DenseLayer(nn.Module):
-
- def __init__(self,
- in_channels,
- out_channels,
- bias=False,
- config_str='batchnorm-relu'):
- super(DenseLayer, self).__init__()
- self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias)
- self.nonlinear = get_nonlinear(config_str, out_channels)
-
- def forward(self, x):
- if len(x.shape) == 2:
- x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1)
- else:
- x = self.linear(x)
- x = self.nonlinear(x)
- return x
def postprocess(segments: list, vad_segments: list,
labels: np.ndarray, embeddings: np.ndarray) -> list:
@@ -592,300 +251,3 @@
d['spk'] = sentence_spk
sd_sentence_list.append(d)
return sd_sentence_list
-
-
-class AFF(nn.Module):
-
- def __init__(self, channels=64, r=4):
- super(AFF, self).__init__()
- inter_channels = int(channels // r)
-
- self.local_att = nn.Sequential(
- nn.Conv2d(channels * 2, inter_channels, kernel_size=1, stride=1, padding=0),
- nn.BatchNorm2d(inter_channels),
- nn.SiLU(inplace=True),
- nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
- nn.BatchNorm2d(channels),
- )
-
- def forward(self, x, ds_y):
- xa = torch.cat((x, ds_y), dim=1)
- x_att = self.local_att(xa)
- x_att = 1.0 + torch.tanh(x_att)
- xo = torch.mul(x, x_att) + torch.mul(ds_y, 2.0 - x_att)
-
- return xo
-
-
-class TSTP(nn.Module):
- """
- Temporal statistics pooling, concatenate mean and std, which is used in
- x-vector
- Comment: simple concatenation can not make full use of both statistics
- """
-
- def __init__(self, **kwargs):
- super(TSTP, self).__init__()
-
- def forward(self, x):
- # The last dimension is the temporal axis
- pooling_mean = x.mean(dim=-1)
- pooling_std = torch.sqrt(torch.var(x, dim=-1) + 1e-8)
- pooling_mean = pooling_mean.flatten(start_dim=1)
- pooling_std = pooling_std.flatten(start_dim=1)
-
- stats = torch.cat((pooling_mean, pooling_std), 1)
- return stats
-
-
-class ReLU(nn.Hardtanh):
-
- def __init__(self, inplace=False):
- super(ReLU, self).__init__(0, 20, inplace)
-
- def __repr__(self):
- inplace_str = 'inplace' if self.inplace else ''
- return self.__class__.__name__ + ' (' \
- + inplace_str + ')'
-
-
-def conv1x1(in_planes, out_planes, stride=1):
- "1x1 convolution without padding"
- return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride,
- padding=0, bias=False)
-
-
-def conv3x3(in_planes, out_planes, stride=1):
- "3x3 convolution with padding"
- return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
- padding=1, bias=False)
-
-
-class BasicBlockERes2Net(nn.Module):
- expansion = 4
-
- def __init__(self, in_planes, planes, stride=1, baseWidth=24, scale=3):
- super(BasicBlockERes2Net, self).__init__()
- width = int(math.floor(planes * (baseWidth / 64.0)))
- self.conv1 = conv1x1(in_planes, width * scale, stride)
- self.bn1 = nn.BatchNorm2d(width * scale)
- self.nums = scale
-
- convs = []
- bns = []
- for i in range(self.nums):
- convs.append(conv3x3(width, width))
- bns.append(nn.BatchNorm2d(width))
- self.convs = nn.ModuleList(convs)
- self.bns = nn.ModuleList(bns)
- self.relu = ReLU(inplace=True)
-
- self.conv3 = conv1x1(width * scale, planes * self.expansion)
- self.bn3 = nn.BatchNorm2d(planes * self.expansion)
- self.shortcut = nn.Sequential()
- if stride != 1 or in_planes != self.expansion * planes:
- self.shortcut = nn.Sequential(
- nn.Conv2d(in_planes,
- self.expansion * planes,
- kernel_size=1,
- stride=stride,
- bias=False),
- nn.BatchNorm2d(self.expansion * planes))
- self.stride = stride
- self.width = width
- self.scale = scale
-
- def forward(self, x):
- residual = x
-
- out = self.conv1(x)
- out = self.bn1(out)
- out = self.relu(out)
- spx = torch.split(out, self.width, 1)
- for i in range(self.nums):
- if i == 0:
- sp = spx[i]
- else:
- sp = sp + spx[i]
- sp = self.convs[i](sp)
- sp = self.relu(self.bns[i](sp))
- if i == 0:
- out = sp
- else:
- out = torch.cat((out, sp), 1)
-
- out = self.conv3(out)
- out = self.bn3(out)
-
- residual = self.shortcut(x)
- out += residual
- out = self.relu(out)
-
- return out
-
-
-class BasicBlockERes2Net_diff_AFF(nn.Module):
- expansion = 4
-
- def __init__(self, in_planes, planes, stride=1, baseWidth=24, scale=3):
- super(BasicBlockERes2Net_diff_AFF, self).__init__()
- width = int(math.floor(planes * (baseWidth / 64.0)))
- self.conv1 = conv1x1(in_planes, width * scale, stride)
- self.bn1 = nn.BatchNorm2d(width * scale)
-
- self.nums = scale
-
- convs = []
- fuse_models = []
- bns = []
- for i in range(self.nums):
- convs.append(conv3x3(width, width))
- bns.append(nn.BatchNorm2d(width))
- for j in range(self.nums - 1):
- fuse_models.append(AFF(channels=width))
-
- self.convs = nn.ModuleList(convs)
- self.bns = nn.ModuleList(bns)
- self.fuse_models = nn.ModuleList(fuse_models)
- self.relu = ReLU(inplace=True)
-
- self.conv3 = conv1x1(width * scale, planes * self.expansion)
- self.bn3 = nn.BatchNorm2d(planes * self.expansion)
- self.shortcut = nn.Sequential()
- if stride != 1 or in_planes != self.expansion * planes:
- self.shortcut = nn.Sequential(
- nn.Conv2d(in_planes,
- self.expansion * planes,
- kernel_size=1,
- stride=stride,
- bias=False),
- nn.BatchNorm2d(self.expansion * planes))
- self.stride = stride
- self.width = width
- self.scale = scale
-
- def forward(self, x):
- residual = x
-
- out = self.conv1(x)
- out = self.bn1(out)
- out = self.relu(out)
- spx = torch.split(out, self.width, 1)
- for i in range(self.nums):
- if i == 0:
- sp = spx[i]
- else:
- sp = self.fuse_models[i - 1](sp, spx[i])
-
- sp = self.convs[i](sp)
- sp = self.relu(self.bns[i](sp))
- if i == 0:
- out = sp
- else:
- out = torch.cat((out, sp), 1)
-
- out = self.conv3(out)
- out = self.bn3(out)
-
- residual = self.shortcut(x)
- out += residual
- out = self.relu(out)
-
- return out
-
-
-class ERes2Net(nn.Module):
- def __init__(self,
- block=BasicBlockERes2Net,
- block_fuse=BasicBlockERes2Net_diff_AFF,
- num_blocks=[3, 4, 6, 3],
- m_channels=64,
- feat_dim=80,
- embedding_size=192,
- pooling_func='TSTP',
- two_emb_layer=False):
- super(ERes2Net, self).__init__()
- self.in_planes = m_channels
- self.feat_dim = feat_dim
- self.embedding_size = embedding_size
- self.stats_dim = int(feat_dim / 8) * m_channels * 8
- self.two_emb_layer = two_emb_layer
-
- self.conv1 = nn.Conv2d(1,
- m_channels,
- kernel_size=3,
- stride=1,
- padding=1,
- bias=False)
- self.bn1 = nn.BatchNorm2d(m_channels)
- self.layer1 = self._make_layer(block,
- m_channels,
- num_blocks[0],
- stride=1)
- self.layer2 = self._make_layer(block,
- m_channels * 2,
- num_blocks[1],
- stride=2)
- self.layer3 = self._make_layer(block_fuse,
- m_channels * 4,
- num_blocks[2],
- stride=2)
- self.layer4 = self._make_layer(block_fuse,
- m_channels * 8,
- num_blocks[3],
- stride=2)
-
- self.layer1_downsample = nn.Conv2d(m_channels * 4, m_channels * 8, kernel_size=3, padding=1, stride=2,
- bias=False)
- self.layer2_downsample = nn.Conv2d(m_channels * 8, m_channels * 16, kernel_size=3, padding=1, stride=2,
- bias=False)
- self.layer3_downsample = nn.Conv2d(m_channels * 16, m_channels * 32, kernel_size=3, padding=1, stride=2,
- bias=False)
- self.fuse_mode12 = AFF(channels=m_channels * 8)
- self.fuse_mode123 = AFF(channels=m_channels * 16)
- self.fuse_mode1234 = AFF(channels=m_channels * 32)
-
- self.n_stats = 1 if pooling_func == 'TAP' or pooling_func == "TSDP" else 2
- self.pool = TSTP(in_dim=self.stats_dim * block.expansion)
- self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats,
- embedding_size)
- if self.two_emb_layer:
- self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False)
- self.seg_2 = nn.Linear(embedding_size, embedding_size)
- else:
- self.seg_bn_1 = nn.Identity()
- self.seg_2 = nn.Identity()
-
- def _make_layer(self, block, planes, num_blocks, stride):
- strides = [stride] + [1] * (num_blocks - 1)
- layers = []
- for stride in strides:
- layers.append(block(self.in_planes, planes, stride))
- self.in_planes = planes * block.expansion
- return nn.Sequential(*layers)
-
- def forward(self, x):
- x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
-
- x = x.unsqueeze_(1)
- out = F.relu(self.bn1(self.conv1(x)))
- out1 = self.layer1(out)
- out2 = self.layer2(out1)
- out1_downsample = self.layer1_downsample(out1)
- fuse_out12 = self.fuse_mode12(out2, out1_downsample)
- out3 = self.layer3(out2)
- fuse_out12_downsample = self.layer2_downsample(fuse_out12)
- fuse_out123 = self.fuse_mode123(out3, fuse_out12_downsample)
- out4 = self.layer4(out3)
- fuse_out123_downsample = self.layer3_downsample(fuse_out123)
- fuse_out1234 = self.fuse_mode1234(out4, fuse_out123_downsample)
- stats = self.pool(fuse_out1234)
-
- embed_a = self.seg_1(stats)
- if self.two_emb_layer:
- out = F.relu(embed_a)
- out = self.seg_bn_1(out)
- embed_b = self.seg_2(out)
- return embed_b
- else:
- return embed_a
--
Gitblit v1.9.1