Initial commit: FunASR Speech Recognition Toolkit
Update API Documentation / build-api-docs (push) Has been cancelled
Update API Documentation / build-api-docs (push) Has been cancelled
Add complete FunASR codebase including models, runtime, and documentation.
This commit is contained in:
freedakgmail
@@ -0,0 +1,2 @@
|
||||
from .eres2net import ERes2Net
|
||||
from .eres2net_aug import ERes2NetAug
|
||||
@@ -0,0 +1,481 @@
|
||||
# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
""" Res2Net implementation is adapted from https://github.com/wenet-e2e/wespeaker.
|
||||
ERes2Net incorporates both local and global feature fusion techniques to improve the performance.
|
||||
The local feature fusion (LFF) fuses the features within one single residual block to extract the local signal.
|
||||
The global feature fusion (GFF) takes acoustic features of different scales as input to aggregate global signal.
|
||||
ERes2Net-Large is an upgraded version of ERes2Net that uses a larger number of parameters to achieve better
|
||||
recognition performance. Parameters expansion, baseWidth, and scale can be modified to obtain optimal performance.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
import funasr.models.sond.pooling.pooling_layers as pooling_layers
|
||||
|
||||
from funasr.models.eres2net.fusion import AFF
|
||||
|
||||
|
||||
class ReLU(nn.Hardtanh):
|
||||
|
||||
def __init__(self, inplace=False):
|
||||
"""Initialize ReLU.
|
||||
|
||||
Args:
|
||||
inplace: TODO.
|
||||
"""
|
||||
super(ReLU, self).__init__(0, 20, inplace)
|
||||
|
||||
def __repr__(self):
|
||||
"""Internal: repr ."""
|
||||
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 = 2
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1, baseWidth=32, scale=2):
|
||||
"""Initialize BasicBlockERes2Net.
|
||||
|
||||
Args:
|
||||
in_planes: TODO.
|
||||
planes: TODO.
|
||||
stride: TODO.
|
||||
baseWidth: TODO.
|
||||
scale: TODO.
|
||||
"""
|
||||
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):
|
||||
"""Forward pass for training.
|
||||
|
||||
Args:
|
||||
x: TODO.
|
||||
"""
|
||||
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 = 2
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1, baseWidth=32, scale=2):
|
||||
"""Initialize BasicBlockERes2Net_diff_AFF.
|
||||
|
||||
Args:
|
||||
in_planes: TODO.
|
||||
planes: TODO.
|
||||
stride: TODO.
|
||||
baseWidth: TODO.
|
||||
scale: TODO.
|
||||
"""
|
||||
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):
|
||||
"""Forward pass for training.
|
||||
|
||||
Args:
|
||||
x: TODO.
|
||||
"""
|
||||
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=32,
|
||||
feat_dim=80,
|
||||
embedding_size=192,
|
||||
pooling_func="TSTP",
|
||||
two_emb_layer=False,
|
||||
):
|
||||
"""Initialize ERes2Net.
|
||||
|
||||
Args:
|
||||
block: TODO.
|
||||
block_fuse: TODO.
|
||||
num_blocks: TODO.
|
||||
m_channels: TODO.
|
||||
feat_dim: Size/dimension parameter.
|
||||
embedding_size: Size/dimension parameter.
|
||||
pooling_func: TODO.
|
||||
two_emb_layer: TODO.
|
||||
"""
|
||||
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)
|
||||
|
||||
# Downsampling module for each layer
|
||||
self.layer1_downsample = nn.Conv2d(
|
||||
m_channels * 2, m_channels * 4, kernel_size=3, stride=2, padding=1, bias=False
|
||||
)
|
||||
self.layer2_downsample = nn.Conv2d(
|
||||
m_channels * 4, m_channels * 8, kernel_size=3, padding=1, stride=2, bias=False
|
||||
)
|
||||
self.layer3_downsample = nn.Conv2d(
|
||||
m_channels * 8, m_channels * 16, kernel_size=3, padding=1, stride=2, bias=False
|
||||
)
|
||||
|
||||
# Bottom-up fusion module
|
||||
self.fuse_mode12 = AFF(channels=m_channels * 4)
|
||||
self.fuse_mode123 = AFF(channels=m_channels * 8)
|
||||
self.fuse_mode1234 = AFF(channels=m_channels * 16)
|
||||
|
||||
self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2
|
||||
self.pool = getattr(pooling_layers, pooling_func)(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):
|
||||
"""Internal: make layer.
|
||||
|
||||
Args:
|
||||
block: TODO.
|
||||
planes: TODO.
|
||||
num_blocks: TODO.
|
||||
stride: TODO.
|
||||
"""
|
||||
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):
|
||||
"""Forward pass for training.
|
||||
|
||||
Args:
|
||||
x: TODO.
|
||||
"""
|
||||
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
|
||||
|
||||
|
||||
class BasicBlockRes2Net(nn.Module):
|
||||
expansion = 2
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1, baseWidth=32, scale=2):
|
||||
"""Initialize BasicBlockRes2Net.
|
||||
|
||||
Args:
|
||||
in_planes: TODO.
|
||||
planes: TODO.
|
||||
stride: TODO.
|
||||
baseWidth: TODO.
|
||||
scale: TODO.
|
||||
"""
|
||||
super(BasicBlockRes2Net, 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 - 1
|
||||
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):
|
||||
"""Forward pass for training.
|
||||
|
||||
Args:
|
||||
x: TODO.
|
||||
"""
|
||||
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 = torch.cat((out, spx[self.nums]), 1)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
residual = self.shortcut(x)
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class Res2Net(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
block=BasicBlockRes2Net,
|
||||
num_blocks=[3, 4, 6, 3],
|
||||
m_channels=32,
|
||||
feat_dim=80,
|
||||
embedding_size=192,
|
||||
pooling_func="TSTP",
|
||||
two_emb_layer=False,
|
||||
):
|
||||
"""Initialize Res2Net.
|
||||
|
||||
Args:
|
||||
block: TODO.
|
||||
num_blocks: TODO.
|
||||
m_channels: TODO.
|
||||
feat_dim: Size/dimension parameter.
|
||||
embedding_size: Size/dimension parameter.
|
||||
pooling_func: TODO.
|
||||
two_emb_layer: TODO.
|
||||
"""
|
||||
super(Res2Net, 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, m_channels * 4, num_blocks[2], stride=2)
|
||||
self.layer4 = self._make_layer(block, m_channels * 8, num_blocks[3], stride=2)
|
||||
|
||||
self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2
|
||||
self.pool = getattr(pooling_layers, pooling_func)(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):
|
||||
"""Internal: make layer.
|
||||
|
||||
Args:
|
||||
block: TODO.
|
||||
planes: TODO.
|
||||
num_blocks: TODO.
|
||||
stride: TODO.
|
||||
"""
|
||||
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):
|
||||
"""Forward pass for training.
|
||||
|
||||
Args:
|
||||
x: TODO.
|
||||
"""
|
||||
x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
|
||||
|
||||
x = x.unsqueeze_(1)
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out = self.layer1(out)
|
||||
out = self.layer2(out)
|
||||
out = self.layer3(out)
|
||||
out = self.layer4(out)
|
||||
|
||||
stats = self.pool(out)
|
||||
|
||||
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
|
||||
@@ -0,0 +1,314 @@
|
||||
# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
""" Res2Net implementation is adapted from https://github.com/wenet-e2e/wespeaker.
|
||||
ERes2Net incorporates both local and global feature fusion techniques to improve the performance.
|
||||
The local feature fusion (LFF) fuses the features within one single residual block to extract the local signal.
|
||||
The global feature fusion (GFF) takes acoustic features of different scales as input to aggregate global signal.
|
||||
ERes2Net-Large is an upgraded version of ERes2Net that uses a larger number of parameters to achieve better
|
||||
recognition performance. Parameters expansion, baseWidth, and scale can be modified to obtain optimal performance.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
import funasr.models.sond.pooling.pooling_layers as pooling_layers
|
||||
|
||||
from funasr.models.eres2net.fusion import AFF
|
||||
|
||||
|
||||
class ReLU(nn.Hardtanh):
|
||||
|
||||
def __init__(self, inplace=False):
|
||||
"""Initialize ReLU.
|
||||
|
||||
Args:
|
||||
inplace: TODO.
|
||||
"""
|
||||
super(ReLU, self).__init__(0, 20, inplace)
|
||||
|
||||
def __repr__(self):
|
||||
"""Internal: repr ."""
|
||||
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):
|
||||
"""Initialize BasicBlockERes2Net.
|
||||
|
||||
Args:
|
||||
in_planes: TODO.
|
||||
planes: TODO.
|
||||
stride: TODO.
|
||||
baseWidth: TODO.
|
||||
scale: TODO.
|
||||
"""
|
||||
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):
|
||||
"""Forward pass for training.
|
||||
|
||||
Args:
|
||||
x: TODO.
|
||||
"""
|
||||
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):
|
||||
"""Initialize BasicBlockERes2Net_diff_AFF.
|
||||
|
||||
Args:
|
||||
in_planes: TODO.
|
||||
planes: TODO.
|
||||
stride: TODO.
|
||||
baseWidth: TODO.
|
||||
scale: TODO.
|
||||
"""
|
||||
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):
|
||||
"""Forward pass for training.
|
||||
|
||||
Args:
|
||||
x: TODO.
|
||||
"""
|
||||
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 ERes2NetAug(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,
|
||||
):
|
||||
"""Initialize ERes2NetAug.
|
||||
|
||||
Args:
|
||||
block: TODO.
|
||||
block_fuse: TODO.
|
||||
num_blocks: TODO.
|
||||
m_channels: TODO.
|
||||
feat_dim: Size/dimension parameter.
|
||||
embedding_size: Size/dimension parameter.
|
||||
pooling_func: TODO.
|
||||
two_emb_layer: TODO.
|
||||
"""
|
||||
super(ERes2NetAug, 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 = getattr(pooling_layers, pooling_func)(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):
|
||||
"""Internal: make layer.
|
||||
|
||||
Args:
|
||||
block: TODO.
|
||||
planes: TODO.
|
||||
num_blocks: TODO.
|
||||
stride: TODO.
|
||||
"""
|
||||
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):
|
||||
"""Forward pass for training.
|
||||
|
||||
Args:
|
||||
x: TODO.
|
||||
"""
|
||||
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
|
||||
@@ -0,0 +1,290 @@
|
||||
# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
import funasr.models.sond.pooling.pooling_layers as pooling_layers
|
||||
from funasr.models.eres2net.fusion import AFF
|
||||
|
||||
|
||||
class ReLU(nn.Hardtanh):
|
||||
|
||||
def __init__(self, inplace=False):
|
||||
"""Initialize ReLU.
|
||||
|
||||
Args:
|
||||
inplace: TODO.
|
||||
"""
|
||||
super(ReLU, self).__init__(0, 20, inplace)
|
||||
|
||||
def __repr__(self):
|
||||
"""Internal: repr ."""
|
||||
inplace_str = "inplace" if self.inplace else ""
|
||||
return self.__class__.__name__ + " (" + inplace_str + ")"
|
||||
|
||||
|
||||
class BasicBlockERes2NetV2(nn.Module):
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1, baseWidth=26, scale=2, expansion=2):
|
||||
"""Initialize BasicBlockERes2NetV2.
|
||||
|
||||
Args:
|
||||
in_planes: TODO.
|
||||
planes: TODO.
|
||||
stride: TODO.
|
||||
baseWidth: TODO.
|
||||
scale: TODO.
|
||||
expansion: TODO.
|
||||
"""
|
||||
super(BasicBlockERes2NetV2, self).__init__()
|
||||
width = int(math.floor(planes * (baseWidth / 64.0)))
|
||||
self.conv1 = nn.Conv2d(in_planes, width * scale, kernel_size=1, stride=stride, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(width * scale)
|
||||
self.nums = scale
|
||||
self.expansion = expansion
|
||||
|
||||
convs = []
|
||||
bns = []
|
||||
for i in range(self.nums):
|
||||
convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False))
|
||||
bns.append(nn.BatchNorm2d(width))
|
||||
self.convs = nn.ModuleList(convs)
|
||||
self.bns = nn.ModuleList(bns)
|
||||
self.relu = ReLU(inplace=True)
|
||||
|
||||
self.conv3 = nn.Conv2d(width * scale, planes * self.expansion, kernel_size=1, bias=False)
|
||||
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):
|
||||
"""Forward pass for training.
|
||||
|
||||
Args:
|
||||
x: TODO.
|
||||
"""
|
||||
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 BasicBlockERes2NetV2AFF(nn.Module):
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1, baseWidth=26, scale=2, expansion=2):
|
||||
"""Initialize BasicBlockERes2NetV2AFF.
|
||||
|
||||
Args:
|
||||
in_planes: TODO.
|
||||
planes: TODO.
|
||||
stride: TODO.
|
||||
baseWidth: TODO.
|
||||
scale: TODO.
|
||||
expansion: TODO.
|
||||
"""
|
||||
super(BasicBlockERes2NetV2AFF, self).__init__()
|
||||
width = int(math.floor(planes * (baseWidth / 64.0)))
|
||||
self.conv1 = nn.Conv2d(in_planes, width * scale, kernel_size=1, stride=stride, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(width * scale)
|
||||
self.nums = scale
|
||||
self.expansion = expansion
|
||||
|
||||
convs = []
|
||||
fuse_models = []
|
||||
bns = []
|
||||
for i in range(self.nums):
|
||||
convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False))
|
||||
bns.append(nn.BatchNorm2d(width))
|
||||
for j in range(self.nums - 1):
|
||||
fuse_models.append(AFF(channels=width, r=4))
|
||||
|
||||
self.convs = nn.ModuleList(convs)
|
||||
self.bns = nn.ModuleList(bns)
|
||||
self.fuse_models = nn.ModuleList(fuse_models)
|
||||
self.relu = ReLU(inplace=True)
|
||||
|
||||
self.conv3 = nn.Conv2d(width * scale, planes * self.expansion, kernel_size=1, bias=False)
|
||||
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):
|
||||
"""Forward pass for training.
|
||||
|
||||
Args:
|
||||
x: TODO.
|
||||
"""
|
||||
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 ERes2NetV2(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
block=BasicBlockERes2NetV2,
|
||||
block_fuse=BasicBlockERes2NetV2AFF,
|
||||
num_blocks=[3, 4, 6, 3],
|
||||
m_channels=64,
|
||||
feat_dim=80,
|
||||
embedding_size=192,
|
||||
baseWidth=26,
|
||||
scale=2,
|
||||
expansion=2,
|
||||
pooling_func="TSTP",
|
||||
two_emb_layer=False,
|
||||
):
|
||||
"""Initialize ERes2NetV2.
|
||||
|
||||
Args:
|
||||
block: TODO.
|
||||
block_fuse: TODO.
|
||||
num_blocks: TODO.
|
||||
m_channels: TODO.
|
||||
feat_dim: Size/dimension parameter.
|
||||
embedding_size: Size/dimension parameter.
|
||||
baseWidth: TODO.
|
||||
scale: TODO.
|
||||
expansion: TODO.
|
||||
pooling_func: TODO.
|
||||
two_emb_layer: TODO.
|
||||
"""
|
||||
super(ERes2NetV2, 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.baseWidth = baseWidth
|
||||
self.scale = scale
|
||||
self.expansion = expansion
|
||||
|
||||
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.layer3_ds = nn.Conv2d(
|
||||
m_channels * 4 * self.expansion, m_channels * 8 * self.expansion,
|
||||
kernel_size=3, padding=1, stride=2, bias=False,
|
||||
)
|
||||
self.fuse34 = AFF(channels=m_channels * 8 * self.expansion, r=4)
|
||||
|
||||
self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2
|
||||
self.pool = getattr(pooling_layers, pooling_func)(in_dim=self.stats_dim * self.expansion)
|
||||
self.seg_1 = nn.Linear(self.stats_dim * self.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):
|
||||
"""Internal: make layer.
|
||||
|
||||
Args:
|
||||
block: TODO.
|
||||
planes: TODO.
|
||||
num_blocks: TODO.
|
||||
stride: TODO.
|
||||
"""
|
||||
strides = [stride] + [1] * (num_blocks - 1)
|
||||
layers = []
|
||||
for stride in strides:
|
||||
layers.append(
|
||||
block(self.in_planes, planes, stride, baseWidth=self.baseWidth, scale=self.scale, expansion=self.expansion)
|
||||
)
|
||||
self.in_planes = planes * self.expansion
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward pass for training.
|
||||
|
||||
Args:
|
||||
x: TODO.
|
||||
"""
|
||||
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)
|
||||
out3 = self.layer3(out2)
|
||||
out4 = self.layer4(out3)
|
||||
out3_ds = self.layer3_ds(out3)
|
||||
fuse_out34 = self.fuse34(out4, out3_ds)
|
||||
stats = self.pool(fuse_out34)
|
||||
|
||||
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
|
||||
@@ -0,0 +1,40 @@
|
||||
# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class AFF(nn.Module):
|
||||
|
||||
def __init__(self, channels=64, r=4):
|
||||
"""Initialize AFF.
|
||||
|
||||
Args:
|
||||
channels: TODO.
|
||||
r: TODO.
|
||||
"""
|
||||
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):
|
||||
"""Forward pass for training.
|
||||
|
||||
Args:
|
||||
x: TODO.
|
||||
ds_y: TODO.
|
||||
"""
|
||||
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
|
||||
@@ -0,0 +1,137 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- encoding: utf-8 -*-
|
||||
# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
|
||||
# MIT License (https://opensource.org/licenses/MIT)
|
||||
# Modified from 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker)
|
||||
|
||||
import os
|
||||
import time
|
||||
import logging
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from funasr.register import tables
|
||||
from funasr.models.campplus.utils import extract_feature
|
||||
from funasr.utils.load_utils import load_audio_text_image_video
|
||||
from funasr.models.eres2net.eres2netv2 import ERes2NetV2
|
||||
|
||||
|
||||
@tables.register("model_classes", "ERes2NetV2")
|
||||
@tables.register("model_classes", "iic/speech_eres2netv2_sv_zh-cn_16k-common")
|
||||
class ERes2NetV2SV(torch.nn.Module):
|
||||
"""ERes2NetV2: Enhanced Res2Net v2 for Speaker Verification.
|
||||
|
||||
Improved speaker embedding model based on Res2Net architecture with
|
||||
multi-scale feature aggregation. Provides 192-dim speaker embeddings
|
||||
for speaker verification and diarization.
|
||||
|
||||
Better than CAM++ for short-duration audio (< 3s) speaker feature extraction.
|
||||
|
||||
Output: {"spk_embedding": Tensor of shape (1, 192)}
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
feat_dim=80,
|
||||
embedding_size=192,
|
||||
m_channels=64,
|
||||
baseWidth=26,
|
||||
scale=2,
|
||||
expansion=2,
|
||||
num_blocks=[3, 4, 6, 3],
|
||||
pooling_func="TSTP",
|
||||
two_emb_layer=False,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize ERes2NetV2SV.
|
||||
|
||||
Args:
|
||||
feat_dim: Size/dimension parameter.
|
||||
embedding_size: Size/dimension parameter.
|
||||
m_channels: TODO.
|
||||
baseWidth: TODO.
|
||||
scale: TODO.
|
||||
expansion: TODO.
|
||||
num_blocks: TODO.
|
||||
pooling_func: TODO.
|
||||
two_emb_layer: TODO.
|
||||
**kwargs: Additional keyword arguments.
|
||||
"""
|
||||
super().__init__()
|
||||
self.model = ERes2NetV2(
|
||||
feat_dim=feat_dim,
|
||||
embedding_size=embedding_size,
|
||||
m_channels=m_channels,
|
||||
baseWidth=baseWidth,
|
||||
scale=scale,
|
||||
expansion=expansion,
|
||||
num_blocks=num_blocks,
|
||||
pooling_func=pooling_func,
|
||||
two_emb_layer=two_emb_layer,
|
||||
)
|
||||
self.embedding_size = embedding_size
|
||||
|
||||
model_path = kwargs.get("model_path", None)
|
||||
init_param = kwargs.get("init_param", None)
|
||||
if init_param is None and model_path is not None:
|
||||
ckpt = os.path.join(model_path, "pretrained_eres2netv2.ckpt")
|
||||
if os.path.exists(ckpt):
|
||||
init_param = ckpt
|
||||
if init_param is not None and os.path.exists(init_param):
|
||||
self._load_pretrained(init_param)
|
||||
|
||||
def _load_pretrained(self, path):
|
||||
"""Internal: load pretrained.
|
||||
|
||||
Args:
|
||||
path: TODO.
|
||||
"""
|
||||
state_dict = torch.load(path, map_location="cpu")
|
||||
if "state_dict" in state_dict:
|
||||
state_dict = state_dict["state_dict"]
|
||||
missing, unexpected = self.model.load_state_dict(state_dict, strict=False)
|
||||
if missing:
|
||||
logging.warning(f"ERes2NetV2 missing keys: {missing[:5]}...")
|
||||
logging.info(f"ERes2NetV2 loaded pretrained weights from {path}")
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward pass for training.
|
||||
|
||||
Args:
|
||||
x: TODO.
|
||||
"""
|
||||
return self.model(x)
|
||||
|
||||
def inference(
|
||||
self,
|
||||
data_in,
|
||||
data_lengths=None,
|
||||
key: list = None,
|
||||
tokenizer=None,
|
||||
frontend=None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Run inference on input data.
|
||||
|
||||
Args:
|
||||
data_in: Input data (audio samples, file paths, or text).
|
||||
data_lengths: Lengths of each input sample in the batch.
|
||||
key: Sample identifiers.
|
||||
tokenizer: Tokenizer instance for text encoding/decoding.
|
||||
frontend: Audio frontend for feature extraction.
|
||||
**kwargs: Additional keyword arguments.
|
||||
"""
|
||||
meta_data = {}
|
||||
time1 = time.perf_counter()
|
||||
audio_sample_list = load_audio_text_image_video(
|
||||
data_in, fs=16000, audio_fs=kwargs.get("fs", 16000), data_type="sound"
|
||||
)
|
||||
time2 = time.perf_counter()
|
||||
meta_data["load_data"] = f"{time2 - time1:0.3f}"
|
||||
speech, speech_lengths, speech_times = extract_feature(audio_sample_list)
|
||||
speech = speech.to(device=kwargs["device"])
|
||||
time3 = time.perf_counter()
|
||||
meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
|
||||
meta_data["batch_data_time"] = np.array(speech_times).sum().item() / 16000.0
|
||||
results = [{"spk_embedding": self.forward(speech.to(torch.float32))}]
|
||||
return results, meta_data
|
||||
Reference in New Issue
Block a user