Spaces:
Running
Running
| # PANNs: Large-Scale Pretrained Audio Neural Networks for Audio Pattern Recognition | |
| # Reference from https://github.com/qiuqiangkong/audioset_tagging_cnn | |
| # Some layers are re-designed for CLAP | |
| import os | |
| os.environ["NUMBA_CACHE_DIR"] = "/tmp/" | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torchlibrosa.stft import Spectrogram, LogmelFilterBank | |
| from torchlibrosa.augmentation import SpecAugmentation | |
| from .utils import do_mixup, interpolate, pad_framewise_output | |
| from .feature_fusion import iAFF, AFF, DAF | |
| def init_layer(layer): | |
| """Initialize a Linear or Convolutional layer.""" | |
| nn.init.xavier_uniform_(layer.weight) | |
| if hasattr(layer, "bias"): | |
| if layer.bias is not None: | |
| layer.bias.data.fill_(0.0) | |
| def init_bn(bn): | |
| """Initialize a Batchnorm layer.""" | |
| bn.bias.data.fill_(0.0) | |
| bn.weight.data.fill_(1.0) | |
| class ConvBlock(nn.Module): | |
| def __init__(self, in_channels, out_channels): | |
| super(ConvBlock, self).__init__() | |
| self.conv1 = nn.Conv2d( | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=(1, 1), | |
| bias=False, | |
| ) | |
| self.conv2 = nn.Conv2d( | |
| in_channels=out_channels, | |
| out_channels=out_channels, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=(1, 1), | |
| bias=False, | |
| ) | |
| self.bn1 = nn.BatchNorm2d(out_channels) | |
| self.bn2 = nn.BatchNorm2d(out_channels) | |
| self.init_weight() | |
| def init_weight(self): | |
| init_layer(self.conv1) | |
| init_layer(self.conv2) | |
| init_bn(self.bn1) | |
| init_bn(self.bn2) | |
| def forward(self, input, pool_size=(2, 2), pool_type="avg"): | |
| x = input | |
| x = F.relu_(self.bn1(self.conv1(x))) | |
| x = F.relu_(self.bn2(self.conv2(x))) | |
| if pool_type == "max": | |
| x = F.max_pool2d(x, kernel_size=pool_size) | |
| elif pool_type == "avg": | |
| x = F.avg_pool2d(x, kernel_size=pool_size) | |
| elif pool_type == "avg+max": | |
| x1 = F.avg_pool2d(x, kernel_size=pool_size) | |
| x2 = F.max_pool2d(x, kernel_size=pool_size) | |
| x = x1 + x2 | |
| else: | |
| raise Exception("Incorrect argument!") | |
| return x | |
| class ConvBlock5x5(nn.Module): | |
| def __init__(self, in_channels, out_channels): | |
| super(ConvBlock5x5, self).__init__() | |
| self.conv1 = nn.Conv2d( | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| kernel_size=(5, 5), | |
| stride=(1, 1), | |
| padding=(2, 2), | |
| bias=False, | |
| ) | |
| self.bn1 = nn.BatchNorm2d(out_channels) | |
| self.init_weight() | |
| def init_weight(self): | |
| init_layer(self.conv1) | |
| init_bn(self.bn1) | |
| def forward(self, input, pool_size=(2, 2), pool_type="avg"): | |
| x = input | |
| x = F.relu_(self.bn1(self.conv1(x))) | |
| if pool_type == "max": | |
| x = F.max_pool2d(x, kernel_size=pool_size) | |
| elif pool_type == "avg": | |
| x = F.avg_pool2d(x, kernel_size=pool_size) | |
| elif pool_type == "avg+max": | |
| x1 = F.avg_pool2d(x, kernel_size=pool_size) | |
| x2 = F.max_pool2d(x, kernel_size=pool_size) | |
| x = x1 + x2 | |
| else: | |
| raise Exception("Incorrect argument!") | |
| return x | |
| class AttBlock(nn.Module): | |
| def __init__(self, n_in, n_out, activation="linear", temperature=1.0): | |
| super(AttBlock, self).__init__() | |
| self.activation = activation | |
| self.temperature = temperature | |
| self.att = nn.Conv1d( | |
| in_channels=n_in, | |
| out_channels=n_out, | |
| kernel_size=1, | |
| stride=1, | |
| padding=0, | |
| bias=True, | |
| ) | |
| self.cla = nn.Conv1d( | |
| in_channels=n_in, | |
| out_channels=n_out, | |
| kernel_size=1, | |
| stride=1, | |
| padding=0, | |
| bias=True, | |
| ) | |
| self.bn_att = nn.BatchNorm1d(n_out) | |
| self.init_weights() | |
| def init_weights(self): | |
| init_layer(self.att) | |
| init_layer(self.cla) | |
| init_bn(self.bn_att) | |
| def forward(self, x): | |
| # x: (n_samples, n_in, n_time) | |
| norm_att = torch.softmax(torch.clamp(self.att(x), -10, 10), dim=-1) | |
| cla = self.nonlinear_transform(self.cla(x)) | |
| x = torch.sum(norm_att * cla, dim=2) | |
| return x, norm_att, cla | |
| def nonlinear_transform(self, x): | |
| if self.activation == "linear": | |
| return x | |
| elif self.activation == "sigmoid": | |
| return torch.sigmoid(x) | |
| class Cnn14(nn.Module): | |
| def __init__( | |
| self, | |
| sample_rate, | |
| window_size, | |
| hop_size, | |
| mel_bins, | |
| fmin, | |
| fmax, | |
| classes_num, | |
| enable_fusion=False, | |
| fusion_type="None", | |
| ): | |
| super(Cnn14, self).__init__() | |
| window = "hann" | |
| center = True | |
| pad_mode = "reflect" | |
| ref = 1.0 | |
| amin = 1e-10 | |
| top_db = None | |
| self.enable_fusion = enable_fusion | |
| self.fusion_type = fusion_type | |
| # Spectrogram extractor | |
| self.spectrogram_extractor = Spectrogram( | |
| n_fft=window_size, | |
| hop_length=hop_size, | |
| win_length=window_size, | |
| window=window, | |
| center=center, | |
| pad_mode=pad_mode, | |
| freeze_parameters=True, | |
| ) | |
| # Logmel feature extractor | |
| self.logmel_extractor = LogmelFilterBank( | |
| sr=sample_rate, | |
| n_fft=window_size, | |
| n_mels=mel_bins, | |
| fmin=fmin, | |
| fmax=fmax, | |
| ref=ref, | |
| amin=amin, | |
| top_db=top_db, | |
| freeze_parameters=True, | |
| ) | |
| # Spec augmenter | |
| self.spec_augmenter = SpecAugmentation( | |
| time_drop_width=64, | |
| time_stripes_num=2, | |
| freq_drop_width=8, | |
| freq_stripes_num=2, | |
| ) | |
| self.bn0 = nn.BatchNorm2d(64) | |
| if (self.enable_fusion) and (self.fusion_type == "channel_map"): | |
| self.conv_block1 = ConvBlock(in_channels=4, out_channels=64) | |
| else: | |
| self.conv_block1 = ConvBlock(in_channels=1, out_channels=64) | |
| self.conv_block2 = ConvBlock(in_channels=64, out_channels=128) | |
| self.conv_block3 = ConvBlock(in_channels=128, out_channels=256) | |
| self.conv_block4 = ConvBlock(in_channels=256, out_channels=512) | |
| self.conv_block5 = ConvBlock(in_channels=512, out_channels=1024) | |
| self.conv_block6 = ConvBlock(in_channels=1024, out_channels=2048) | |
| self.fc1 = nn.Linear(2048, 2048, bias=True) | |
| self.fc_audioset = nn.Linear(2048, classes_num, bias=True) | |
| if (self.enable_fusion) and ( | |
| self.fusion_type in ["daf_1d", "aff_1d", "iaff_1d"] | |
| ): | |
| self.mel_conv1d = nn.Sequential( | |
| nn.Conv1d(64, 64, kernel_size=5, stride=3, padding=2), | |
| nn.BatchNorm1d(64), # No Relu | |
| ) | |
| if self.fusion_type == "daf_1d": | |
| self.fusion_model = DAF() | |
| elif self.fusion_type == "aff_1d": | |
| self.fusion_model = AFF(channels=64, type="1D") | |
| elif self.fusion_type == "iaff_1d": | |
| self.fusion_model = iAFF(channels=64, type="1D") | |
| if (self.enable_fusion) and ( | |
| self.fusion_type in ["daf_2d", "aff_2d", "iaff_2d"] | |
| ): | |
| self.mel_conv2d = nn.Sequential( | |
| nn.Conv2d(1, 64, kernel_size=(5, 5), stride=(6, 2), padding=(2, 2)), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True), | |
| ) | |
| if self.fusion_type == "daf_2d": | |
| self.fusion_model = DAF() | |
| elif self.fusion_type == "aff_2d": | |
| self.fusion_model = AFF(channels=64, type="2D") | |
| elif self.fusion_type == "iaff_2d": | |
| self.fusion_model = iAFF(channels=64, type="2D") | |
| self.init_weight() | |
| def init_weight(self): | |
| init_bn(self.bn0) | |
| init_layer(self.fc1) | |
| init_layer(self.fc_audioset) | |
| def forward(self, input, mixup_lambda=None, device=None): | |
| """ | |
| Input: (batch_size, data_length)""" | |
| if self.enable_fusion and input["longer"].sum() == 0: | |
| # if no audio is longer than 10s, then randomly select one audio to be longer | |
| input["longer"][torch.randint(0, input["longer"].shape[0], (1,))] = True | |
| if not self.enable_fusion: | |
| x = self.spectrogram_extractor( | |
| input["waveform"].to(device=device, non_blocking=True) | |
| ) # (batch_size, 1, time_steps, freq_bins) | |
| x = self.logmel_extractor(x) # (batch_size, 1, time_steps, mel_bins) | |
| x = x.transpose(1, 3) | |
| x = self.bn0(x) | |
| x = x.transpose(1, 3) | |
| else: | |
| longer_list = input["longer"].to(device=device, non_blocking=True) | |
| x = input["mel_fusion"].to(device=device, non_blocking=True) | |
| longer_list_idx = torch.where(longer_list)[0] | |
| x = x.transpose(1, 3) | |
| x = self.bn0(x) | |
| x = x.transpose(1, 3) | |
| if self.fusion_type in ["daf_1d", "aff_1d", "iaff_1d"]: | |
| new_x = x[:, 0:1, :, :].clone().contiguous() | |
| # local processing | |
| if len(longer_list_idx) > 0: | |
| fusion_x_local = x[longer_list_idx, 1:, :, :].clone().contiguous() | |
| FB, FC, FT, FF = fusion_x_local.size() | |
| fusion_x_local = fusion_x_local.view(FB * FC, FT, FF) | |
| fusion_x_local = torch.permute( | |
| fusion_x_local, (0, 2, 1) | |
| ).contiguous() | |
| fusion_x_local = self.mel_conv1d(fusion_x_local) | |
| fusion_x_local = fusion_x_local.view( | |
| FB, FC, FF, fusion_x_local.size(-1) | |
| ) | |
| fusion_x_local = ( | |
| torch.permute(fusion_x_local, (0, 2, 1, 3)) | |
| .contiguous() | |
| .flatten(2) | |
| ) | |
| if fusion_x_local.size(-1) < FT: | |
| fusion_x_local = torch.cat( | |
| [ | |
| fusion_x_local, | |
| torch.zeros( | |
| (FB, FF, FT - fusion_x_local.size(-1)), | |
| device=device, | |
| ), | |
| ], | |
| dim=-1, | |
| ) | |
| else: | |
| fusion_x_local = fusion_x_local[:, :, :FT] | |
| # 1D fusion | |
| new_x = new_x.squeeze(1).permute((0, 2, 1)).contiguous() | |
| new_x[longer_list_idx] = self.fusion_model( | |
| new_x[longer_list_idx], fusion_x_local | |
| ) | |
| x = new_x.permute((0, 2, 1)).contiguous()[:, None, :, :] | |
| else: | |
| x = new_x | |
| elif self.fusion_type in ["daf_2d", "aff_2d", "iaff_2d", "channel_map"]: | |
| x = x # no change | |
| if self.training: | |
| x = self.spec_augmenter(x) | |
| # Mixup on spectrogram | |
| if self.training and mixup_lambda is not None: | |
| x = do_mixup(x, mixup_lambda) | |
| if (self.enable_fusion) and ( | |
| self.fusion_type in ["daf_2d", "aff_2d", "iaff_2d"] | |
| ): | |
| global_x = x[:, 0:1, :, :] | |
| # global processing | |
| B, C, H, W = global_x.shape | |
| global_x = self.conv_block1(global_x, pool_size=(2, 2), pool_type="avg") | |
| if len(longer_list_idx) > 0: | |
| local_x = x[longer_list_idx, 1:, :, :].contiguous() | |
| TH = global_x.size(-2) | |
| # local processing | |
| B, C, H, W = local_x.shape | |
| local_x = local_x.view(B * C, 1, H, W) | |
| local_x = self.mel_conv2d(local_x) | |
| local_x = local_x.view( | |
| B, C, local_x.size(1), local_x.size(2), local_x.size(3) | |
| ) | |
| local_x = local_x.permute((0, 2, 1, 3, 4)).contiguous().flatten(2, 3) | |
| TB, TC, _, TW = local_x.size() | |
| if local_x.size(-2) < TH: | |
| local_x = torch.cat( | |
| [ | |
| local_x, | |
| torch.zeros( | |
| (TB, TC, TH - local_x.size(-2), TW), | |
| device=global_x.device, | |
| ), | |
| ], | |
| dim=-2, | |
| ) | |
| else: | |
| local_x = local_x[:, :, :TH, :] | |
| global_x[longer_list_idx] = self.fusion_model( | |
| global_x[longer_list_idx], local_x | |
| ) | |
| x = global_x | |
| else: | |
| x = self.conv_block1(x, pool_size=(2, 2), pool_type="avg") | |
| x = F.dropout(x, p=0.2, training=self.training) | |
| x = self.conv_block2(x, pool_size=(2, 2), pool_type="avg") | |
| x = F.dropout(x, p=0.2, training=self.training) | |
| x = self.conv_block3(x, pool_size=(2, 2), pool_type="avg") | |
| x = F.dropout(x, p=0.2, training=self.training) | |
| x = self.conv_block4(x, pool_size=(2, 2), pool_type="avg") | |
| x = F.dropout(x, p=0.2, training=self.training) | |
| x = self.conv_block5(x, pool_size=(2, 2), pool_type="avg") | |
| x = F.dropout(x, p=0.2, training=self.training) | |
| x = self.conv_block6(x, pool_size=(1, 1), pool_type="avg") | |
| x = F.dropout(x, p=0.2, training=self.training) | |
| x = torch.mean(x, dim=3) | |
| latent_x1 = F.max_pool1d(x, kernel_size=3, stride=1, padding=1) | |
| latent_x2 = F.avg_pool1d(x, kernel_size=3, stride=1, padding=1) | |
| latent_x = latent_x1 + latent_x2 | |
| latent_x = latent_x.transpose(1, 2) | |
| latent_x = F.relu_(self.fc1(latent_x)) | |
| latent_output = interpolate(latent_x, 32) | |
| (x1, _) = torch.max(x, dim=2) | |
| x2 = torch.mean(x, dim=2) | |
| x = x1 + x2 | |
| x = F.dropout(x, p=0.5, training=self.training) | |
| x = F.relu_(self.fc1(x)) | |
| embedding = F.dropout(x, p=0.5, training=self.training) | |
| clipwise_output = torch.sigmoid(self.fc_audioset(x)) | |
| output_dict = { | |
| "clipwise_output": clipwise_output, | |
| "embedding": embedding, | |
| "fine_grained_embedding": latent_output, | |
| } | |
| return output_dict | |
| class Cnn6(nn.Module): | |
| def __init__( | |
| self, | |
| sample_rate, | |
| window_size, | |
| hop_size, | |
| mel_bins, | |
| fmin, | |
| fmax, | |
| classes_num, | |
| enable_fusion=False, | |
| fusion_type="None", | |
| ): | |
| super(Cnn6, self).__init__() | |
| window = "hann" | |
| center = True | |
| pad_mode = "reflect" | |
| ref = 1.0 | |
| amin = 1e-10 | |
| top_db = None | |
| self.enable_fusion = enable_fusion | |
| self.fusion_type = fusion_type | |
| # Spectrogram extractor | |
| self.spectrogram_extractor = Spectrogram( | |
| n_fft=window_size, | |
| hop_length=hop_size, | |
| win_length=window_size, | |
| window=window, | |
| center=center, | |
| pad_mode=pad_mode, | |
| freeze_parameters=True, | |
| ) | |
| # Logmel feature extractor | |
| self.logmel_extractor = LogmelFilterBank( | |
| sr=sample_rate, | |
| n_fft=window_size, | |
| n_mels=mel_bins, | |
| fmin=fmin, | |
| fmax=fmax, | |
| ref=ref, | |
| amin=amin, | |
| top_db=top_db, | |
| freeze_parameters=True, | |
| ) | |
| # Spec augmenter | |
| self.spec_augmenter = SpecAugmentation( | |
| time_drop_width=64, | |
| time_stripes_num=2, | |
| freq_drop_width=8, | |
| freq_stripes_num=2, | |
| ) | |
| self.bn0 = nn.BatchNorm2d(64) | |
| self.conv_block1 = ConvBlock5x5(in_channels=1, out_channels=64) | |
| self.conv_block2 = ConvBlock5x5(in_channels=64, out_channels=128) | |
| self.conv_block3 = ConvBlock5x5(in_channels=128, out_channels=256) | |
| self.conv_block4 = ConvBlock5x5(in_channels=256, out_channels=512) | |
| self.fc1 = nn.Linear(512, 512, bias=True) | |
| self.fc_audioset = nn.Linear(512, classes_num, bias=True) | |
| self.init_weight() | |
| def init_weight(self): | |
| init_bn(self.bn0) | |
| init_layer(self.fc1) | |
| init_layer(self.fc_audioset) | |
| def forward(self, input, mixup_lambda=None, device=None): | |
| """ | |
| Input: (batch_size, data_length)""" | |
| x = self.spectrogram_extractor(input) # (batch_size, 1, time_steps, freq_bins) | |
| x = self.logmel_extractor(x) # (batch_size, 1, time_steps, mel_bins) | |
| x = x.transpose(1, 3) | |
| x = self.bn0(x) | |
| x = x.transpose(1, 3) | |
| if self.training: | |
| x = self.spec_augmenter(x) | |
| # Mixup on spectrogram | |
| if self.training and mixup_lambda is not None: | |
| x = do_mixup(x, mixup_lambda) | |
| x = self.conv_block1(x, pool_size=(2, 2), pool_type="avg") | |
| x = F.dropout(x, p=0.2, training=self.training) | |
| x = self.conv_block2(x, pool_size=(2, 2), pool_type="avg") | |
| x = F.dropout(x, p=0.2, training=self.training) | |
| x = self.conv_block3(x, pool_size=(2, 2), pool_type="avg") | |
| x = F.dropout(x, p=0.2, training=self.training) | |
| x = self.conv_block4(x, pool_size=(2, 2), pool_type="avg") | |
| x = F.dropout(x, p=0.2, training=self.training) | |
| x = torch.mean(x, dim=3) | |
| latent_x1 = F.max_pool1d(x, kernel_size=3, stride=1, padding=1) | |
| latent_x2 = F.avg_pool1d(x, kernel_size=3, stride=1, padding=1) | |
| latent_x = latent_x1 + latent_x2 | |
| latent_x = latent_x.transpose(1, 2) | |
| latent_x = F.relu_(self.fc1(latent_x)) | |
| latent_output = interpolate(latent_x, 16) | |
| (x1, _) = torch.max(x, dim=2) | |
| x2 = torch.mean(x, dim=2) | |
| x = x1 + x2 | |
| x = F.dropout(x, p=0.5, training=self.training) | |
| x = F.relu_(self.fc1(x)) | |
| embedding = F.dropout(x, p=0.5, training=self.training) | |
| clipwise_output = torch.sigmoid(self.fc_audioset(x)) | |
| output_dict = { | |
| "clipwise_output": clipwise_output, | |
| "embedding": embedding, | |
| "fine_grained_embedding": latent_output, | |
| } | |
| return output_dict | |
| class Cnn10(nn.Module): | |
| def __init__( | |
| self, | |
| sample_rate, | |
| window_size, | |
| hop_size, | |
| mel_bins, | |
| fmin, | |
| fmax, | |
| classes_num, | |
| enable_fusion=False, | |
| fusion_type="None", | |
| ): | |
| super(Cnn10, self).__init__() | |
| window = "hann" | |
| center = True | |
| pad_mode = "reflect" | |
| ref = 1.0 | |
| amin = 1e-10 | |
| top_db = None | |
| self.enable_fusion = enable_fusion | |
| self.fusion_type = fusion_type | |
| # Spectrogram extractor | |
| self.spectrogram_extractor = Spectrogram( | |
| n_fft=window_size, | |
| hop_length=hop_size, | |
| win_length=window_size, | |
| window=window, | |
| center=center, | |
| pad_mode=pad_mode, | |
| freeze_parameters=True, | |
| ) | |
| # Logmel feature extractor | |
| self.logmel_extractor = LogmelFilterBank( | |
| sr=sample_rate, | |
| n_fft=window_size, | |
| n_mels=mel_bins, | |
| fmin=fmin, | |
| fmax=fmax, | |
| ref=ref, | |
| amin=amin, | |
| top_db=top_db, | |
| freeze_parameters=True, | |
| ) | |
| # Spec augmenter | |
| self.spec_augmenter = SpecAugmentation( | |
| time_drop_width=64, | |
| time_stripes_num=2, | |
| freq_drop_width=8, | |
| freq_stripes_num=2, | |
| ) | |
| self.bn0 = nn.BatchNorm2d(64) | |
| self.conv_block1 = ConvBlock(in_channels=1, out_channels=64) | |
| self.conv_block2 = ConvBlock(in_channels=64, out_channels=128) | |
| self.conv_block3 = ConvBlock(in_channels=128, out_channels=256) | |
| self.conv_block4 = ConvBlock(in_channels=256, out_channels=512) | |
| self.conv_block5 = ConvBlock(in_channels=512, out_channels=1024) | |
| self.fc1 = nn.Linear(1024, 1024, bias=True) | |
| self.fc_audioset = nn.Linear(1024, classes_num, bias=True) | |
| self.init_weight() | |
| def init_weight(self): | |
| init_bn(self.bn0) | |
| init_layer(self.fc1) | |
| init_layer(self.fc_audioset) | |
| def forward(self, input, mixup_lambda=None, device=None): | |
| """ | |
| Input: (batch_size, data_length)""" | |
| x = self.spectrogram_extractor(input) # (batch_size, 1, time_steps, freq_bins) | |
| x = self.logmel_extractor(x) # (batch_size, 1, time_steps, mel_bins) | |
| x = x.transpose(1, 3) | |
| x = self.bn0(x) | |
| x = x.transpose(1, 3) | |
| if self.training: | |
| x = self.spec_augmenter(x) | |
| # Mixup on spectrogram | |
| if self.training and mixup_lambda is not None: | |
| x = do_mixup(x, mixup_lambda) | |
| x = self.conv_block1(x, pool_size=(2, 2), pool_type="avg") | |
| x = F.dropout(x, p=0.2, training=self.training) | |
| x = self.conv_block2(x, pool_size=(2, 2), pool_type="avg") | |
| x = F.dropout(x, p=0.2, training=self.training) | |
| x = self.conv_block3(x, pool_size=(2, 2), pool_type="avg") | |
| x = F.dropout(x, p=0.2, training=self.training) | |
| x = self.conv_block4(x, pool_size=(2, 2), pool_type="avg") | |
| x = F.dropout(x, p=0.2, training=self.training) | |
| x = self.conv_block5(x, pool_size=(2, 2), pool_type="avg") | |
| x = F.dropout(x, p=0.2, training=self.training) | |
| x = torch.mean(x, dim=3) | |
| latent_x1 = F.max_pool1d(x, kernel_size=3, stride=1, padding=1) | |
| latent_x2 = F.avg_pool1d(x, kernel_size=3, stride=1, padding=1) | |
| latent_x = latent_x1 + latent_x2 | |
| latent_x = latent_x.transpose(1, 2) | |
| latent_x = F.relu_(self.fc1(latent_x)) | |
| latent_output = interpolate(latent_x, 32) | |
| (x1, _) = torch.max(x, dim=2) | |
| x2 = torch.mean(x, dim=2) | |
| x = x1 + x2 | |
| x = F.dropout(x, p=0.5, training=self.training) | |
| x = F.relu_(self.fc1(x)) | |
| embedding = F.dropout(x, p=0.5, training=self.training) | |
| clipwise_output = torch.sigmoid(self.fc_audioset(x)) | |
| output_dict = { | |
| "clipwise_output": clipwise_output, | |
| "embedding": embedding, | |
| "fine_grained_embedding": latent_output, | |
| } | |
| return output_dict | |
| def create_pann_model(audio_cfg, enable_fusion=False, fusion_type="None"): | |
| try: | |
| ModelProto = eval(audio_cfg.model_name) | |
| model = ModelProto( | |
| sample_rate=audio_cfg.sample_rate, | |
| window_size=audio_cfg.window_size, | |
| hop_size=audio_cfg.hop_size, | |
| mel_bins=audio_cfg.mel_bins, | |
| fmin=audio_cfg.fmin, | |
| fmax=audio_cfg.fmax, | |
| classes_num=audio_cfg.class_num, | |
| enable_fusion=enable_fusion, | |
| fusion_type=fusion_type, | |
| ) | |
| return model | |
| except: | |
| raise RuntimeError( | |
| f"Import Model for {audio_cfg.model_name} not found, or the audio cfg parameters are not enough." | |
| ) | |