1/mvector/models/ecapa_tdnn.py

import torch
import torch.nn as nn
import torch.nn.functional as F

from mvector.models.pooling import AttentiveStatsPool, TemporalAveragePooling
from mvector.models.pooling import SelfAttentivePooling, TemporalStatisticsPooling


class Res2Conv1dReluBn(nn.Module):
    def __init__(self, channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False, scale=4):
        super().__init__()
        assert channels % scale == 0, "{} % {} != 0".format(channels, scale)
        self.scale = scale
        self.width = channels // scale
        self.nums = scale if scale == 1 else scale - 1

        self.convs = []
        self.bns = []
        for i in range(self.nums):
            self.convs.append(nn.Conv1d(self.width, self.width, kernel_size, stride, padding, dilation, bias=bias))
            self.bns.append(nn.BatchNorm1d(self.width))
        self.convs = nn.ModuleList(self.convs)
        self.bns = nn.ModuleList(self.bns)

    def forward(self, x):
        out = []
        spx = torch.split(x, self.width, 1)
        # 遍历每个分支
        for i in range(self.nums):
            if i == 0:
                sp = spx[i]
            else:
                 # 其他分支则将当前子特征与前面所有子特征相加，形成残差连接
                sp = sp + spx[i]
            # Order: conv -> relu -> bn
            sp = self.convs[i](sp)
            sp = self.bns[i](F.relu(sp))
            out.append(sp)
        if self.scale != 1:
            out.append(spx[self.nums])
            
         # 将所有子分支的结果在通道维度上合并
        out = torch.cat(out, dim=1)
        return out


class Conv1dReluBn(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False):
        super().__init__()
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)
        self.bn = nn.BatchNorm1d(out_channels)

    def forward(self, x):
        return self.bn(F.relu(self.conv(x)))


class SE_Connect(nn.Module):
    def __init__(self, channels, s=2):
        super().__init__()
        assert channels % s == 0, "{} % {} != 0".format(channels, s)
        self.linear1 = nn.Linear(channels, channels // s)
        self.linear2 = nn.Linear(channels // s, channels)

    def forward(self, x):
        out = x.mean(dim=2)
        out = F.relu(self.linear1(out))
        out = torch.sigmoid(self.linear2(out))
        out = x * out.unsqueeze(2)
        return out

def SE_Res2Block(channels, kernel_size, stride, padding, dilation, scale):
    """
    初始化函数。
        
    参数:
    - input_size: 输入尺寸，默认为80。
    - channels: 通道数，默认为512。
    - kernel_size: 卷积核大小, 默认为3。
    - embd_dim: 嵌入维度，默认为192。
    - pooling_type: 池化类型，默认为"ASP"，可选值包括"ASP"、"SAP"、"TAP"、"TSP"。
    - dilation : 空洞卷积的空洞率，默认为1。
    - scale: SE模块的缩放比例，默认为8。

    返回值:
    - 无。
    """
    return nn.Sequential(
        Conv1dReluBn(channels, channels, kernel_size=1, stride=1, padding=0),
        Res2Conv1dReluBn(channels, kernel_size, stride, padding, dilation, scale=scale),
        Conv1dReluBn(channels, channels, kernel_size=1, stride=1, padding=0),
        SE_Connect(channels)
    )


class EcapaTdnn(nn.Module):
    """
    初始化函数。
        
    参数:
    - input_size: 输入尺寸，默认为80。
    - channels: 通道数，默认为512。
    - embd_dim: 嵌入维度，默认为192。
    - pooling_type: 池化类型，默认为"ASP"，可选值包括"ASP"、"SAP"、"TAP"、"TSP"。
     """
    def __init__(self, input_size=80, channels=512, embd_dim=192, pooling_type="ASP"):
        super().__init__()
        self.layer1 = Conv1dReluBn(input_size, channels, kernel_size=5, padding=2, dilation=1)
        self.layer2 = SE_Res2Block(channels, kernel_size=3, stride=1, padding=2, dilation=2, scale=8)
        self.layer3 = SE_Res2Block(channels, kernel_size=3, stride=1, padding=3, dilation=3, scale=8)
        self.layer4 = SE_Res2Block(channels, kernel_size=3, stride=1, padding=4, dilation=4, scale=8)

        cat_channels = channels * 3
        self.emb_size = embd_dim
        self.conv = nn.Conv1d(cat_channels, cat_channels, kernel_size=1)
        if pooling_type == "ASP":
            self.pooling = AttentiveStatsPool(cat_channels, 128)
            self.bn1 = nn.BatchNorm1d(cat_channels * 2)
            self.linear = nn.Linear(cat_channels * 2, embd_dim)
            self.bn2 = nn.BatchNorm1d(embd_dim)
        elif pooling_type == "SAP":
            self.pooling = SelfAttentivePooling(cat_channels, 128)
            self.bn1 = nn.BatchNorm1d(cat_channels)
            self.linear = nn.Linear(cat_channels, embd_dim)
            self.bn2 = nn.BatchNorm1d(embd_dim)
        elif pooling_type == "TAP":
            self.pooling = TemporalAveragePooling()
            self.bn1 = nn.BatchNorm1d(cat_channels)
            self.linear = nn.Linear(cat_channels, embd_dim)
            self.bn2 = nn.BatchNorm1d(embd_dim)
        elif pooling_type == "TSP":
            self.pooling = TemporalStatisticsPooling()
            self.bn1 = nn.BatchNorm1d(cat_channels * 2)
            self.linear = nn.Linear(cat_channels * 2, embd_dim)
            self.bn2 = nn.BatchNorm1d(embd_dim)
        else:
            raise Exception(f'没有{pooling_type}池化层！')

    # def forward(self, x):
    #     """
    #     Compute embeddings.

    #     Args:
    #         x (torch.Tensor): Input data with shape (N, time, freq).

    #     Returns:
    #         torch.Tensor: Output embeddings with shape (N, self.emb_size, 1)
    #     """
    #     x = x.transpose(2, 1)
    #     out1 = self.layer1(x)
    #     out2 = self.layer2(out1) + out1
    #     out3 = self.layer3(out1 + out2) + out1 + out2
    #     out4 = self.layer4(out1 + out2 + out3) + out1 + out2 + out3

    #     out = torch.cat([out2, out3, out4], dim=1)
    #     out = F.relu(self.conv(out))
    #     out = self.bn1(self.pooling(out))
    #     out = self.bn2(self.linear(out))
    #     return out
        
    def forward(self, x):
        """
        计算嵌入向量。

        参数:
            x (torch.Tensor): 输入数据，形状为 (N, time, freq)，其中N为样本数量，time为时间维度，freq为频率维度。

        返回值:
            torch.Tensor: 输出嵌入向量，形状为 (N, self.emb_size, 1)
        """
        # 将输入数据的频率和时间维度交换
        x = x.transpose(2, 1)
        # 通过第一层卷积层
        out1 = self.layer1(x)
        # 通过第二层卷积层，并与第一层输出相加
        out2 = self.layer2(out1) + out1
        # 通过第三层卷积层，并依次与前两层输出相加
        out3 = self.layer3(out1 + out2) + out1 + out2
        # 通过第四层卷积层，并依次与前三层输出相加
        out4 = self.layer4(out1 + out2 + out3) + out1 + out2 + out3

        # 将第二、三、四层的输出在特征维度上连接
        out = torch.cat([out2, out3, out4], dim=1)
        # 应用ReLU激活函数，并通过卷积层处理
        out = F.relu(self.conv(out))
        # 经过批归一化和池化操作
        out = self.bn1(self.pooling(out))
        # 经过线性变换和批归一化
        out = self.bn2(self.linear(out))
        return out