{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Single-Channel ResNeSt for SED architecture\n---\n\nThis notebook contains code for modifying ResNeSt models for single-channel use. This permits their inclusion as the encoder in a modified PANNs SED architecture. This method was used as part of my submission to the RFCX audio classifier competition.\n\n\nI am indebted to Ryan Epp for his [excellent notebook](https://www.kaggle.com/reppic/mean-teachers-find-more-birds) and to Hidehisa Arai for the [original PANNS SED notebook](https://www.kaggle.com/hidehisaarai1213/introduction-to-sound-event-detection) on Kaggle."},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install torchlibrosa > /dev/null\n!pip install git+https://github.com/zhanghang1989/ResNeSt.git > /dev/null","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Params"},{"metadata":{"trusted":true},"cell_type":"code","source":"RESNEST_TYPE = '50'\n\nRESIZE_DICT = {'50' : 224, \n               '101' : 256, \n               '200' : 320}\n\n\nN_CLASSES = 24\nN_CHANNELS = 1\nRESIZE = RESIZE_DICT[RESNEST_TYPE]\nENCODER_LEN = 2048\nDROPOUT = 0.5\n\n\nFFT = 4096\nHOP = 512\nF_MIN = 60\nF_MAX = 14000\nSR = 36000","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Libraries"},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom torchlibrosa.stft import Spectrogram, LogmelFilterBank\nfrom resnest.torch import resnest50, resnest101, resnest200\nfrom resnest.torch.resnet import ResNet, Bottleneck","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Single-channel pretrained ResNeSt feature extractor\n---\n\nThis model removes the average-pooling, flatten and fully-connected layers from the original ResNeSt model. This allows it to be used as a feature extractor, since the existing code does not have an inbuilt function for this. If you wish to use it as a full CNN model, uncomment the lines in `forward()`.\n\nAdditionally, the initial convolutional layer is modified to allow it to take single-channel image input.\n\nAll credit to the original author *zhanghang1989* at github: https://github.com/zhanghang1989/ResNeSt "},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"class ResNestEncoder(ResNet):\n    def __init__(self):\n        super(ResNestEncoder, self).__init__(Bottleneck, [3, 4, 6, 3], radix=2, groups=1,\n                                          bottleneck_width=64, deep_stem=True, stem_width=32,\n                                          avg_down=True, avd=True, avd_first=False)\n        \n    def forward(self, x):\n        x = self.conv1(x)\n        x = self.bn1(x)\n        x = self.relu(x)\n        x = self.maxpool(x)\n\n        x = self.layer1(x)\n        x = self.layer2(x)\n        x = self.layer3(x)\n        x = self.layer4(x)\n        \n        '''\n        # if using full model\n        x = self.avgpool(x)\n        x = torch.flatten(x, 1)\n        if self.drop:\n            x = self.drop(x)\n        x = self.fc(x)\n        '''\n        \n        return x       \n    \n\ndef get_model():    \n    model = ResNestEncoder()\n    model.load_state_dict(torch.hub.load_state_dict_from_url('https://s3.us-west-1.wasabisys.com/resnest/torch/resnest50-528c19ca.pth',\n                                                             progress=True, check_hash=True))\n    \n    # modify initial convolutional layer to use a single channel\n    model.conv1 = nn.Sequential(\n    nn.Conv2d(N_CHANNELS, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False), #<-- in_channels specified here\n    nn.BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),\n    nn.ReLU(inplace=True),\n    nn.Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),\n    nn.BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),\n    nn.ReLU(inplace=True),\n    nn.Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n  )        \n                \n    return model\n\nresnest_feature_extractor = get_model()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We can test it by passing a single-channel image tensor:"},{"metadata":{"trusted":true},"cell_type":"code","source":"a = torch.rand([2, 1, 64, 64]) # a batch of 2 single-channel images of size 64px by 64px\nresnest_feature_extractor(a)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# PANNs SED Architecture using ResNeSt\n---\nThe only pertinent change is that size of `ENCODER_FEATURES` is different to that of other CNNs. The current value works for `resnest50` - you may have to adjust it for other variants."},{"metadata":{"trusted":true},"cell_type":"code","source":"def init_layer(layer):\n    nn.init.xavier_uniform_(layer.weight)\n\n    if hasattr(layer, \"bias\"):\n        if layer.bias is not None:\n            layer.bias.data.fill_(0.)\n\n\ndef init_bn(bn):\n    bn.bias.data.fill_(0.)\n    bn.weight.data.fill_(1.0)\n\ndef interpolate(x: torch.Tensor, ratio: int):\n    \"\"\"Interpolate data in time domain. This is used to compensate the\n    resolution reduction in downsampling of a CNN.\n\n    Args:\n      x: (batch_size, time_steps, classes_num)\n      ratio: int, ratio to interpolate\n    Returns:\n      upsampled: (batch_size, time_steps * ratio, classes_num)\n    \"\"\"\n    (batch_size, time_steps, classes_num) = x.shape\n    upsampled = x[:, :, None, :].repeat(1, 1, ratio, 1)\n    upsampled = upsampled.reshape(batch_size, time_steps * ratio, classes_num)\n    return upsampled\n\nclass AttentionHead(nn.Module):\n    \n    def __init__(self, in_features: int, out_features: int):\n        super().__init__()\n        self.conv_attention = nn.Conv1d(in_channels=in_features, \n                                        out_channels=out_features,\n                                        kernel_size=1, stride=1, \n                                        padding=0, bias=True)\n        self.conv_classes = nn.Conv1d(in_channels=in_features, \n                                      out_channels=out_features,\n                                      kernel_size=1, stride=1, \n                                      padding=0, bias=True)\n        self.batch_norm_attention = nn.BatchNorm1d(out_features)\n        self.init_weights()\n\n    def init_weights(self):\n        init_layer(self.conv_attention)\n        init_layer(self.conv_classes)\n        init_bn(self.batch_norm_attention)\n\n    def forward(self, x):\n        norm_att = torch.softmax(torch.tanh(self.conv_attention(x)), dim=-1)\n        classes = self.conv_classes(x)\n        x = torch.sum(norm_att * classes, dim=2)\n        return x, norm_att, classes\n\n\nclass SEDAudioClassifier(nn.Module):\n    def __init__(self, sample_rate=SR, n_fft=FFT, hop_length=HOP, \n                 mel_bins=RESIZE, fmin=F_MIN, fmax=F_MAX,\n                 n_classes=N_CLASSES, dropout=DROPOUT):\n        super().__init__()\n        self.interpolate_ratio = 32\n\n        self.spectrogram_extractor = Spectrogram(n_fft=n_fft, \n                                                 hop_length=hop_length,\n                                                 win_length=None, \n                                                 window='hann',\n                                                 center=True,\n                                                 pad_mode='reflect', \n                                                 freeze_parameters=True)\n        self.logmel_extractor = LogmelFilterBank(sr=sample_rate,\n                                                 n_fft=n_fft,\n                                                 n_mels=mel_bins,\n                                                 fmin=fmin, \n                                                 fmax=fmax,\n                                                 ref=1.0, \n                                                 amin=1e-10,\n                                                 top_db=None, \n                                                 freeze_parameters=True)\n\n        self.batch_norm = nn.BatchNorm2d(mel_bins)\n        self.encoder = get_model()\n        self.fc = nn.Linear(ENCODER_LEN, \n                            ENCODER_LEN, bias=True)\n        self.att_head = AttentionHead(ENCODER_LEN, n_classes)\n        self.avg_pool = nn.modules.pooling.AdaptiveAvgPool2d((1, 1))\n        self.dropout = nn.Dropout(dropout)\n        self.init_weight()\n\n    def init_weight(self):\n        init_bn(self.batch_norm)\n        init_layer(self.fc)\n        self.att_head.init_weights()\n\n    def forward(self, input, spec_aug=False, \n                mixup_lambda=None, return_encoding=False):\n        x = self.spectrogram_extractor(input.float())\n        x = self.logmel_extractor(x)\n        \n        x = x.transpose(1, 3)\n        x = self.batch_norm(x)\n        x = x.transpose(1, 3)\n        x = self.encoder(x)\n        x = torch.mean(x, dim=3)\n        x1 = F.max_pool1d(x, kernel_size=3, stride=1, padding=1)\n        x2 = F.avg_pool1d(x, kernel_size=3, stride=1, padding=1)\n        x = x1 + x2\n\n        x = self.dropout(x)\n        x = x.transpose(1, 2)\n        x = F.relu_(self.fc(x))\n        x = x.transpose(1, 2)\n        x = self.dropout(x)\n\n        (clipwise_output, norm_att, segmentwise_output) = self.att_head(x)\n        segmentwise_output = segmentwise_output.transpose(1, 2)\n        framewise_output = interpolate(segmentwise_output, self.interpolate_ratio)\n        \n        return clipwise_output, framewise_output    \n        ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We can test this by passing it an example tensor that is the same shape as a batch of audio vectors. This is the input that PANNs SED expects:"},{"metadata":{"trusted":true},"cell_type":"code","source":"model = SEDAudioClassifier()\na = torch.rand([2, 25600]) # a batch of 2 flat audio tensors\nmodel(a)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Success!"}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}