{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":70203,"databundleVersionId":8068726,"sourceType":"competition"},{"sourceId":2132865,"sourceType":"datasetVersion","datasetId":1279829},{"sourceId":2132868,"sourceType":"datasetVersion","datasetId":1279831},{"sourceId":2132870,"sourceType":"datasetVersion","datasetId":1279833},{"sourceId":2132872,"sourceType":"datasetVersion","datasetId":1279834},{"sourceId":2273501,"sourceType":"datasetVersion","datasetId":1369048},{"sourceId":177386312,"sourceType":"kernelVersion"}],"dockerImageVersionId":30699,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Unsupervised Contrastive Learning Pipeline\n\nHey! This notebook is a pipeline to pretrain models with unsupervised contrastive learning. This is an implementation of the paper Contrastive Learning of General-Purpose Audio Representations ([COLA](https://arxiv.org/abs/2010.10915)). Most of the implementation is inspired from this [repo](https://github.com/CVxTz/COLA_pytorch). Idea of this notebook is: any amount of mel spectogram data can be added and trained given they are stored in .npy format.  \n\n** Please note, this is not trained on the full data. Please remove 'limit_train_batches' and 'limit_val_batches' in the trainer instance to train on full data. ","metadata":{}},{"cell_type":"code","source":"!pip install torchtoolbox -q\n!pip install timm -q","metadata":{"execution":{"iopub.status.busy":"2024-05-22T07:50:19.194550Z","iopub.execute_input":"2024-05-22T07:50:19.195458Z","iopub.status.idle":"2024-05-22T07:50:48.654850Z","shell.execute_reply.started":"2024-05-22T07:50:19.195420Z","shell.execute_reply":"2024-05-22T07:50:48.653756Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport torch\nimport torch.nn as nn\nimport os\nimport pytorch_lightning as pl\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn import model_selection\nimport torchvision.transforms as transforms\nimport librosa\nfrom PIL import Image\nimport albumentations as alb\nimport warnings\nimport sklearn.metrics\nimport timm \nfrom torchtoolbox.tools import mixup_data, mixup_criterion\nfrom torch.nn.functional import cross_entropy\nimport torchmetrics\nfrom sklearn.model_selection import train_test_split\nfrom pytorch_lightning.callbacks import  ModelCheckpoint, BackboneFinetuning, EarlyStopping \nfrom __future__ import print_function\nimport random\nimport librosa\nfrom torch.nn import functional as F\nimport torch.utils.data as data\nimport os\nimport glob\nimport numpy as np\nimport pandas as pd \nfrom collections import Counter\n\n\nwarnings.filterwarnings('ignore')","metadata":{"execution":{"iopub.status.busy":"2024-05-22T07:50:48.656812Z","iopub.execute_input":"2024-05-22T07:50:48.657111Z","iopub.status.idle":"2024-05-22T07:51:06.114616Z","shell.execute_reply.started":"2024-05-22T07:50:48.657081Z","shell.execute_reply":"2024-05-22T07:51:06.113857Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Config:\n    use_aug = False\n    num_classes = 182\n    batch_size = 64\n    epochs = 5\n    PRECISION = 16    \n    PATIENCE = 8  \n    seed = 2024\n    model = \"tf_efficientnet_b0_ns\"\n    pretrained = True            \n    weight_decay = 1e-3\n    use_mixup = False\n    mixup_alpha = 0.2   \n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')    \n    n_views = 2\n    \n    # SET THESE TO THE \n    data_root = \"/kaggle/input/birdclef-2024-numpy-spectrogram-split/\"\n    train_images = \"/kaggle/input/birdclef-2024-numpy-spectrogram-split/train_images/specs/train/\"\n    valid_images = \"/kaggle/input/birdclef-2024-numpy-spectrogram-split/train_images/specs/valid/\"\n    train_path = \"/kaggle/input/birdclef-2024-numpy-spectrogram-split/train_images/train.csv\"\n    valid_path = \"/kaggle/input/birdclef-2024-numpy-spectrogram-split/train_images/valid.csv\"\n    \n    \n    SR = 32000\n    DURATION = 5\n    MAX_READ_SAMPLES = 5\n    LR = 1e-4\n\n    \ndef config_to_dict(cfg):\n    return dict((name, getattr(cfg, name)) for name in dir(cfg) if not name.startswith('__'))","metadata":{"execution":{"iopub.status.busy":"2024-05-22T07:53:15.772097Z","iopub.execute_input":"2024-05-22T07:53:15.772457Z","iopub.status.idle":"2024-05-22T07:53:15.780173Z","shell.execute_reply.started":"2024-05-22T07:53:15.772428Z","shell.execute_reply":"2024-05-22T07:53:15.779310Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\ndef get_npy_files(directory):\n    # Use glob to find all .npy files in the specified directory\n    npy_files = glob.glob(os.path.join(directory, '**', '*.npy'), recursive=True)\n    return npy_files\n# Specify the directory path\ndirectory_path = '/kaggle/input'\n\n# Get all .npy files in the directory\nnpy_files = get_npy_files(directory_path)\n\nout_df = pd.DataFrame()\ndf_train = pd.DataFrame()\ndf_train['file_path'] = npy_files\ndf_train['primary_label'] = 'UNK'\ndf_train['data_source'] = df_train.file_path.apply(lambda x: x.split('/')[3])\n","metadata":{"execution":{"iopub.status.busy":"2024-05-22T07:53:15.922650Z","iopub.execute_input":"2024-05-22T07:53:15.923008Z","iopub.status.idle":"2024-05-22T07:53:22.329691Z","shell.execute_reply.started":"2024-05-22T07:53:15.922968Z","shell.execute_reply":"2024-05-22T07:53:22.328736Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del_list = []\nsource_shape_mapper = {}\n\nfor source in df_train['data_source'].unique():\n    \n    shape_li = []\n    # Create a subset of the dataframe for each unique value\n    subset = df_train[df_train['data_source'] == source]\n    # Randomly select one row from the subset\n    selected_row = subset.sample(n=10)\n    # Append the selected row to the list\n    \n    # Sanity - Shape check     \n    for x in selected_row.file_path:\n        test_im =  np.load(x)\n        shape_li.append(test_im.shape[1:])\n        if len(test_im.shape)!=3:\n            del_list.append(x.split('/')[3])\n            \n    counter = Counter(shape_li)\n\n    desired_shape = counter.most_common(1)[0][0]\n    \n    source_shape_mapper[source] = desired_shape\n","metadata":{"execution":{"iopub.status.busy":"2024-05-22T07:53:22.331877Z","iopub.execute_input":"2024-05-22T07:53:22.332227Z","iopub.status.idle":"2024-05-22T07:53:23.071404Z","shell.execute_reply.started":"2024-05-22T07:53:22.332194Z","shell.execute_reply":"2024-05-22T07:53:23.070515Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(f'del list: {del_list}')","metadata":{"execution":{"iopub.status.busy":"2024-05-22T07:53:23.072989Z","iopub.execute_input":"2024-05-22T07:53:23.073793Z","iopub.status.idle":"2024-05-22T07:53:23.078678Z","shell.execute_reply.started":"2024-05-22T07:53:23.073757Z","shell.execute_reply":"2024-05-22T07:53:23.077828Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train = df_train[~df_train['data_source'].isin(del_list)]","metadata":{"execution":{"iopub.status.busy":"2024-05-22T07:53:23.080603Z","iopub.execute_input":"2024-05-22T07:53:23.081320Z","iopub.status.idle":"2024-05-22T07:53:23.118523Z","shell.execute_reply.started":"2024-05-22T07:53:23.081283Z","shell.execute_reply":"2024-05-22T07:53:23.117586Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from sklearn.preprocessing import LabelEncoder\npl.seed_everything(Config.seed, workers=True)\n\n\nConfig.num_classes = len(df_train.primary_label.unique())\n# Assuming df_train is your DataFrame\nlabel_encoder = LabelEncoder()\ndf_train['primary_label'] = label_encoder.fit_transform(df_train['primary_label'])\n\n\n\nval_size = 0.05\n\ndf_train, df_valid = train_test_split(df_train, test_size=val_size, random_state=42)\n\nprint(\"Shape of training set:\", df_train.shape)\nprint(\"Shape of validation set:\", df_valid.shape)","metadata":{"execution":{"iopub.status.busy":"2024-05-22T07:53:23.119767Z","iopub.execute_input":"2024-05-22T07:53:23.120700Z","iopub.status.idle":"2024-05-22T07:53:23.188032Z","shell.execute_reply.started":"2024-05-22T07:53:23.120647Z","shell.execute_reply":"2024-05-22T07:53:23.186924Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import cv2\ndef resize_spectrogram(spectrogram, target_shape):\n    \n    if spectrogram.shape != target_shape:\n        return cv2.resize(spectrogram, target_shape, interpolation=cv2.INTER_LINEAR)\n    else: return spectrogram\n\ndef random_mask(data, rate_start=0.1, rate_seq=0.2):\n    new_data = data.copy()\n    mean = new_data.mean()\n    prev_zero = False\n    for i in range(new_data.shape[0]):\n        if random.random() < rate_start or (\n            prev_zero and random.random() < rate_seq\n        ):\n            prev_zero = True\n            new_data[i, :] = mean\n        else:\n            prev_zero = False\n\n    return new_data\n\n\ndef random_crop(data, crop_size=128):\n    start = int(random.random() * (data.shape[0] - crop_size))\n    return data[start : (start + crop_size), :]\n\n\ndef random_multiply(data):\n    new_data = data.copy()\n    return new_data * (0.9 + random.random() / 5.)","metadata":{"execution":{"iopub.status.busy":"2024-05-22T07:53:23.189198Z","iopub.execute_input":"2024-05-22T07:53:23.189844Z","iopub.status.idle":"2024-05-22T07:53:23.200578Z","shell.execute_reply.started":"2024-05-22T07:53:23.189815Z","shell.execute_reply":"2024-05-22T07:53:23.199700Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class AudioDataset(torch.utils.data.Dataset):\n    def __init__(self, data, max_len=100, augment=True, im_size = (128, 281)):\n        self.data = data\n        self.max_len = max_len\n        self.augment = augment\n        self.img_dir = Config.train_images\n        self.im_size = im_size\n\n    def __len__(self):\n        return len(self.data)\n\n    def __getitem__(self, idx):\n        row = self.data.iloc[idx]\n#         npy_path = self.img_dir + f\"{row.filename}.npy\"\n        npy_path = row.file_path\n        \n\n        x = np.load(npy_path)\n                \n        rn = np.random.choice(len(x))\n        x1 = x[rn]\n        x1 = resize_spectrogram(x1, self.im_size)\n\n        rn = np.random.choice(len(x))\n        x2 = x[rn]\n        x2 = resize_spectrogram(x2, self.im_size)\n        \n        x1 = np.expand_dims(x1, axis=0)\n        x2 = np.expand_dims(x2, axis=0)\n\n        if self.augment:\n            x1 = random_mask(x1)\n            x2 = random_mask(x2)\n\n\n        x1 = random_crop(x1, crop_size=self.max_len)\n        x2 = random_crop(x2, crop_size=self.max_len)\n\n\n        if self.augment:\n            x1 = random_multiply(x1)\n            x2 = random_multiply(x2)\n\n        x1 = torch.tensor(x1, dtype=torch.float)\n        x2 = torch.tensor(x2, dtype=torch.float)\n        \n        \n\n        return x1, x2","metadata":{"execution":{"iopub.status.busy":"2024-05-22T07:53:23.202093Z","iopub.execute_input":"2024-05-22T07:53:23.202430Z","iopub.status.idle":"2024-05-22T07:53:23.217242Z","shell.execute_reply.started":"2024-05-22T07:53:23.202398Z","shell.execute_reply":"2024-05-22T07:53:23.216339Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DecayLearningRate(pl.Callback):\n    def __init__(self):\n        self.old_lrs = []\n\n    def on_train_start(self, trainer, pl_module):\n        # track the initial learning rates\n        for opt_idx, optimizer in enumerate(trainer.optimizers):\n            group = []\n            for param_group in optimizer.param_groups:\n                group.append(param_group[\"lr\"])\n            self.old_lrs.append(group)\n\n    def on_train_epoch_end(self, trainer, pl_module):\n        for opt_idx, optimizer in enumerate(trainer.optimizers):\n            old_lr_group = self.old_lrs[opt_idx]\n            new_lr_group = []\n            for p_idx, param_group in enumerate(optimizer.param_groups):\n                old_lr = old_lr_group[p_idx]\n                new_lr = old_lr * 0.99\n                new_lr_group.append(new_lr)\n                param_group[\"lr\"] = new_lr\n            self.old_lrs[opt_idx] = new_lr_group","metadata":{"execution":{"iopub.status.busy":"2024-05-22T07:53:23.218762Z","iopub.execute_input":"2024-05-22T07:53:23.219038Z","iopub.status.idle":"2024-05-22T07:53:23.235846Z","shell.execute_reply.started":"2024-05-22T07:53:23.219013Z","shell.execute_reply":"2024-05-22T07:53:23.234973Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Encoder(pl.LightningModule):\n    def __init__(self, model_name=Config.model, num_classes=Config.num_classes, \n                 pretrained=Config.pretrained):\n        super().__init__()\n        self.num_classes = num_classes\n\n        self.backbone = timm.create_model(model_name, pretrained=pretrained)\n\n    def forward(self, images):\n        # If the input images are grayscale, add a singleton dimension for the channel\n        if images.ndim == 3:\n            images = images.unsqueeze(1).repeat(1, 3, 1, 1)\n        logits = self.backbone(images)\n        return logits\n    ","metadata":{"execution":{"iopub.status.busy":"2024-05-22T07:53:23.237180Z","iopub.execute_input":"2024-05-22T07:53:23.237802Z","iopub.status.idle":"2024-05-22T07:53:23.254972Z","shell.execute_reply.started":"2024-05-22T07:53:23.237772Z","shell.execute_reply":"2024-05-22T07:53:23.254035Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Cola(pl.LightningModule):\n    def __init__(self, p=0.1):\n        super().__init__()\n        self.save_hyperparameters()\n\n        self.p = p\n\n        self.do = torch.nn.Dropout(p=self.p)\n\n#         self.encoder = Encoder(drop_connect_rate=p)\n        self.encoder = Encoder()\n\n        self.g = torch.nn.Linear(1000, 512)\n        self.layer_norm = torch.nn.LayerNorm(normalized_shape=512)\n        self.linear = torch.nn.Linear(512, 512, bias=False)\n        \n        self.validation_step_outputs = []\n\n    def forward(self, x):\n        x1, x2 = x\n\n        x1 = self.do(self.encoder(x1.repeat(1, 3, 1, 1)))\n        x1 = self.do(self.g(x1))\n        x1 = self.do(torch.tanh(self.layer_norm(x1)))\n\n        x2 = self.do(self.encoder(x2.repeat(1, 3, 1, 1)))\n        x2 = self.do(self.g(x2))\n        x2 = self.do(torch.tanh(self.layer_norm(x2)))\n\n        x1 = self.linear(x1)\n\n        return x1, x2\n\n    def training_step(self, x, batch_idx):\n        x1, x2 = self(x)\n\n        y = torch.arange(x1.size(0), device=x1.device)\n\n        y_hat = torch.mm(x1, x2.t())\n\n        loss = F.cross_entropy(y_hat, y)\n\n        _, predicted = torch.max(y_hat, 1)\n        acc = (predicted == y).double().mean()\n\n        self.log(\"train_loss\", loss, on_step=True, on_epoch=True, prog_bar=True)\n        self.log(\"train_acc\", acc, on_step=True, on_epoch=True, prog_bar=True)\n        \n        return loss\n\n    def validation_step(self, x, batch_idx):\n        x1, x2 = self(x)\n\n        y = torch.arange(x1.size(0), device=x1.device)\n\n        y_hat = torch.mm(x1, x2.t())\n\n        loss = F.cross_entropy(y_hat, y)\n\n        _, predicted = torch.max(y_hat, 1)\n        acc = (predicted == y).double().mean()\n        \n        loss_array = {\"val_loss\": loss}\n        \n        self.validation_step_outputs.append(loss_array)\n\n        self.log(\"valid_loss\", loss, on_step=True, on_epoch=True, prog_bar=True)\n        self.log(\"valid_acc\", acc, on_step=True, on_epoch=True, prog_bar=True)\n\n    def test_step(self, x, batch_idx):\n        x1, x2 = self(x)\n\n        y = torch.arange(x1.size(0), device=x1.device)\n\n        y_hat = torch.mm(x1, x2.t())\n\n        loss = F.cross_entropy(y_hat, y)\n\n        _, predicted = torch.max(y_hat, 1)\n        acc = (predicted == y).double().mean()\n\n        self.log(\"test_loss\", loss)\n        self.log(\"test_acc\", acc)\n\n    def configure_optimizers(self):\n        return torch.optim.Adam(self.parameters(), lr=1e-4)\n    \n    def on_validation_epoch_end(self):\n        avg_loss = torch.stack([x['val_loss'] for x in self.validation_step_outputs]).mean()\n        \n        print(f'Epoch {self.current_epoch} Validation Loss: {avg_loss}')\n        \n        \n        self.save_encoder_weights(f'efficientnet_backbone_weights_{self.current_epoch}.pth')\n        self.save_full_weights(f'cola_weights_{self.current_epoch}.pth')\n        return {'val_loss': avg_loss}\n    \n    def save_encoder_weights(self, path):\n        torch.save(self.encoder.state_dict(), path)\n    \n    def save_full_weights(self, path):\n        torch.save(self.state_dict(), path)\n#     def load_encoder_weights(self, path):\n#         self.encoder.load_state_dict(torch.load(path))\n","metadata":{"execution":{"iopub.status.busy":"2024-05-22T07:53:23.258292Z","iopub.execute_input":"2024-05-22T07:53:23.258563Z","iopub.status.idle":"2024-05-22T07:53:23.284311Z","shell.execute_reply.started":"2024-05-22T07:53:23.258540Z","shell.execute_reply":"2024-05-22T07:53:23.283167Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = Cola()","metadata":{"execution":{"iopub.status.busy":"2024-05-22T07:53:23.285927Z","iopub.execute_input":"2024-05-22T07:53:23.286304Z","iopub.status.idle":"2024-05-22T07:53:23.692807Z","shell.execute_reply.started":"2024-05-22T07:53:23.286273Z","shell.execute_reply":"2024-05-22T07:53:23.691835Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_source_dict = {}\n\n# Iterate over unique values in the 'data_source' column\nfor source in df_train['data_source'].unique():\n    # Create a subset of the dataframe for each unique value\n    subset = df_train[df_train['data_source'] == source]\n    # Store the subset in the dictionary\n    data_source_dict[source] = subset\n\n# Print the dictionary keys and corresponding dataframes (optional)\nfor key, value in data_source_dict.items():\n    print(f'Processing Data: {key}')\n    \n    train, val = train_test_split(value, test_size=0.05, random_state=1337)\n\n    train_data = AudioDataset(train, augment=True, im_size = source_shape_mapper[key])\n    val_data = AudioDataset(val, augment=False, im_size = source_shape_mapper[key])\n\n\n    train_loader = DataLoader(\n        train_data, batch_size=Config.batch_size, num_workers=8, shuffle=True\n    )\n    val_loader = DataLoader(\n        val_data, batch_size=Config.batch_size, num_workers=8, shuffle=True\n    )\n    checkpoint_callback = ModelCheckpoint(\n        monitor=\"valid_acc\", mode=\"max\", \n        dirpath=\"models/\", \n        filename= f\"encoder-latest\"\n    )\n\n\n    trainer = pl.Trainer(\n        max_epochs=Config.epochs,\n    #     gpus=1,\n        logger=None,\n        callbacks=[checkpoint_callback, DecayLearningRate()],\n#         limit_train_batches = 0.01, # Remove for whole run\n#         limit_val_batches= 0.1\n    )\n\n    trainer.fit(model, train_loader, val_loader)","metadata":{"execution":{"iopub.status.busy":"2024-05-22T07:53:35.079037Z","iopub.execute_input":"2024-05-22T07:53:35.079445Z","iopub.status.idle":"2024-05-22T07:54:22.879290Z","shell.execute_reply.started":"2024-05-22T07:53:35.079409Z","shell.execute_reply":"2024-05-22T07:54:22.878161Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}