import torch import torch.nn as nn import torch.nn.functional as F import torchaudio from torchinfo import summary ''' Res2Conv1d + BatchNorm1d + ReLU ''' class Res2Conv1dReluBn(nn.Module): ''' in_channels == out_channels == channels ''' 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 ''' Conv1d + BatchNorm1d + ReLU ''' 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))) ''' The SE connection of 1D case. ''' 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 ''' SE-Res2Block. Note: residual connection is implemented in the ECAPA_TDNN.yaml model, not here. ''' class SE_Res2Block(nn.Module): def __init__(self, channels, kernel_size, stride, padding, dilation, scale): super().__init__() self.block = 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) ) def forward(self, x): out = self.block(x) return out + x ''' Attentive weighted mean and standard deviation pooling. ''' class AttentiveStatsPool(nn.Module): def __init__(self, in_dim, bottleneck_dim): super().__init__() # Use Conv1d with stride == 1 rather than Linear, then we don't need to transpose inputs. self.linear1 = nn.Conv1d(in_dim, bottleneck_dim, kernel_size=1) # equals W and b in the paper self.linear2 = nn.Conv1d(bottleneck_dim, in_dim, kernel_size=1) # equals V and k in the paper def forward(self, x): # DON'T use ReLU here! In experiments, I find ReLU hard to converge. alpha = torch.tanh(self.linear1(x)) alpha = torch.softmax(self.linear2(alpha), dim=2) mean = torch.sum(alpha * x, dim=2) residuals = torch.sum(alpha * x ** 2, dim=2) - mean ** 2 std = torch.sqrt(residuals.clamp(min=1e-9)) return torch.cat([mean, std], dim=1) ''' Implementation of "ECAPA-TDNN: Emphasized Channel Attention, Propagation and Aggregation in TDNN Based Speaker Verification". Note that we DON'T concatenate the last frame-wise layer with non-weighted mean and standard deviation, because it brings little improvment but significantly increases model parameters. As a result, this implementation basically equals the A.2 of Table 2 in the paper. ''' class ECAPA_TDNN(nn.Module): def __init__(self, in_channels=80, channels=512, embd_dim=192): super().__init__() self.torchfb = torchaudio.transforms.MelSpectrogram(sample_rate=16000, n_fft=512, win_length=400, hop_length=160, f_min=0.0, f_max=8000, pad=0, n_mels=80) self.instancenorm = nn.InstanceNorm1d(80) self.layer1 = Conv1dReluBn(in_channels, channels, kernel_size=5, padding=2) 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.conv = nn.Conv1d(cat_channels, cat_channels, kernel_size=1) 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) def forward(self, x): x = self.torchfb(x) + 1e-6 x = x.log() x = self.instancenorm(x) # print(x.shape) # x = x.transpose(1, 2) 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)) # print(out.shape) out = self.bn1(self.pooling(out)) # print(out.shape) out = self.bn2(self.linear(out)) return out class ECAPA_TDNN_ks5(nn.Module): def __init__(self, in_channels=80, channels=512, embd_dim=192): super().__init__() self.torchfb = torchaudio.transforms.MelSpectrogram(sample_rate=16000, n_fft=512, win_length=400, hop_length=160, f_min=0.0, f_max=8000, pad=0, n_mels=80) self.instancenorm = nn.InstanceNorm1d(40) self.layer1 = Conv1dReluBn(in_channels, channels, kernel_size=7, padding=3) self.layer2 = SE_Res2Block(channels, kernel_size=5, stride=1, padding=4, dilation=2, scale=8) self.layer3 = SE_Res2Block(channels, kernel_size=5, stride=1, padding=6, dilation=3, scale=8) self.layer4 = SE_Res2Block(channels, kernel_size=5, stride=1, padding=8, dilation=4, scale=8) cat_channels = channels * 3 self.conv = nn.Conv1d(cat_channels, cat_channels, kernel_size=1) 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) def forward(self, x): x = self.torchfb(x) + 1e-6 x = x.log() x = self.instancenorm(x) # print(x.shape) # x = x.transpose(1, 2) 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 class ECAPA_TDNN_L2(nn.Module): def __init__(self, in_channels=80, channels=512, embd_dim=192): super().__init__() self.torchfb = torchaudio.transforms.MelSpectrogram(sample_rate=16000, n_fft=512, win_length=400, hop_length=160, f_min=0.0, f_max=8000, pad=0, n_mels=80) self.instancenorm = nn.InstanceNorm1d(40) self.layer1 = Conv1dReluBn(in_channels, channels, kernel_size=5, padding=2) 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.conv = nn.Conv1d(cat_channels, cat_channels, kernel_size=1) 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) def forward(self, x): x = self.torchfb(x) + 1e-6 x = x.log() x = self.instancenorm(x) # print(x.shape) # x = x.transpose(1, 2) 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)) out_l2 = out / torch.norm(out, dim=1, keepdim=True) return out_l2*512 if __name__ == '__main__': # Input size: batch_size * seq_len * feat_dim 32240 => 202, 35760=>224 x = torch.zeros(32, 35760).cuda() model = ECAPA_TDNN(in_channels=80, channels=512, embd_dim=192) # print(model) summary(model, input_size=(tuple(x.shape))) out = model(x) print(out.shape) # should be [2, 192]