| 1 | # just for speaker similarity evaluation, third-party code |
| 2 | |
| 3 | # From https://github.com/microsoft/UniSpeech/blob/main/downstreams/speaker_verification/models/ |
| 4 | # part of the code is borrowed from https://github.com/lawlict/ECAPA-TDNN |
| 5 | |
| 6 | import os |
| 7 | |
| 8 | import torch |
| 9 | import torch.nn as nn |
| 10 | import torch.nn.functional as F |
| 11 | |
| 12 | |
| 13 | """ Res2Conv1d + BatchNorm1d + ReLU |
| 14 | """ |
| 15 | |
| 16 | |
| 17 | class Res2Conv1dReluBn(nn.Module): |
| 18 | """ |
| 19 | in_channels == out_channels == channels |
| 20 | """ |
| 21 | |
| 22 | def __init__(self, channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True, scale=4): |
| 23 | super().__init__() |
| 24 | assert channels % scale == 0, "{} % {} != 0".format(channels, scale) |
| 25 | self.scale = scale |
| 26 | self.width = channels // scale |
| 27 | self.nums = scale if scale == 1 else scale - 1 |
| 28 | |
| 29 | self.convs = [] |
| 30 | self.bns = [] |
| 31 | for i in range(self.nums): |
| 32 | self.convs.append(nn.Conv1d(self.width, self.width, kernel_size, stride, padding, dilation, bias=bias)) |
| 33 | self.bns.append(nn.BatchNorm1d(self.width)) |
| 34 | self.convs = nn.ModuleList(self.convs) |
| 35 | self.bns = nn.ModuleList(self.bns) |
| 36 | |
| 37 | def forward(self, x): |
| 38 | out = [] |
| 39 | spx = torch.split(x, self.width, 1) |
| 40 | for i in range(self.nums): |
| 41 | if i == 0: |
| 42 | sp = spx[i] |
| 43 | else: |
| 44 | sp = sp + spx[i] |
| 45 | # Order: conv -> relu -> bn |
| 46 | sp = self.convs[i](sp) |
| 47 | sp = self.bns[i](F.relu(sp)) |
| 48 | out.append(sp) |
| 49 | if self.scale != 1: |
| 50 | out.append(spx[self.nums]) |
| 51 | out = torch.cat(out, dim=1) |
| 52 | |
| 53 | return out |
| 54 | |
| 55 | |
| 56 | """ Conv1d + BatchNorm1d + ReLU |
| 57 | """ |
| 58 | |
| 59 | |
| 60 | class Conv1dReluBn(nn.Module): |
| 61 | def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True): |
| 62 | super().__init__() |
| 63 | self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias) |
| 64 | self.bn = nn.BatchNorm1d(out_channels) |
| 65 | |
| 66 | def forward(self, x): |
| 67 | return self.bn(F.relu(self.conv(x))) |
| 68 | |
| 69 | |
| 70 | """ The SE connection of 1D case. |
| 71 | """ |
| 72 | |
| 73 | |
| 74 | class SE_Connect(nn.Module): |
| 75 | def __init__(self, channels, se_bottleneck_dim=128): |
| 76 | super().__init__() |
| 77 | self.linear1 = nn.Linear(channels, se_bottleneck_dim) |
| 78 | self.linear2 = nn.Linear(se_bottleneck_dim, channels) |
| 79 | |
| 80 | def forward(self, x): |
| 81 | out = x.mean(dim=2) |
| 82 | out = F.relu(self.linear1(out)) |
| 83 | out = torch.sigmoid(self.linear2(out)) |
| 84 | out = x * out.unsqueeze(2) |
| 85 | |
| 86 | return out |
| 87 | |
| 88 | |
| 89 | """ SE-Res2Block of the ECAPA-TDNN architecture. |
| 90 | """ |
| 91 | |
| 92 | # def SE_Res2Block(channels, kernel_size, stride, padding, dilation, scale): |
| 93 | # return nn.Sequential( |
| 94 | # Conv1dReluBn(channels, 512, kernel_size=1, stride=1, padding=0), |
| 95 | # Res2Conv1dReluBn(512, kernel_size, stride, padding, dilation, scale=scale), |
| 96 | # Conv1dReluBn(512, channels, kernel_size=1, stride=1, padding=0), |
| 97 | # SE_Connect(channels) |
| 98 | # ) |
| 99 | |
| 100 | |
| 101 | class SE_Res2Block(nn.Module): |
| 102 | def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, scale, se_bottleneck_dim): |
| 103 | super().__init__() |
| 104 | self.Conv1dReluBn1 = Conv1dReluBn(in_channels, out_channels, kernel_size=1, stride=1, padding=0) |
| 105 | self.Res2Conv1dReluBn = Res2Conv1dReluBn(out_channels, kernel_size, stride, padding, dilation, scale=scale) |
| 106 | self.Conv1dReluBn2 = Conv1dReluBn(out_channels, out_channels, kernel_size=1, stride=1, padding=0) |
| 107 | self.SE_Connect = SE_Connect(out_channels, se_bottleneck_dim) |
| 108 | |
| 109 | self.shortcut = None |
| 110 | if in_channels != out_channels: |
| 111 | self.shortcut = nn.Conv1d( |
| 112 | in_channels=in_channels, |
| 113 | out_channels=out_channels, |
| 114 | kernel_size=1, |
| 115 | ) |
| 116 | |
| 117 | def forward(self, x): |
| 118 | residual = x |
| 119 | if self.shortcut: |
| 120 | residual = self.shortcut(x) |
| 121 | |
| 122 | x = self.Conv1dReluBn1(x) |
| 123 | x = self.Res2Conv1dReluBn(x) |
| 124 | x = self.Conv1dReluBn2(x) |
| 125 | x = self.SE_Connect(x) |
| 126 | |
| 127 | return x + residual |
| 128 | |
| 129 | |
| 130 | """ Attentive weighted mean and standard deviation pooling. |
| 131 | """ |
| 132 | |
| 133 | |
| 134 | class AttentiveStatsPool(nn.Module): |
| 135 | def __init__(self, in_dim, attention_channels=128, global_context_att=False): |
| 136 | super().__init__() |
| 137 | self.global_context_att = global_context_att |
| 138 | |
| 139 | # Use Conv1d with stride == 1 rather than Linear, then we don't need to transpose inputs. |
| 140 | if global_context_att: |
| 141 | self.linear1 = nn.Conv1d(in_dim * 3, attention_channels, kernel_size=1) # equals W and b in the paper |
| 142 | else: |
| 143 | self.linear1 = nn.Conv1d(in_dim, attention_channels, kernel_size=1) # equals W and b in the paper |
| 144 | self.linear2 = nn.Conv1d(attention_channels, in_dim, kernel_size=1) # equals V and k in the paper |
| 145 | |
| 146 | def forward(self, x): |
| 147 | if self.global_context_att: |
| 148 | context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x) |
| 149 | context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x) |
| 150 | x_in = torch.cat((x, context_mean, context_std), dim=1) |
| 151 | else: |
| 152 | x_in = x |
| 153 | |
| 154 | # DON'T use ReLU here! In experiments, I find ReLU hard to converge. |
| 155 | alpha = torch.tanh(self.linear1(x_in)) |
| 156 | # alpha = F.relu(self.linear1(x_in)) |
| 157 | alpha = torch.softmax(self.linear2(alpha), dim=2) |
| 158 | mean = torch.sum(alpha * x, dim=2) |
| 159 | residuals = torch.sum(alpha * (x**2), dim=2) - mean**2 |
| 160 | std = torch.sqrt(residuals.clamp(min=1e-9)) |
| 161 | return torch.cat([mean, std], dim=1) |
| 162 | |
| 163 | |
| 164 | class ECAPA_TDNN(nn.Module): |
| 165 | def __init__( |
| 166 | self, |
| 167 | feat_dim=80, |
| 168 | channels=512, |
| 169 | emb_dim=192, |
| 170 | global_context_att=False, |
| 171 | feat_type="wavlm_large", |
| 172 | sr=16000, |
| 173 | feature_selection="hidden_states", |
| 174 | update_extract=False, |
| 175 | config_path=None, |
| 176 | ): |
| 177 | super().__init__() |
| 178 | |
| 179 | self.feat_type = feat_type |
| 180 | self.feature_selection = feature_selection |
| 181 | self.update_extract = update_extract |
| 182 | self.sr = sr |
| 183 | |
| 184 | torch.hub._validate_not_a_forked_repo = lambda a, b, c: True |
| 185 | try: |
| 186 | local_s3prl_path = os.path.expanduser("~/.cache/torch/hub/s3prl_s3prl_main") |
| 187 | self.feature_extract = torch.hub.load(local_s3prl_path, feat_type, source="local", config_path=config_path) |
| 188 | except: # noqa: E722 |
| 189 | self.feature_extract = torch.hub.load("s3prl/s3prl", feat_type) |
| 190 | |
| 191 | if len(self.feature_extract.model.encoder.layers) == 24 and hasattr( |
| 192 | self.feature_extract.model.encoder.layers[23].self_attn, "fp32_attention" |
| 193 | ): |
| 194 | self.feature_extract.model.encoder.layers[23].self_attn.fp32_attention = False |
| 195 | if len(self.feature_extract.model.encoder.layers) == 24 and hasattr( |
| 196 | self.feature_extract.model.encoder.layers[11].self_attn, "fp32_attention" |
| 197 | ): |
| 198 | self.feature_extract.model.encoder.layers[11].self_attn.fp32_attention = False |
| 199 | |
| 200 | self.feat_num = self.get_feat_num() |
| 201 | self.feature_weight = nn.Parameter(torch.zeros(self.feat_num)) |
| 202 | |
| 203 | if feat_type != "fbank" and feat_type != "mfcc": |
| 204 | freeze_list = ["final_proj", "label_embs_concat", "mask_emb", "project_q", "quantizer"] |
| 205 | for name, param in self.feature_extract.named_parameters(): |
| 206 | for freeze_val in freeze_list: |
| 207 | if freeze_val in name: |
| 208 | param.requires_grad = False |
| 209 | break |
| 210 | |
| 211 | if not self.update_extract: |
| 212 | for param in self.feature_extract.parameters(): |
| 213 | param.requires_grad = False |
| 214 | |
| 215 | self.instance_norm = nn.InstanceNorm1d(feat_dim) |
| 216 | # self.channels = [channels] * 4 + [channels * 3] |
| 217 | self.channels = [channels] * 4 + [1536] |
| 218 | |
| 219 | self.layer1 = Conv1dReluBn(feat_dim, self.channels[0], kernel_size=5, padding=2) |
| 220 | self.layer2 = SE_Res2Block( |
| 221 | self.channels[0], |
| 222 | self.channels[1], |
| 223 | kernel_size=3, |
| 224 | stride=1, |
| 225 | padding=2, |
| 226 | dilation=2, |
| 227 | scale=8, |
| 228 | se_bottleneck_dim=128, |
| 229 | ) |
| 230 | self.layer3 = SE_Res2Block( |
| 231 | self.channels[1], |
| 232 | self.channels[2], |
| 233 | kernel_size=3, |
| 234 | stride=1, |
| 235 | padding=3, |
| 236 | dilation=3, |
| 237 | scale=8, |
| 238 | se_bottleneck_dim=128, |
| 239 | ) |
| 240 | self.layer4 = SE_Res2Block( |
| 241 | self.channels[2], |
| 242 | self.channels[3], |
| 243 | kernel_size=3, |
| 244 | stride=1, |
| 245 | padding=4, |
| 246 | dilation=4, |
| 247 | scale=8, |
| 248 | se_bottleneck_dim=128, |
| 249 | ) |
| 250 | |
| 251 | # self.conv = nn.Conv1d(self.channels[-1], self.channels[-1], kernel_size=1) |
| 252 | cat_channels = channels * 3 |
| 253 | self.conv = nn.Conv1d(cat_channels, self.channels[-1], kernel_size=1) |
| 254 | self.pooling = AttentiveStatsPool( |
| 255 | self.channels[-1], attention_channels=128, global_context_att=global_context_att |
| 256 | ) |
| 257 | self.bn = nn.BatchNorm1d(self.channels[-1] * 2) |
| 258 | self.linear = nn.Linear(self.channels[-1] * 2, emb_dim) |
| 259 | |
| 260 | def get_feat_num(self): |
| 261 | self.feature_extract.eval() |
| 262 | wav = [torch.randn(self.sr).to(next(self.feature_extract.parameters()).device)] |
| 263 | with torch.no_grad(): |
| 264 | features = self.feature_extract(wav) |
| 265 | select_feature = features[self.feature_selection] |
| 266 | if isinstance(select_feature, (list, tuple)): |
| 267 | return len(select_feature) |
| 268 | else: |
| 269 | return 1 |
| 270 | |
| 271 | def get_feat(self, x): |
| 272 | if self.update_extract: |
| 273 | x = self.feature_extract([sample for sample in x]) |
| 274 | else: |
| 275 | with torch.no_grad(): |
| 276 | if self.feat_type == "fbank" or self.feat_type == "mfcc": |
| 277 | x = self.feature_extract(x) + 1e-6 # B x feat_dim x time_len |
| 278 | else: |
| 279 | x = self.feature_extract([sample for sample in x]) |
| 280 | |
| 281 | if self.feat_type == "fbank": |
| 282 | x = x.log() |
| 283 | |
| 284 | if self.feat_type != "fbank" and self.feat_type != "mfcc": |
| 285 | x = x[self.feature_selection] |
| 286 | if isinstance(x, (list, tuple)): |
| 287 | x = torch.stack(x, dim=0) |
| 288 | else: |
| 289 | x = x.unsqueeze(0) |
| 290 | norm_weights = F.softmax(self.feature_weight, dim=-1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) |
| 291 | x = (norm_weights * x).sum(dim=0) |
| 292 | x = torch.transpose(x, 1, 2) + 1e-6 |
| 293 | |
| 294 | x = self.instance_norm(x) |
| 295 | return x |
| 296 | |
| 297 | def forward(self, x): |
| 298 | x = self.get_feat(x) |
| 299 | |
| 300 | out1 = self.layer1(x) |
| 301 | out2 = self.layer2(out1) |
| 302 | out3 = self.layer3(out2) |
| 303 | out4 = self.layer4(out3) |
| 304 | |
| 305 | out = torch.cat([out2, out3, out4], dim=1) |
| 306 | out = F.relu(self.conv(out)) |
| 307 | out = self.bn(self.pooling(out)) |
| 308 | out = self.linear(out) |
| 309 | |
| 310 | return out |
| 311 | |
| 312 | |
| 313 | def ECAPA_TDNN_SMALL( |
| 314 | feat_dim, |
| 315 | emb_dim=256, |
| 316 | feat_type="wavlm_large", |
| 317 | sr=16000, |
| 318 | feature_selection="hidden_states", |
| 319 | update_extract=False, |
| 320 | config_path=None, |
| 321 | ): |
| 322 | return ECAPA_TDNN( |
| 323 | feat_dim=feat_dim, |
| 324 | channels=512, |
| 325 | emb_dim=emb_dim, |
| 326 | feat_type=feat_type, |
| 327 | sr=sr, |
| 328 | feature_selection=feature_selection, |
| 329 | update_extract=update_extract, |
| 330 | config_path=config_path, |
| 331 | ) |
| 332 |