{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Introduction\nI wanted to share something that worked pretty well for me early on in this competition. The idea comes from a [2018 paper](https://arxiv.org/pdf/1703.01780.pdf) titled *Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results* by Antti Tarvainen and Harri Valpola. \n\n### Mean Teacher\nBiefly, the idea is to use two models. A student model with weights trained the standard way, using backprop. And a teacher model with weights that are an exponential moving average of the student's weights. The teacher is the *mean* of the student \\*ba dum tss\\*. The student is then trained using two different losses, a standard classification loss and a consistency loss that penalizes student predictions that deviate from the teaher's. \n\n![](https://raw.githubusercontent.com/CuriousAI/mean-teacher/master/mean_teacher.png)\n\nMean teachers are useful in a semi-supervised context where we have both labeled and unlabeled samples. The consistency loss on the unlabeled samples acts as a form of regularization and helps the model generalize better. As an added bonus the final teacher model is a temporal ensemble which tends to perform better than the results at the end of a single epoch. \n\n### Missing Labels\nAs a few others have pointed out, there are a lot of missing labels. If we were to randomly sample a segment from the training data, we might consider it completely unlabeled rather than rely on the provided labels. We'll train our mean teacher model(s) on two classes of data, carefully selected positive samples and randomly selected unlabeled samples. The classification loss won't apply to the unlabeled samples. \n\n![](https://www.googleapis.com/download/storage/v1/b/kaggle-forum-message-attachments/o/inbox%2F4704212%2F9ca088bb386abf7114543c019c1d8a5f%2Ffig.png?generation=1609892974092435&alt=media)\n\n*Thanks to [shinmura0](https://www.kaggle.com/shinmurashinmura) for the great visualization!*\n\n### Results\nFor me, mean teacher worked a good bit better than baseline models with similar configurations. \n\n|                                         | Baseline | Mean Teacher |\n|-----------------------------------------|----------|--------------|\n| Well Tuned, 5 fold, from my local setup | 0.847        | **0.865**            |\n| Single fold Expt1 on Kaggle                   | 0.592**        | **0.786**            |\n| Single fold Expt2 on Kaggle                   | 0.826        | **0.830**            |\n| 5 Fold on Kaggle***                        | 0.844        | **0.857**           |\n\n\\*\\* I might have accidentally sabatoged this run.\n\n\\*\\*\\* There was a major bug in v21 of the notebook where the consistence_ramp was set to 1000 which means it was just normal / non-mean-teacher training. Setting consisteny_ramp to 6 and using the mean teacher, we get an improvement of 0.13."},{"metadata":{"trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"!pip -q install --upgrade pip\n!pip -q install timm\n!pip -q install torchlibrosa\n!pip -q install audiomentations","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"import audiomentations as A\nimport os, time, librosa, random\nfrom functools import partial\nimport numpy as np\nimport pandas as pd\nfrom sklearn.model_selection import StratifiedKFold\nimport torch\nimport torch.nn.functional as F\nimport torch.nn as nn\nfrom torch.utils.data import Dataset\nfrom timm.models import resnet34d\nfrom torchlibrosa.stft import Spectrogram, LogmelFilterBank\nfrom torchlibrosa.augmentation import SpecAugmentation\nfrom tqdm import tqdm\nimport soundfile as sf\nfrom contextlib import nullcontext","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Config\nWe'll start by setting up some global config variable that we'll access later."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# Global Vars\nNO_LABEL = -1\nNUM_CLASSES = 24\n\n\nclass config:\n    seed = 42\n    device = \"cuda:0\"\n    \n    train_tp_csv = '../input/rfcx-species-audio-detection/train_tp.csv'\n    test_csv = '../input/rfcx-species-audio-detection/sample_submission.csv'\n    save_path = './'\n    \n    encoder = resnet34d\n    encoder_features = 512\n    \n    percent_unlabeled = 1.0\n    consistency_weight = 100.0\n    consistency_rampup = 6\n    ema_decay = 0.995\n    positive_weight = 2.0\n    \n    lr = 1e-3\n    epochs = 25\n    batch_size = 32\n    num_workers = 4\n    train_5_folds = True\n    \n    period = 6 # 6 second clips\n    step = 1\n    model_params = {\n        'sample_rate': 48000,\n        'window_size': 2048,\n        'hop_size': 512,\n        'mel_bins': 384,\n        'fmin': 20,\n        'fmax': 48000 // 2,\n        'classes_num': NUM_CLASSES\n    }\n    \n    augmenter = A.Compose([\n        A.AddGaussianNoise(p=0.33, max_amplitude=0.02),\n        A.AddGaussianSNR(p=0.33),\n        A.FrequencyMask(min_frequency_band=0.01,  max_frequency_band=0.25, p=0.33),\n        A.TimeMask(min_band_part=0.01, max_band_part=0.25, p=0.33),\n        A.Gain(p=0.33)\n    ])\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"## Utils - Not much interesting going on here.\n\ndef get_n_fold_df(csv_path, folds=5):\n    df = pd.read_csv(csv_path)\n    df_group = df.groupby(\"recording_id\")[[\"species_id\"]].first().reset_index()\n    df_group = df_group.sample(frac=1, random_state=config.seed).reset_index(drop=True)\n    df_group.loc[:, 'fold'] = -1\n\n    X = df_group[\"recording_id\"].values\n    y = df_group[\"species_id\"].values\n\n    kfold = StratifiedKFold(n_splits=folds, random_state=config.seed)\n    for fold, (t_idx, v_idx) in enumerate(kfold.split(X, y)):\n        df_group.loc[v_idx, \"fold\"] = fold\n\n    return df.merge(df_group[['recording_id', 'fold']], on=\"recording_id\", how=\"left\")\n    \n\ndef 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\n\ndef sigmoid_rampup(current, rampup_length):\n    \"\"\"Exponential rampup from https://arxiv.org/abs/1610.02242\"\"\"\n    if rampup_length == 0:\n        return 1.0\n    else:\n        current = np.clip(current, 0.0, rampup_length)\n        phase = 1.0 - current / rampup_length\n        return float(np.exp(-5.0 * phase * phase))\n\n\nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\n\nclass MetricMeter(object):\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.y_true = []\n        self.y_pred = []\n\n    def update(self, y_true, y_pred):\n        try:\n            self.y_true.extend(y_true.detach().cpu().numpy().tolist())\n            self.y_pred.extend(torch.sigmoid(y_pred).cpu().detach().numpy().tolist())\n        except:\n            print(\"UPDATE FAILURE\")\n\n    def update_list(self, y_true, y_pred):\n        self.y_true.extend(y_true)\n        self.y_pred.extend(y_pred)\n\n    @property\n    def avg(self):\n        score_class, weight = lwlrap(np.array(self.y_true), np.array(self.y_pred))\n        self.score = (score_class * weight).sum()\n\n        return self.score\n    \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\ndef _one_sample_positive_class_precisions(scores, truth):\n    num_classes = scores.shape[0]\n    pos_class_indices = np.flatnonzero(truth > 0)\n\n    if not len(pos_class_indices):\n        return pos_class_indices, np.zeros(0)\n\n    retrieved_classes = np.argsort(scores)[::-1]\n\n    class_rankings = np.zeros(num_classes, dtype=np.int)\n    class_rankings[retrieved_classes] = range(num_classes)\n\n    retrieved_class_true = np.zeros(num_classes, dtype=np.bool)\n    retrieved_class_true[class_rankings[pos_class_indices]] = True\n\n    retrieved_cumulative_hits = np.cumsum(retrieved_class_true)\n\n    precision_at_hits = (\n            retrieved_cumulative_hits[class_rankings[pos_class_indices]] /\n            (1 + class_rankings[pos_class_indices].astype(np.float)))\n    return pos_class_indices, precision_at_hits\n\n\ndef lwlrap(truth, scores):\n    assert truth.shape == scores.shape\n    num_samples, num_classes = scores.shape\n    precisions_for_samples_by_classes = np.zeros((num_samples, num_classes))\n    for sample_num in range(num_samples):\n        pos_class_indices, precision_at_hits = _one_sample_positive_class_precisions(scores[sample_num, :],\n                                                                                     truth[sample_num, :])\n        precisions_for_samples_by_classes[sample_num, pos_class_indices] = precision_at_hits\n\n    labels_per_class = np.sum(truth > 0, axis=0)\n    weight_per_class = labels_per_class / float(np.sum(labels_per_class))\n\n    per_class_lwlrap = (np.sum(precisions_for_samples_by_classes, axis=0) /\n                        np.maximum(1, labels_per_class))\n    return per_class_lwlrap, weight_per_class\n\n\ndef pretty_print_metrics(fold, epoch, optimizer, train_loss_metrics, val_loss_metrics):\n    print(f\"\"\"\n    {time.ctime()} \\n\n    Fold:{fold}, Epoch:{epoch}, LR:{optimizer.param_groups[0]['lr']:.7}, Cons. Weight: {train_loss_metrics['consistency_weight']}\\n\n    --------------------------------------------------------\n    Metric:              Train    |   Val\n    --------------------------------------------------------\n    Loss:                {train_loss_metrics['loss']:0.4f}   |   {val_loss_metrics['loss']:0.4f}\\n\n    LWLRAP:              {train_loss_metrics['lwlrap']:0.4f}   |   {val_loss_metrics['lwlrap']:0.4f}\\n\n    Class Loss:          {train_loss_metrics['class_loss']:0.4f}   |   {val_loss_metrics['class_loss']:0.4f}\\n\n    Consistency Loss:    {train_loss_metrics['consistency_loss']:0.4f}   |   {val_loss_metrics['consistency_loss']:0.4f}\\n\n    --------------------------------------------------------\\n\n    \"\"\")\n    \n\nclass TestDataset(Dataset):\n    def __init__(self, df, data_path, period=10, step=1):\n        self.data_path = data_path\n        self.period = period\n        self.step = step\n        self.recording_ids = list(df[\"recording_id\"].unique())\n\n    def __len__(self):\n        return len(self.recording_ids)\n\n    def __getitem__(self, idx):\n        recording_id = self.recording_ids[idx]\n\n        y, sr = sf.read(f\"{self.data_path}/{recording_id}.flac\")\n\n        len_y = len(y)\n        effective_length = sr * self.period\n        effective_step = sr * self.step\n\n        y_ = []\n        i = 0\n        while i+effective_length <= len_y:\n            y__ = y[i:i + effective_length]\n\n            y_.append(y__)\n            i = i + effective_step\n\n        y = np.stack(y_)\n\n        label = np.zeros(NUM_CLASSES, dtype='f')\n\n        return {\n            \"waveform\": y,\n            \"target\": torch.tensor(label, dtype=torch.float),\n            \"id\": recording_id\n        }\n\n\ndef predict_on_test(model, test_loader):\n    model.eval()\n    pred_list = []\n    id_list = []\n    with torch.no_grad():\n        t = tqdm(test_loader)\n        for i, sample in enumerate(t):\n            input = sample[\"waveform\"].to(config.device)\n            bs, seq, w = input.shape\n            input = input.reshape(bs * seq, w)\n            id = sample[\"id\"]\n            output, _ = model(input)\n            output = output.reshape(bs, seq, -1)\n            output, _ = torch.max(output, dim=1)\n            \n            output = output.cpu().detach().numpy().tolist()\n            pred_list.extend(output)\n            id_list.extend(id)\n\n    return pred_list, id_list","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model\nThe model should look pretty familiar if you're using [SED](https://arxiv.org/abs/1912.04761). (Huge thanks to [Hidehisa Arai](https://www.kaggle.com/hidehisaarai1213) and their [SED Notebook](https://www.kaggle.com/hidehisaarai1213/introduction-to-sound-event-detection)!) You could use any model you'd like here. There's just one small tweak we need to make for our mean teacher setup. We need to \"detach\" the teacher's parameters so they aren't updated by the optimizer."},{"metadata":{"trusted":true},"cell_type":"code","source":"class 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\n    def __init__(self, sample_rate, window_size, hop_size, \n                 mel_bins, fmin, fmax, classes_num):\n        super().__init__()\n        self.interpolate_ratio = 32\n\n        self.spectrogram_extractor = Spectrogram(n_fft=window_size, \n                                                 hop_length=hop_size,\n                                                 win_length=window_size, \n                                                 window='hann', center=True,\n                                                 pad_mode='reflect', \n                                                 freeze_parameters=True)\n        self.logmel_extractor = LogmelFilterBank(sr=sample_rate, n_fft=window_size,\n                                                 n_mels=mel_bins, fmin=fmin, \n                                                 fmax=fmax, ref=1.0, \n                                                 amin=1e-10, top_db=None, \n                                                 freeze_parameters=True)\n\n        self.batch_norm = nn.BatchNorm2d(mel_bins)\n        self.encoder = partial(config.encoder, pretrained=True, in_chans=1)()\n        self.fc = nn.Linear(config.encoder_features, \n                            config.encoder_features, bias=True)\n        self.att_head = AttentionHead(config.encoder_features, classes_num)\n        self.avg_pool = nn.modules.pooling.AdaptiveAvgPool2d((1, 1))\n\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\n        x = self.encoder.forward_features(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 = F.dropout(x, p=0.5, training=self.training)\n        x = x.transpose(1, 2)\n        x = F.relu_(self.fc(x))\n        x = x.transpose(1, 2)\n        x = F.dropout(x, p=0.5, training=self.training)\n\n        (clipwise_output, norm_att, segmentwise_output) = self.att_head(x)\n        segmentwise_output = segmentwise_output.transpose(1, 2)\n\n        framewise_output = interpolate(segmentwise_output, self.interpolate_ratio)\n        return clipwise_output, framewise_output\n\n\ndef get_model(is_mean_teacher=False):\n    model = SEDAudioClassifier(**config.model_params)\n    model = model.to(config.device)\n    \n    # Detach params for Exponential Moving Average Model (aka the Mean Teacher).\n    # We'll manually update these params instead of using backprop.\n    if is_mean_teacher:\n        for param in model.parameters():\n            param.detach_()\n    return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Loss Function\nThe loss function has 2 components:\n\n1. A classification loss that only applies to labeled samples.\n2. A consistency loss that applies to all samples. \n\nFor the consistency loss we'll use the mean square error between the student and teacher predictions. We'll slowly ramp up the influence of the consistency loss since we don't want bad, early predictions having too much influence. \n\nNotice that we're weighting the positive samples for the classification loss. This is because we know the positives are correct while we're less sure about the negatives due to the missing labels issue. I found that this works better in practice. "},{"metadata":{"trusted":true},"cell_type":"code","source":"def sigmoid_mse_loss(input_logits, target_logits):\n    assert input_logits.size() == target_logits.size()\n    input_softmax = torch.sigmoid(input_logits)\n    target_softmax = torch.sigmoid(target_logits)\n    num_classes = input_logits.size()[1]\n    return F.mse_loss(input_softmax, target_softmax, size_average=False\n                     ) / num_classes\n\n\nclass MeanTeacherLoss(nn.Module):\n    \n    def __init__(self):\n        super().__init__()\n        self.positive_weight = torch.ones(\n            NUM_CLASSES).to(config.device) * config.positive_weight\n        self.class_criterion = nn.BCEWithLogitsLoss(\n            reduction='none', pos_weight=self.positive_weight)\n        self.consistency_criterion = sigmoid_mse_loss\n\n    def make_safe(self, pred):\n        pred = torch.where(torch.isnan(pred), torch.zeros_like(pred), pred)\n        return torch.where(torch.isinf(pred), torch.zeros_like(pred), pred)\n        \n    def get_consistency_weight(self, epoch):\n        # Consistency ramp-up from https://arxiv.org/abs/1610.02242\n        return config.consistency_weight * sigmoid_rampup(\n            epoch, config.consistency_rampup)\n    \n    def forward(self, student_pred, teacher_pred, target, classif_weights, epoch):\n        student_pred = self.make_safe(student_pred)\n        teacher_pred = self.make_safe(teacher_pred).detach().data\n\n        batch_size = len(target)\n        labeled_batch_size = target.ne(NO_LABEL).all(axis=1).sum().item() + 1e-3\n\n        student_classif, student_consistency = student_pred, student_pred\n        student_class_loss = (self.class_criterion(\n            student_classif, target) * classif_weights / labeled_batch_size).sum()\n\n        consistency_weights = self.get_consistency_weight(epoch)\n        consistency_loss = consistency_weights * self.consistency_criterion(\n            student_consistency, teacher_pred) / batch_size\n        loss = student_class_loss + consistency_loss\n        return loss, student_class_loss, consistency_loss, consistency_weights","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data Loader\nThe data loader produces two types of samples:\n\n1. Labeled samples with the audio centered in the clip.\n2. Random unlabeled clips without labels selected from files with at least one true positive label.\n\nEach sample contains 2 different inputs, one for the student and one for the teacher. Different augmentations are applied to each input."},{"metadata":{"trusted":true},"cell_type":"code","source":"class MeanTeacherDataset(Dataset):\n    \n    def __init__(self, df, transforms, period=5, \n                 data_path=\"../input/rfcx-species-audio-detection/train\", \n                 val=False, percent_unlabeled=0.0):\n        self.period = period\n        self.transforms = transforms\n        self.data_path = data_path\n        self.val = val\n        self.percent_unlabeled = percent_unlabeled\n\n        dfgby = df.groupby(\"recording_id\").agg(lambda x: list(x)).reset_index()\n        self.recording_ids = dfgby[\"recording_id\"].values\n        self.species_ids = dfgby[\"species_id\"].values\n        self.t_mins = dfgby[\"t_min\"].values\n        self.t_maxs = dfgby[\"t_max\"].values\n\n    def __len__(self):\n        return int(len(self.recording_ids) * (1 + self.percent_unlabeled))\n\n    def __getitem__(self, idx):\n        if idx >= len(self.recording_ids):\n            audio, label, rec_id, sr = self.get_unlabeled_item(idx)\n            # For unlabeled samples, we zero out the classification loss.\n            classif_weights = np.zeros(NUM_CLASSES, dtype='f')\n        else:\n            audio, label, rec_id, sr = self.get_labeled_item(idx)\n            classif_weights = np.ones(NUM_CLASSES, dtype='f')\n\n        audio_teacher = np.copy(audio)\n\n        # The 2 samples fed to the 2 models have should have different augmentations.\n        audio = self.transforms(samples=audio, sample_rate=sr)\n        audio_teacher = self.transforms(samples=audio_teacher, sample_rate=sr)\n        # assert (audio != audio_teacher).any()\n        \n        return {\n            \"waveform\": audio,\n            \"teacher_waveform\": audio_teacher,\n            \"target\": torch.tensor(label, dtype=torch.float),\n            \"classification_weights\": classif_weights,\n            \"id\": rec_id\n        }\n\n    def get_labeled_item(self, idx):\n        recording_id = self.recording_ids[idx]\n        species_id = self.species_ids[idx]\n        t_min, t_max = self.t_mins[idx], self.t_maxs[idx]\n\n        rec, sr = sf.read(f\"{self.data_path}/{recording_id}.flac\")\n\n        len_rec = len(rec)\n        effective_length = sr * self.period\n        rint = np.random.randint(len(t_min))\n        tmin, tmax = round(sr * t_min[rint]), round(sr * t_max[rint])\n        dur = tmax - tmin\n        min_dur = min(dur, round(sr * self.period))\n\n        center = round((tmin + tmax) / 2)\n        rand_start = center - effective_length + max(min_dur - dur//2, 0)\n        if rand_start < 0:\n            rand_start = 0\n        rand_end = center - max(min_dur - dur//2, 0)\n        start = np.random.randint(rand_start, rand_end)\n        rec = rec[start:start + effective_length]\n        if len(rec) < effective_length:\n            new_rec = np.zeros(effective_length, dtype=rec.dtype)\n            start1 = np.random.randint(effective_length - len(rec))\n            new_rec[start1:start1 + len(rec)] = rec\n            rec = new_rec.astype(np.float32)\n        else:\n            rec = rec.astype(np.float32)\n\n        start_time = start / sr\n        end_time = (start + effective_length) / sr\n\n        label = np.zeros(NUM_CLASSES, dtype='f')\n\n        for i in range(len(t_min)):\n            if (t_min[i] >= start_time) & (t_max[i] <= end_time):\n                label[species_id[i]] = 1\n            elif start_time <= ((t_min[i] + t_max[i]) / 2) <= end_time:\n                label[species_id[i]] = 1\n\n        return rec, label, recording_id, sr\n\n    def get_unlabeled_item(self, idx, random_sample=False):\n        real_idx = idx - len(self.recording_ids)\n        # We want our validation set to be fixed.\n        if self.val:\n            rec_id = self.recording_ids[real_idx]\n        else:\n            rec_id = random.sample(list(self.recording_ids), 1)[0]\n\n        rec, sr = sf.read(f\"{self.data_path}/{rec_id}.flac\")\n        effective_length = int(sr * self.period)\n        max_end = len(rec) - effective_length\n        if self.val:\n            # Fixed start for validation. Probaably a better way to do this.\n            start = int(idx * 16963 % max_end)\n        else:\n            start = np.random.randint(0, max_end)\n        rec = rec[start:(start+effective_length)]\n        rec = rec.astype(np.float32)\n\n        label = np.ones(NUM_CLASSES, dtype='f') * NO_LABEL\n\n        return rec, label, rec_id, sr\n\n    \ndef get_data_loader(df, is_val=False):\n    dataset = MeanTeacherDataset(\n        df=df,\n        transforms=config.augmenter,\n        period=config.period,\n        percent_unlabeled=config.percent_unlabeled\n    )\n    return torch.utils.data.DataLoader(\n        dataset,\n        batch_size=config.batch_size,\n        shuffle=not is_val,\n        drop_last=not is_val,\n        num_workers=config.num_workers\n    )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training\nAt the end of each training step we update the teacher weights by averaging in the latest student weights."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Update teacher to be exponential moving average of student params.\ndef update_teacher_params(student, teacher, alpha, global_step):\n    # Use the true average until the exponential average is more correct\n    alpha = min(1 - 1 / (global_step + 1), alpha)\n    for ema_param, param in zip(teacher.parameters(), student.parameters()):\n        ema_param.data.mul_(alpha).add_(1 - alpha, param.data)\n\n\ndef train_one_epoch(student, mean_teacher, loader, \n                    criterion, optimizer, scheduler, epoch, is_val=False):\n    global_step = 0\n    losses = AverageMeter()\n    consistency_loss_avg = AverageMeter()\n    class_loss_avg = AverageMeter()\n    comp_metric = MetricMeter()\n    \n    if is_val:\n        student.eval()\n        mean_teacher.eval()\n        context = torch.no_grad()\n    else:\n        student.train()\n        mean_teacher.train()\n        context = nullcontext()\n    \n    with context:\n        t = tqdm(loader)\n        for i, sample in enumerate(t):\n            student_input = sample['waveform'].to(config.device)\n            teacher_input = sample['teacher_waveform'].to(config.device)\n            target = sample['target'].to(config.device)\n            classif_weights = sample['classification_weights'].to(config.device)\n            batch_size = len(target)\n\n            student_pred, _  = student(student_input)\n            teacher_pred, _  = mean_teacher(teacher_input)\n\n            loss, class_loss, consistency_loss, consistency_weight = criterion(\n                student_pred, teacher_pred, target, classif_weights, epoch)\n\n            if not is_val:\n                optimizer.zero_grad()\n                loss.backward()\n                optimizer.step()\n                update_teacher_params(student, mean_teacher, \n                                      config.ema_decay, global_step)\n\n                scheduler.step()\n\n            comp_metric.update(target, student_pred)\n            losses.update(loss.item(), batch_size)\n            consistency_loss_avg.update(consistency_loss.item(), batch_size)\n            class_loss_avg.update(class_loss.item(), batch_size)\n            global_step += 1\n\n            t.set_description(f\"Epoch:{epoch} - Loss:{losses.avg:0.4f}\")\n        t.close()\n    return {'lwlrap':comp_metric.avg, \n            'loss':losses.avg, \n            'consistency_loss':consistency_loss_avg.avg, \n            'class_loss':class_loss_avg.avg, \n            'consistency_weight':consistency_weight}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Finally putting everything together..."},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"def train(df, fold):\n    train_df = df[df.fold != fold]\n    val_df = df[df.fold == fold]\n    train_loader = get_data_loader(train_df)\n    val_loader = get_data_loader(val_df)\n\n    student_model = get_model()\n    teacher_model = get_model(is_mean_teacher=True)\n\n    optimizer = torch.optim.AdamW(student_model.parameters(), lr=config.lr)\n    num_train_steps = int(len(train_loader) * config.epochs)\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n        optimizer, T_max=num_train_steps)\n    criterion = MeanTeacherLoss()\n\n    best_val_metric = -np.inf\n    val_metrics = []\n    train_metrics = []\n    for epoch in range(0, config.epochs):\n        train_loss_metrics = train_one_epoch(\n            student_model, teacher_model, train_loader, \n            criterion, optimizer, scheduler, epoch)\n        val_loss_metrics = train_one_epoch(\n            student_model, teacher_model, val_loader, \n            criterion, optimizer, scheduler, epoch, is_val=True)\n\n        train_metrics.append(train_loss_metrics)\n        val_metrics.append(val_loss_metrics)\n        pretty_print_metrics(fold, epoch, optimizer, \n                             train_loss_metrics, val_loss_metrics)\n        \n        if val_loss_metrics['lwlrap'] > best_val_metric:\n            print(f\"    LWLRAP Improved from {best_val_metric} --> {val_loss_metrics['lwlrap']}\\n\")\n            torch.save(teacher_model.state_dict(), \n                       os.path.join(config.save_path, f'fold-{fold}.bin'))\n            best_val_metric = val_loss_metrics['lwlrap']\n    \n\n\ndf = get_n_fold_df(config.train_tp_csv)\nfor fold in range(5 if config.train_5_folds else 1):\n    train(df, fold)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Predict on Test Set\nWe'll predict using the teacher model but you could also use the student or a combination of the two. Inference works just like it would for a vanilla baseline model."},{"metadata":{"trusted":true},"cell_type":"code","source":"def test(test_df, train_fold):\n    test_dataset = TestDataset(\n        df=test_df,\n        data_path=\"../input/rfcx-species-audio-detection/test\",\n        period=config.period,\n        step=config.step\n    )\n    test_loader = torch.utils.data.DataLoader(\n        test_dataset,\n        batch_size=1,\n        shuffle=False,\n        drop_last=False,\n        num_workers=config.num_workers\n    )\n    \n    weights_path = os.path.join(config.save_path, f'fold-{train_fold}.bin')\n    model = get_model()\n    model.load_state_dict(torch.load(weights_path, map_location=config.device), strict=False)\n    \n    test_pred, ids = predict_on_test(model, test_loader)\n\n    # Build Submission File\n    test_pred_df = pd.DataFrame({\n        \"recording_id\": test_df.recording_id.values\n    })\n    target_cols = test_df.columns[1:].values.tolist()\n    test_pred_df = test_pred_df.join(pd.DataFrame(np.array(test_pred), \n                                                  columns=target_cols))\n    test_pred_df.to_csv(os.path.join(config.save_path, \n                                     f\"fold-{train_fold}-submission.csv\"), \n                        index=False)\n    \n    \ntest_df = pd.read_csv(config.test_csv)\nfor fold in range(5 if config.train_5_folds else 1):\n    test(test_df, fold)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 5 Fold Ensemble\nFor 5 fold runs, we'll create a single ensemble prediction by simply averaging all of the folds."},{"metadata":{"trusted":true},"cell_type":"code","source":"def ensemble(submission_path):\n    dfs = [pd.read_csv(os.path.join(\n        config.save_path, f\"fold-{i}-submission.csv\")) for i in range(5)]\n    anchor = dfs[0].copy()\n    cols = anchor.columns[1:]\n   \n    for c in cols:\n        total = 0\n        for df in dfs:\n            total += df[c]\n        anchor[c] = total / len(dfs)\n    anchor.to_csv(submission_path, index=False)\n\n\nsubmission_path = os.path.join(config.save_path, f\"submission.csv\")\nif config.train_5_folds:\n    ensemble(submission_path)\nelse:\n    fold0_submission = os.path.join(config.save_path, f\"fold-0-submission.csv\")\n    os.rename(fold0_submission, submission_path)\n    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Conclusion \nThanks for reading! I dropped some unrelated tricks from this and didn't spend much time tuning so there's almost definetely room for improvement.\n\nI know it's pretty late in the competition for new notebooks, but considering that there are a few other public notebooks that score higher, I'm hoping this won't cause a significant shakeup. "}],"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}