# 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."
)