{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Description\n\n* This notebook demonstrates how a few-shot learning on BirdCLEF 2022 dataset is going.\n* For metric learning, I used Prototypical network[1] since it is one of basic approaches on few-shot learning.\n* Most of the code in this notebook was adapted from [2].\n\n## What this notebook supports\n\n- learning classifier without overfitting to train dataset even if samples are few (few-shot learning)\n    - classifying unseen 5 classes by feeding only 5 samples per each classes\n\n## What this notebook doesn't support\n\n- call/no call classifier\n    - since the models are trained only by positive (bird call) samples, it can't distinguish background from bird call (; maybe you need another classifier, or you need to input background samples).\n- multi-label classification\n    - it only classifies most probable 1 class with inputting 5 second audio frame\n    - `secondary_labels` are completely ignored\n\n# Reference\n\n* [1] https://arxiv.org/abs/1703.05175\n* [2] https://github.com/Frankluox/LightningFSL","metadata":{}},{"cell_type":"markdown","source":"# Environment Setup","metadata":{}},{"cell_type":"code","source":"!pip install nb_black > /dev/null\n!pip install torchinfo > /dev/null","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!nvidia-smi","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip freeze | grep torch","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip freeze | grep librosa","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from collections import Counter\nimport math\nimport os\n\nfrom argparse import Namespace\nfrom typing import List, Union, TypeVar, Iterator, Optional, Dict, Tuple, Any\n\nfrom IPython.core.debugger import set_trace, Pdb\nimport librosa\nimport librosa.display\nimport numpy as np\nimport pandas as pd\nimport pytorch_lightning as pl\nimport torch\nimport torch.nn.functional as F\nimport torchaudio\nimport torchaudio.transforms as T\nimport matplotlib.pyplot as plt\nimport wandb\n\nfrom kaggle_secrets import UserSecretsClient\nfrom pytorch_lightning import LightningDataModule, LightningModule\nfrom pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint\nfrom pytorch_lightning.loggers import WandbLogger\nfrom pytorch_lightning.utilities.cli import (\n    LightningCLI,\n    LightningArgumentParser,\n    SaveConfigCallback,\n)\nfrom pytorch_lightning.utilities.seed import seed_everything\n\nfrom pytorch_lightning.trainer.trainer import Trainer\nfrom sklearn.model_selection import train_test_split\nfrom tqdm.notebook import tqdm\nfrom torch import nn, Tensor\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\nfrom torch.utils.data import Dataset, DataLoader, Sampler\n\nfrom torchmetrics import Accuracy, MeanMetric, Metric\nfrom torchinfo import summary\n\nplt.style.use(\"ggplot\")\n\n%load_ext lab_black\n%load_ext autoreload\n%autoreload 2","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-04-19T01:55:00.205832Z","iopub.execute_input":"2022-04-19T01:55:00.206574Z","iopub.status.idle":"2022-04-19T01:55:04.935402Z","shell.execute_reply.started":"2022-04-19T01:55:00.206452Z","shell.execute_reply":"2022-04-19T01:55:04.934443Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Configuration","metadata":{}},{"cell_type":"code","source":"cfg = dict(\n    sample_rate=32_000,\n    hop_length=500,\n    n_fft=2_000,\n    fmin=0,\n    fmax=16_000,\n    wav_crop_sec=7,\n    random_seed=3154,\n    max_epochs=10,\n)\nmetadata_path = (\n    \"../input/birdclef-2022-precomputed-melspec-hop-size500/spec_metadata.csv\"\n)\nimg_path = \"../input/birdclef-2022-precomputed-melspec-hop-size500/train_audio\"\ntrain_transforms = [\n    # T.TimeMasking(time_mask_param=80),\n    # T.FrequencyMasking(freq_mask_param=32),\n    # T.PitchShift(\n    #    sample_rate=cfg[\"sample_rate\"],\n    #    n_steps=5,\n    #    n_fft=cfg[\"n_fft\"],\n    #    hop_length=cfg[\"hop_length\"],\n    # ),\n]\nval_transforms = []\ntrain_transform = nn.Sequential(*train_transforms)\nval_transform = nn.Sequential(*val_transforms)","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:55:38.813624Z","iopub.execute_input":"2022-04-19T01:55:38.814178Z","iopub.status.idle":"2022-04-19T01:55:38.886000Z","shell.execute_reply.started":"2022-04-19T01:55:38.814133Z","shell.execute_reply":"2022-04-19T01:55:38.885138Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# W&B login","metadata":{}},{"cell_type":"code","source":"user_secrets = UserSecretsClient()\n\npersonal_key_for_api = user_secrets.get_secret(\"ke\")\n! wandb login $personal_key_for_api","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:55:39.576748Z","iopub.execute_input":"2022-04-19T01:55:39.577016Z","iopub.status.idle":"2022-04-19T01:55:42.046789Z","shell.execute_reply.started":"2022-04-19T01:55:39.576984Z","shell.execute_reply":"2022-04-19T01:55:42.045227Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Split data","metadata":{}},{"cell_type":"code","source":"trainval = pd.read_csv(\n    \"../input/birdclef-2022-precomputed-melspec-hop-size500/spec_metadata.csv\"\n)","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:55:42.049654Z","iopub.execute_input":"2022-04-19T01:55:42.050197Z","iopub.status.idle":"2022-04-19T01:55:42.264893Z","shell.execute_reply.started":"2022-04-19T01:55:42.050153Z","shell.execute_reply":"2022-04-19T01:55:42.264124Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainval_species = (\n    trainval.groupby(\"primary_label\")\n    .agg(n_samples=(\"length\", \"count\"))\n    .query(\"n_samples >= 20\")\n)\ntrainval_species = set(trainval_species.index)\nprint(f\"* {len(trainval_species)} species which have sample >= 20\")","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:55:42.266254Z","iopub.execute_input":"2022-04-19T01:55:42.266676Z","iopub.status.idle":"2022-04-19T01:55:42.415324Z","shell.execute_reply.started":"2022-04-19T01:55:42.266638Z","shell.execute_reply":"2022-04-19T01:55:42.414336Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_val_split(trainval_species, train_ratio=0.8):\n    train_species = set(\n        np.random.choice(\n            list(trainval_species),\n            int(len(trainval_species) * train_ratio),\n            replace=False,\n        )\n    )\n    val_species = trainval_species - train_species\n    return train_species, val_species","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:55:42.417998Z","iopub.execute_input":"2022-04-19T01:55:42.420113Z","iopub.status.idle":"2022-04-19T01:55:42.521640Z","shell.execute_reply.started":"2022-04-19T01:55:42.420074Z","shell.execute_reply":"2022-04-19T01:55:42.520066Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seed_everything(cfg[\"random_seed\"])\ntrain_species, val_species = train_val_split(trainval_species)\nprint(len(train_species), len(val_species))\nprint(list(train_species)[:3], list(val_species)[:3])","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:55:42.524419Z","iopub.execute_input":"2022-04-19T01:55:42.524660Z","iopub.status.idle":"2022-04-19T01:55:42.701688Z","shell.execute_reply.started":"2022-04-19T01:55:42.524631Z","shell.execute_reply":"2022-04-19T01:55:42.700142Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train = trainval.query(\"primary_label in @train_species\")\nval = trainval.query(\"primary_label in @val_species\")\nlen(train) + len(val), len(trainval)","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:55:42.716251Z","iopub.execute_input":"2022-04-19T01:55:42.718948Z","iopub.status.idle":"2022-04-19T01:55:42.946004Z","shell.execute_reply.started":"2022-04-19T01:55:42.718901Z","shell.execute_reply":"2022-04-19T01:55:42.945250Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"class BirdClefDataset(Dataset):\n    def __init__(\n        self,\n        metadata_df,\n        img_dir,\n        cfg,\n        transform=None,\n        epsilon=1e-12,\n    ):\n        self.metadata_df = metadata_df\n        self.img_dir = img_dir\n        self.transform = transform\n        self.cfg = cfg\n        label = self.metadata_df.primary_label.to_numpy()\n        self.species = np.unique(label)\n        self.species2index = {s: i for i, s in enumerate(self.species)}\n        self.label = np.array(list(map(lambda x: self.species2index[x], label)))\n        self.epsilon = 1e-12\n\n    def __len__(self):\n        return len(self.metadata_df)\n\n    def __getitem__(self, idx):\n        row = self.metadata_df.iloc[idx]\n        cfg = self.cfg\n\n        data_path = os.path.join(self.img_dir, row[\"filename\"])\n        spec = np.load(data_path).astype(np.float32)\n\n        # calc_offset\n        wav_len = spec.shape[1]\n        duration = int(\n            cfg[\"wav_crop_sec\"] / cfg[\"hop_length\"] * cfg[\"sample_rate\"]\n        )  # 30 / 500 * 32_000 = 1920\n        max_offset = max(wav_len - duration, 0)\n        offset = np.random.randint(max_offset + 1)\n\n        # print(f\"========== {row['filename']} ==========\")\n        # print(f\"offset: {offset}\")\n        # print(f\"wav_len: {wav_len}\")\n        # print(f\"duration: {duration}\")\n        spec = spec[:, offset:]\n\n        # crop spectrogram in time dimension\n        if wav_len < duration:\n            pad = duration - wav_len\n            if pad >= wav_len:  # repeat the same timeframe\n                n_repeat = duration // wav_len + 1\n                spec = np.tile(spec, [1, n_repeat])\n            else:\n                spec_orig = (\n                    spec.copy()\n                )  # concat original timeframe with randomly cropped timeframe\n                max_offs = wav_len - pad\n                offs = np.random.randint(max_offs + 1)\n                spec = np.concatenate([spec, spec_orig[:, offs : offs + pad]], axis=1)\n        spec = spec[:, :duration]\n        # print(f\"spec.shape: {spec.shape}\")\n\n        # channel normalization\n        mean, std = spec.mean(), spec.std()\n        spec = (spec - mean) / max(std, self.epsilon)\n\n        # (F, T) -> (1, F, T)\n        spec = np.expand_dims(spec, 0)\n\n        # set label\n        label = row[\"primary_label\"]\n        # label = self.species2index[label]\n\n        # numpy to tensor\n        spec = torch.from_numpy(spec).float()\n\n        # transform\n        if self.transform:\n            spec = self.transform(spec)\n        return spec, label","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:55:44.411354Z","iopub.execute_input":"2022-04-19T01:55:44.411649Z","iopub.status.idle":"2022-04-19T01:55:44.524450Z","shell.execute_reply.started":"2022-04-19T01:55:44.411614Z","shell.execute_reply":"2022-04-19T01:55:44.523716Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data Module","metadata":{}},{"cell_type":"code","source":"T_co = TypeVar(\"T_co\", covariant=True)\n\n\nclass CategoriesSampler(Sampler[T_co]):\n    r\"\"\"Sampler that collects data into several few-shot learning tasks\"\"\"\n\n    def __init__(\n        self,\n        labels: Union[List, \"np.ndarray\"],\n        num_task: int,\n        way: int,\n        total_sample_per_class: int,\n        total_batch_size: int = 4,\n        drop_last: bool = False,\n    ) -> None:\n        \"\"\"\n        Args:\n           labels: The corresponding labels of the whole dataset .\n           num_task: The number of tasks within one epoch.\n           way: The number of classes within one task.\n           total_sample_per_class: The number of samples within each few-shot class(all samples from support and query).\n           total_batch_size: The number of tasks to handle per iteration.\n           drop_last (bool, optional): if ``True``, then the sampler will drop the\n               tail of the data to make it evenly divisible across the number of\n               replicas. If ``False``, the sampler will add extra indices to make\n               the data evenly divisible across the replicas. Default: ``False``.\n        \"\"\"\n        self.num_task = num_task\n        self.way = way\n        self.total_sample_per_class = total_sample_per_class\n        self.drop_last = drop_last\n        self.total_batch_size = total_batch_size\n        self.per_gpu_batch_size = self.total_batch_size\n        self.m_ind = None\n\n        if self.drop_last:\n            self.num_iteration = math.floor(self.num_task / self.total_batch_size)\n        else:\n            self.num_iteration = math.ceil(self.num_task / self.total_batch_size)\n\n        labels = np.array(labels)  # all data labels\n        self.m_ind = {}  # the data index of each class\n        classes = np.unique(labels)\n        for c in classes:\n            ind = np.argwhere(labels == c).reshape(-1)  # all data index of this class\n            # ind = torch.from_numpy(ind)\n            self.m_ind[c] = ind\n\n    def __len__(self) -> int:\n        return self.num_iteration\n\n    def __iter__(self) -> Iterator[T_co]:\n        # print(self.num_iteration)\n        for _ in range(self.num_iteration):\n            tasks = []\n            for i in range(self.per_gpu_batch_size):\n                task = []\n                # random sample num_class indexs,e.g. 5\n                classes = torch.randperm(len(self.m_ind))[: self.way].numpy()\n                # print(f\"{j}: {i}: {dist.get_rank()}: {classes}\")\n                # print(classes)\n                for c in classes:\n                    # sample total_sample_per_class data index of this class\n                    l = self.m_ind[c]  # all data indexs of this class\n                    pos = torch.randperm(len(l))[: self.total_sample_per_class]\n                    task.append(l[pos])\n                tasks.append(np.stack(task).transpose().reshape(-1))\n            tasks = np.stack(tasks).reshape(-1)\n            yield tasks","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:55:45.247433Z","iopub.execute_input":"2022-04-19T01:55:45.247853Z","iopub.status.idle":"2022-04-19T01:55:45.342333Z","shell.execute_reply.started":"2022-04-19T01:55:45.247815Z","shell.execute_reply":"2022-04-19T01:55:45.341516Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class FewShotDataModule(LightningDataModule):\n    def __init__(\n        self,\n        train,\n        val,\n        img_path,\n        cfg,\n        train_transform,\n        val_transform,\n        way: int = 5,\n        num_query: int = 15,\n        drop_last: Optional[bool] = None,\n        num_gpus: int = 1,\n        train_batch_size: int = 1,\n        val_batch_size: int = 4,\n        train_num_workers: int = 2,\n        val_num_workers: int = 2,\n        train_num_task_per_epoch: Optional[int] = 1000,\n        val_num_task: int = 200,\n        train_shot: Optional[int] = 5,\n        val_shot: int = 5,\n        train_dataset_params: Dict = {},\n        val_test_dataset_params: Dict = {},\n    ) -> None:\n        super().__init__()\n        self.train = train\n        self.val = val\n        self.img_path = img_path\n        self.cfg = cfg\n        self.train_transform = train_transform\n        self.val_transform = val_transform\n\n        self.way = way\n        self.num_query = num_query\n        self.drop_last = drop_last\n        self.num_gpus = num_gpus\n\n        self.train_num_workers = train_num_workers\n        self.train_batch_size = train_batch_size\n        self.train_batch_sampler = None\n        self.train_num_task_per_epoch = train_num_task_per_epoch\n        self.train_shot = train_shot\n        self.train_dataset_params = train_dataset_params\n\n        self.val_num_workers = val_num_workers\n        self.val_batch_size = val_batch_size\n        self.val_batch_sampler = None\n        self.val_num_task = val_num_task\n        self.val_shot = val_shot\n        self.val_test_dataset_params = val_test_dataset_params\n\n    def set_train_dataset(self):\n        self.train_dataset = BirdClefDataset(\n            self.train, self.img_path, self.cfg, transform=self.train_transform\n        )\n\n    def set_val_dataset(self):\n        self.val_dataset = BirdClefDataset(\n            self.val, self.img_path, self.cfg, transform=self.val_transform\n        )\n\n    def set_sampler(self):\n        self.train_batch_sampler = CategoriesSampler(\n            self.train_dataset.label,\n            self.train_num_task_per_epoch,\n            self.way,\n            self.train_shot + self.num_query,\n            self.train_batch_size,\n            self.drop_last,\n        )\n        self.val_batch_sampler = CategoriesSampler(\n            self.val_dataset.label,\n            self.val_num_task,\n            self.way,\n            self.val_shot + self.num_query,\n            self.val_batch_size,\n            self.drop_last,\n        )\n\n    def setup(self, stage=None):\n        self.set_train_dataset()\n        self.set_val_dataset()\n        self.set_sampler()\n\n    def train_dataloader(self):\n        loader = DataLoader(\n            self.train_dataset,\n            shuffle=False,\n            num_workers=self.train_num_workers,\n            batch_sampler=self.train_batch_sampler,\n            pin_memory=True,\n        )\n        return loader\n\n    def val_dataloader(self):\n        loader = DataLoader(\n            self.val_dataset,\n            shuffle=False,\n            num_workers=self.val_num_workers,\n            batch_sampler=self.val_batch_sampler,\n            pin_memory=True,\n        )\n        return loader","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:55:45.746953Z","iopub.execute_input":"2022-04-19T01:55:45.747206Z","iopub.status.idle":"2022-04-19T01:55:45.859148Z","shell.execute_reply.started":"2022-04-19T01:55:45.747176Z","shell.execute_reply":"2022-04-19T01:55:45.858393Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"datamodule = FewShotDataModule(\n    train, val, img_path, cfg, train_transform, val_transform\n)\ndatamodule.setup()\nfor batch, label in datamodule.train_dataloader():\n    print(batch.shape)\n    break","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:55:46.203381Z","iopub.execute_input":"2022-04-19T01:55:46.205491Z","iopub.status.idle":"2022-04-19T01:55:47.203564Z","shell.execute_reply.started":"2022-04-19T01:55:46.205434Z","shell.execute_reply":"2022-04-19T01:55:47.202583Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"markdown","source":"## Backbone","metadata":{}},{"cell_type":"code","source":"def floor_power(num, divisor, power):\n    \"\"\"Performs what we call a floor power, a recursive fixed division process\n        with a flooring between each time\n\n    Args:\n        num (int or float):The original number to divide from\n        divisor (int or float): The actual divisor for the number\n        power (int): How many times we apply this divide and then floor\n\n    Returns:\n        int: The numerical result of the floor division process\n    \"\"\"\n    for _ in range(power):\n        num = np.floor(num / divisor)\n    return num","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:55:47.600559Z","iopub.execute_input":"2022-04-19T01:55:47.601379Z","iopub.status.idle":"2022-04-19T01:55:47.675308Z","shell.execute_reply.started":"2022-04-19T01:55:47.601328Z","shell.execute_reply":"2022-04-19T01:55:47.674260Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def gem(x, p=3, eps=1e-6):\n    return F.avg_pool2d(x.clamp(min=eps).pow(p), (x.size(-2), x.size(-1))).pow(1.0 / p)\n\n\nclass GeMPooling(nn.Module):\n    def __init__(self, p=3, eps=1e-6):\n        super().__init__()\n        self.p = p\n        self.eps = eps\n\n    def forward(self, x):\n        x = gem(x, p=self.p, eps=self.eps)\n        return x","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:55:48.737614Z","iopub.execute_input":"2022-04-19T01:55:48.737875Z","iopub.status.idle":"2022-04-19T01:55:48.805081Z","shell.execute_reply.started":"2022-04-19T01:55:48.737846Z","shell.execute_reply":"2022-04-19T01:55:48.804171Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def conv3x3(in_planes, out_planes):\n    return nn.Conv2d(in_planes, out_planes, 3, padding=1, bias=False)\n\n\ndef norm_layer(planes):\n    return nn.BatchNorm2d(planes)\n\n\nclass Block(nn.Module):\n    def __init__(self, inplanes, outplanes, pool_dim):\n        super().__init__()\n        self.relu = nn.LeakyReLU(0.1)\n        self.conv = conv3x3(inplanes, outplanes)\n        self.bn = norm_layer(outplanes)\n        self.maxpool = nn.MaxPool2d(pool_dim)\n\n    def forward(self, x):\n        out = self.conv(x)\n        out = self.bn(out)\n        out = self.relu(out)\n        out = self.maxpool(out)\n        return out\n\n\nclass ConvN(nn.Module):\n    def __init__(self, channels=[1, 64, 64, 64, 64], pool_dim=(3, 3)):\n        super().__init__()\n\n        self.nn = nn.Sequential(\n            *[\n                Block(ci, co, pool_dim)\n                for i, (ci, co) in enumerate(zip(channels[:-1], channels[1:]))\n            ]\n        )\n\n    def forward(self, x):\n        out = self.nn(x)\n        return out\n\n\nclass ConvNFeatureExtractor(nn.Module):\n    def __init__(\n        self,\n        trial_shape=[-1, -1, 128, 448],\n        channels=[1, 64, 64, 128, 128],\n        pool_dim=(2, 3),\n        out_dim=128,\n        pool_type=\"flatten\",\n    ):\n        super().__init__()\n        self.pool_type = pool_type\n\n        self.conv_encoder = ConvN(channels=channels, pool_dim=pool_dim)\n        num_convs = len(channels) - 1\n        num_logits = int(\n            channels[-1]\n            * floor_power(trial_shape[2], pool_dim[0], num_convs)\n            * floor_power(trial_shape[3], pool_dim[1], num_convs)\n        )\n        if self.pool_type == \"dense\":\n            self.pool = nn.Sequential(\n                nn.Dropout(p=0.3),\n                nn.BatchNorm1d(num_logits, eps=1e-05, momentum=0.1, affine=True),\n                nn.Linear(in_features=num_logits, out_features=out_dim),\n            )\n        elif self.pool_type == \"gem\":\n            self.pool = GeMPooling(p=3)\n\n    def forward(self, x):\n        x = self.conv_encoder(x)\n        if self.pool_type == \"flatten\":\n            x = x.view(x.size(0), -1)\n        elif self.pool_type == \"dense\":\n            x = x.view(x.size(0), -1)\n            x = self.pool(x)\n        elif self.pool_type == \"adaptive_average\":\n            x = F.adaptive_avg_pool2d(x, 1).view(x.size(0), -1)\n        elif self.pool_type == \"normed_adaptive_average\":\n            F.normalize(x, p=2, dim=1, eps=1e-12)\n            x = F.adaptive_avg_pool2d(x, 1).view(x.size(0), -1)\n        elif self.pool_type == \"gem\":\n            x = self.pool(x).view(x.size(0), -1)\n        else:\n            raise NotImplementedError\n        assert x.dim() == 2, x.dim()\n\n        return x","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:57:37.116289Z","iopub.execute_input":"2022-04-19T01:57:37.116762Z","iopub.status.idle":"2022-04-19T01:57:37.246624Z","shell.execute_reply.started":"2022-04-19T01:57:37.116718Z","shell.execute_reply":"2022-04-19T01:57:37.245895Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"batch_size = 2\nshot = 5\nway = 5\nn_query = 15\nduration = int(cfg[\"wav_crop_sec\"] / cfg[\"hop_length\"] * cfg[\"sample_rate\"])\ntotal_samples = batch_size * (way * (shot + n_query))\nsummary(ConvNFeatureExtractor(), input_size=(total_samples, 1, 128, duration))","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:57:38.813101Z","iopub.execute_input":"2022-04-19T01:57:38.813993Z","iopub.status.idle":"2022-04-19T01:57:38.995884Z","shell.execute_reply.started":"2022-04-19T01:57:38.813945Z","shell.execute_reply":"2022-04-19T01:57:38.995144Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Classifier","metadata":{}},{"cell_type":"code","source":"def L2SquareDist(A: Tensor, B: Tensor, average: bool = True) -> Tensor:\n    r\"\"\"calculate parwise euclidean distance between two batchs of features.\n\n    Args:\n        A: Torch feature tensor. size:[Batch_size, Na, nC]\n        B: Torch feature tensor. size:[Batch_size, Nb, nC]\n    Output:\n        dist: The calculated distance tensor. size:[Batch_size, Na, Nb]\n    \"\"\"\n    assert A.dim() == 3\n    assert B.dim() == 3\n    assert A.size(0) == B.size(0) and A.size(2) == B.size(2)\n    nB = A.size(0)\n    Na = A.size(1)\n    Nb = B.size(1)\n    nC = A.size(2)\n\n    # AB = A * B = [nB x Na x nC] * [nB x nC x Nb] = [nB x Na x Nb]\n    AB = torch.bmm(A, B.transpose(1, 2))\n\n    AA = (A * A).sum(dim=2, keepdim=True).view(nB, Na, 1)  # [nB x Na x 1]\n    BB = (B * B).sum(dim=2, keepdim=True).view(nB, 1, Nb)  # [nB x 1 x Nb]\n    # l2squaredist = A*A + B*B - 2 * A * B\n    dist = AA.expand_as(AB) + BB.expand_as(AB) - 2 * AB\n    if average:\n        dist = dist / nC\n\n    return dist","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:58:53.805126Z","iopub.execute_input":"2022-04-19T01:58:53.805765Z","iopub.status.idle":"2022-04-19T01:58:53.883443Z","shell.execute_reply.started":"2022-04-19T01:58:53.805726Z","shell.execute_reply":"2022-04-19T01:58:53.882490Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class PN_head(nn.Module):\n    r\"\"\"The metric-based protypical classifier from ``Prototypical Networks for Few-shot Learning''.\n\n    Args:\n        metric: Whether use cosine or enclidean distance.\n        scale_cls: The initial scale number which affects the following softmax function.\n        learn_scale: Whether make scale number learnable.\n        normalize: Whether normalize each spatial dimension of image features before average pooling.\n    \"\"\"\n\n    def __init__(\n        self,\n        metric: str = \"euclidean\",\n        scale_cls: int = 10.0,\n        learn_scale: bool = True,\n        normalize: bool = True,\n        pool_feature: bool = False,\n    ) -> None:\n        super().__init__()\n        assert metric in [\"cosine\", \"euclidean\"]\n        if learn_scale:\n            self.scale_cls = nn.Parameter(\n                torch.FloatTensor(1).fill_(scale_cls), requires_grad=True\n            )\n        else:\n            self.scale_cls = scale_cls\n        self.metric = metric\n        self.normalize = normalize\n\n    def forward(\n        self, features_test: Tensor, features_train: Tensor, way: int, shot: int\n    ) -> Tensor:\n        r\"\"\"Take batches of few-shot training examples and testing examples as input,\n            output the logits of each testing examples.\n\n        Args:\n            features_test: Testing examples. size: [batch_size, num_query, d]\n            features_train: Training examples which has labels like:[abcdabcdabcd].\n                            size: [batch_size, way*shot, d]\n            way: The number of classes of each few-shot classification task.\n            shot: The number of training images per class in each few-shot classification\n                  task.\n        Output:\n            classification_scores: The calculated logits of testing examples.\n                                   size: [batch_size, num_query, way]\n        \"\"\"\n        assert features_train.dim() == 3\n        assert features_test.dim() == 3\n\n        batch_size = features_train.size(0)\n        if self.metric == \"cosine\":\n            features_train = F.normalize(features_train, p=2, dim=2, eps=1e-12)\n\n        # prototypes: [batch_size, way, c]\n        prototypes = torch.mean(\n            features_train.reshape(batch_size, shot, way, -1), dim=1\n        )\n        prototypes = F.normalize(prototypes, p=2, dim=2, eps=1e-12)\n\n        if self.metric == \"cosine\":\n            features_test = F.normalize(features_test, p=2, dim=2, eps=1e-12)\n            # [batch_size, num_query, c] * [batch_size, c, way] -> [batch_size, num_query, way]\n            classification_scores = -self.scale_cls * torch.bmm(\n                features_test, prototypes.transpose(1, 2)\n            )\n        elif self.metric == \"euclidean\":\n            classification_scores = self.scale_cls * L2SquareDist(\n                features_test, prototypes\n            )\n        return classification_scores","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:58:54.050744Z","iopub.execute_input":"2022-04-19T01:58:54.050981Z","iopub.status.idle":"2022-04-19T01:58:54.142018Z","shell.execute_reply.started":"2022-04-19T01:58:54.050955Z","shell.execute_reply":"2022-04-19T01:58:54.141178Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train/Evaluation Module","metadata":{}},{"cell_type":"code","source":"class CategoricalAccuracy(Metric):\n    def __init__(self, dist_sync_on_step=False):\n        super().__init__(dist_sync_on_step=dist_sync_on_step)\n        self.add_state(\"correct\", default=torch.tensor(0), dist_reduce_fx=\"sum\")\n        self.add_state(\"total\", default=torch.tensor(0), dist_reduce_fx=\"sum\")\n\n    def update(self, preds: torch.Tensor, target: torch.Tensor):\n        predictions = preds.argmax(dim=-1)\n        assert predictions.shape == target.shape\n        self.correct += torch.sum(predictions == target)\n        self.total += target.numel()\n\n    def compute(self):\n        return self.correct.float() / self.total","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:58:54.443502Z","iopub.execute_input":"2022-04-19T01:58:54.444167Z","iopub.status.idle":"2022-04-19T01:58:54.513071Z","shell.execute_reply.started":"2022-04-19T01:58:54.444130Z","shell.execute_reply":"2022-04-19T01:58:54.512182Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def epoch_wrapup(pl_module: LightningModule, mode: str):\n    r\"\"\"On the end of each epoch, log information of the whole\n        epoch and reset all metrics.\n\n    Args:\n        pl_module: An instance of LightningModule.\n        mode: The current mode (train, val or test).\n    \"\"\"\n    assert mode in [\"train\", \"val\", \"test\"]\n    value = getattr(pl_module, f\"{mode}_loss\").compute()\n    if mode == \"train\":\n        pl_module.log(f\"{mode}/loss_epoch\", value)\n    getattr(pl_module, f\"{mode}_loss\").reset()\n    value = getattr(pl_module, f\"{mode}_acc\").compute()\n    if mode == \"train\":\n        pl_module.log(f\"{mode}/acc_epoch\", value)\n    getattr(pl_module, f\"{mode}_acc\").reset()\n\n\ndef set_schedule(pl_module):\n    r\"\"\"Set the optimizer and scheduler for training.\n\n    Supported optimizer:\n        Adam and SGD\n    Supported scheduler:\n        cosine scheduler and decaying on specified epochs\n\n    Args:\n        pl_module: An instance of LightningModule.\n    \"\"\"\n    lr = pl_module.hparams.lr\n    wd = pl_module.hparams.weight_decay\n    decay_scheduler = pl_module.hparams.decay_scheduler\n    optim_type = pl_module.hparams.optim_type\n\n    if optim_type == \"adamw\":\n        optimizer = AdamW(pl_module.parameters(), weight_decay=wd, lr=lr, amsgrad=True)\n    elif optim_type == \"sgd\":\n        optimizer = SGD(\n            pl_module.parameters(), momentum=0.9, nesterov=True, weight_decay=wd, lr=lr\n        )\n    else:\n        raise RuntimeError(\n            \"optim_type not supported.\\\n                            Try to implement your own optimizer.\"\n        )\n\n    if decay_scheduler == \"cosine\":\n        max_steps = pl_module.trainer.max_steps\n        if (max_steps is None) or (max_steps == -1):\n            length_epoch = len(pl_module.trainer.datamodule.train_dataloader())\n            max_steps = length_epoch * pl_module.trainer.max_epochs\n\n        print(f\"max_steps: {max_steps}\")\n        scheduler = {\n            \"scheduler\": CosineAnnealingLR(optimizer, max_steps),\n            \"interval\": \"step\",\n        }\n        return {\"optimizer\": optimizer, \"lr_scheduler\": scheduler}\n    elif decay_scheduler == \"specified_epochs\":\n        decay_epochs = pl_module.hparams.decay_epochs\n        decay_power = pl_module.hparams.decay_power\n        assert decay_epochs is not None and decay_power is not None\n        scheduler = {\n            \"scheduler\": MultiStepLR(\n                optimizer, milestones=decay_epochs, gamma=decay_power\n            ),\n            \"interval\": \"epoch\",\n        }\n        return {\"optimizer\": optimizer, \"lr_scheduler\": scheduler}\n    elif decay_scheduler is None:\n        return optimizer\n    else:\n        raise RuntimeError(\n            \"decay scheduler not supported.\\\n                            Try to implement your own scheduler.\"\n        )","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:58:54.740893Z","iopub.execute_input":"2022-04-19T01:58:54.741399Z","iopub.status.idle":"2022-04-19T01:58:54.838019Z","shell.execute_reply.started":"2022-04-19T01:58:54.741344Z","shell.execute_reply":"2022-04-19T01:58:54.837234Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class BaseFewShotModule(LightningModule):\n    r\"\"\"Template for all few-shot learning models.\"\"\"\n\n    def __init__(\n        self,\n        backbone_name: str = \"ConvN\",\n        way: int = 5,\n        train_shot: Optional[int] = None,\n        val_shot: int = 5,\n        test_shot: int = 5,\n        num_query: int = 15,\n        train_batch_size_per_gpu: Optional[int] = None,\n        val_batch_size_per_gpu: int = 4,\n        test_batch_size_per_gpu: int = 4,\n        lr: float = 0.001,\n        weight_decay: float = 5e-4,\n        decay_scheduler: Optional[str] = \"cosine\",\n        optim_type: str = \"adamw\",\n        decay_epochs: Union[List, Tuple, None] = None,\n        decay_power: Optional[float] = None,\n        backbone_kwargs: Dict = {},\n    ) -> None:\n        \"\"\"\n        Args:\n            backbone_name: The name of the feature extractor,\n                        which should match the correspond\n                        file name in architectures.feature_extractor\n            way: The number of classes within one task.\n            train_shot: The number of samples within each few-shot\n                        support class during training.\n                        For meta-learning only.\n            val_shot: The number of samples within each few-shot\n                    support class during validation.\n            test_shot: The number of samples within each few-shot\n                    support class during testing.\n            num_query: The number of samples within each few-shot\n                    query class.\n            train_batch_size_per_gpu: The batch size of training per GPU.\n            val_batch_size_per_gpu: The batch size of validation per GPU.\n            test_batch_size_per_gpu: The batch size of testing per GPU.\n            lr: The initial learning rate.\n            weight_decay: The weight decay parameter.\n            decay_scheduler: The scheduler of optimizer.\n                            \"cosine\" or \"specified_epochs\".\n            optim_type: The optimizer type.\n                        \"sgd\" or \"adam\"\n            decay_epochs: The list of decay epochs of decay_scheduler \"specified_epochs\".\n            decay_power: The decay power of decay_scheduler \"specified_epochs\"\n                        at eachspeicified epoch.\n                        i.e., adjusted_lr = lr * decay_power\n            backbone_kwargs: The parameters for creating backbone network.\n        \"\"\"\n        super().__init__()\n        self.save_hyperparameters()\n        self.backbone = ConvNFeatureExtractor()\n        self.label = torch.arange(way, dtype=torch.int8).repeat(num_query)\n        self.label = self.label.type(torch.LongTensor).reshape(-1)\n\n        self.set_metrics()\n\n    def train_forward(self, batch):\n        r\"\"\"Here implements the forward function of training.\n\n        Output: logits\n        Args: (can be dynamically adjusted)\n            batch: a batch from train_dataloader.\n        \"\"\"\n        raise NotImplementedError\n\n    def val_test_forward(self, batch, batch_size, way, shot):\n        r\"\"\"Here implements the forward function of validation and testing.\n\n        Output: logits\n        Args: (can be dynamically adjusted)\n            batch: a batch from val_dataloader.\n            batch_size: number of tasks during one iteration.\n            way: The number of classes within one task.\n            shot: The number of samples within each few-shot support class.\n        \"\"\"\n        raise NotImplementedError\n\n    def shared_step(self, batch, mode):\n        r\"\"\"The shared operation across\n            validation, testing and potentially training (meta-learning).\n\n        Args:\n            batch: a batch from val_dataloader.\n            mode: train, val or test\n        \"\"\"\n        assert mode in [\"train\", \"val\", \"test\"]\n        if mode == \"train\":\n            flag = \"train\"\n        else:\n            flag = \"val_test\"\n        foward_function = getattr(self, f\"{flag}_forward\")\n        batch_size_per_gpu = getattr(self.hparams, f\"{mode}_batch_size_per_gpu\")\n        shot = getattr(self.hparams, f\"{mode}_shot\")\n        # label\n        # print(batch[0].shape)\n        distances = foward_function(batch, batch_size_per_gpu, self.hparams.way, shot)\n        # import pdb\n        # pdb.set_trace()\n        label = (\n            torch.unsqueeze(self.label, 0)\n            .repeat(batch_size_per_gpu, 1)\n            .reshape(-1)\n            .to(distances.device)\n        )\n        distances = distances.reshape(label.size(0), -1)\n\n        # loss = F.cross_entropy(-distances, label)\n        loss = F.nll_loss(F.log_softmax(-distances, dim=1), label)\n\n        y_pred = (-distances).softmax(dim=1)\n\n        log_loss = getattr(self, f\"{mode}_loss\")(loss)\n        accuracy = getattr(self, f\"{mode}_acc\")(y_pred, label)\n        self.log(f\"{mode}/loss\", log_loss)\n        self.log(f\"{mode}/acc\", accuracy)\n        return loss\n\n    def training_step(self, batch, batch_idx):\n        if (\n            self.hparams.train_shot == None\n            or self.hparams.train_batch_size_per_gpu == None\n        ):\n            raise RuntimeError(\n                \"train_shot or train_batch_size not specified.\\\n                                Please implement your own training step if the\\\n                                 training is not meta-learning.\"\n            )\n        return self.shared_step(batch, \"train\")\n\n    def validation_step(self, batch, batch_idx):\n        _ = self.shared_step(batch, \"val\")\n\n    def test_step(self, batch, batch_idx):\n        _ = self.shared_step(batch, \"test\")\n\n    def training_epoch_end(self, outs):\n        epoch_wrapup(self, \"train\")\n\n    def validation_epoch_end(self, outs):\n        epoch_wrapup(self, \"val\")\n\n    def test_epoch_end(self, outs):\n        epoch_wrapup(self, \"test\")\n\n    def configure_optimizers(self):\n        return set_schedule(self)\n\n    def set_metrics(self):\n        r\"\"\"Set basic logging metrics for few-shot learning.\"\"\"\n        for split in [\"train\", \"val\", \"test\"]:\n            setattr(self, f\"{split}_loss\", MeanMetric())\n            setattr(self, f\"{split}_acc\", CategoricalAccuracy())","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:58:55.048955Z","iopub.execute_input":"2022-04-19T01:58:55.049505Z","iopub.status.idle":"2022-04-19T01:58:55.179293Z","shell.execute_reply.started":"2022-04-19T01:58:55.049458Z","shell.execute_reply":"2022-04-19T01:58:55.178508Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ProtoNet(BaseFewShotModule):\n    r\"\"\"The datamodule implementing Prototypical Network.\"\"\"\n\n    def __init__(\n        self,\n        metric: str = \"euclidean\",\n        scale_cls: float = 10.0,\n        normalize: bool = True,\n        backbone_name: str = \"ConvN\",\n        way: int = 5,\n        train_shot: int = 5,\n        val_shot: int = 5,\n        test_shot: int = 5,\n        num_query: int = 15,\n        train_batch_size_per_gpu: int = 1,\n        val_batch_size_per_gpu: int = 4,\n        test_batch_size_per_gpu: int = 4,\n        lr: float = 0.1,\n        weight_decay: float = 5e-4,\n        decay_scheduler: Optional[str] = \"cosine\",\n        optim_type: str = \"adamw\",\n        decay_epochs: Union[List, Tuple, None] = None,\n        decay_power: Optional[float] = None,\n        backbone_kwargs: Dict = {},\n        **kwargs,\n    ) -> None:\n        \"\"\"\n        Args:\n            metric: what metrics applied. \"cosine\" or \"euclidean\".\n            scale_cls: The initial scale number which affects the\n                    following softmax function.\n            normalize: Whether normalize each spatial dimension of image features before average pooling.\n            backbone_name: The name of the feature extractor,\n                        which should match the correspond\n                        file name in architectures.feature_extractor\n            way: The number of classes within one task.\n            train_shot: The number of samples within each few-shot\n                        support class during training.\n                        For meta-learning only.\n            val_shot: The number of samples within each few-shot\n                    support class during validation.\n            test_shot: The number of samples within each few-shot\n                    support class during testing.\n            num_query: The number of samples within each few-shot\n                    query class.\n            train_batch_size_per_gpu: The batch size of training per GPU.\n            val_batch_size_per_gpu: The batch size of validation per GPU.\n            test_batch_size_per_gpu: The batch size of testing per GPU.\n            lr: The initial learning rate.\n            weight_decay: The weight decay parameter.\n            decay_scheduler: The scheduler of optimizer.\n                            \"cosine\" or \"specified_epochs\".\n            optim_type: The optimizer type.\n                        \"sgd\" or \"adam\"\n            decay_epochs: The list of decay epochs of decay_scheduler \"specified_epochs\".\n            decay_power: The decay power of decay_scheduler \"specified_epochs\"\n                        at eachspeicified epoch.\n                        i.e., adjusted_lr = lr * decay_power\n            backbone_kwargs: The parameters for creating backbone network.\n        \"\"\"\n        super().__init__(\n            backbone_name,\n            way,\n            train_shot,\n            val_shot,\n            test_shot,\n            num_query,\n            train_batch_size_per_gpu,\n            val_batch_size_per_gpu,\n            test_batch_size_per_gpu,\n            lr,\n            weight_decay,\n            decay_scheduler,\n            optim_type,\n            decay_epochs,\n            decay_power,\n            backbone_kwargs,\n        )\n        self.classifier = PN_head(metric, scale_cls, normalize=normalize)\n\n    def forward(self, batch, batch_size, way, shot):\n        r\"\"\"Since PN is a meta-learning method,\n            the model forward process is the same for train, val and test.\n\n        Args:\n            batch: a batch from val_dataloader.\n            batch_size: number of tasks during one iteration.\n            way: The number of classes within one task.\n            shot: The number of samples within each few-shot support class.\n        \"\"\"\n        num_support_samples = way * shot\n        data, _ = batch\n        data = self.backbone(data)  # (B * (N + Q), D)\n\n        assert data.dim() == 2\n        data = data.reshape([batch_size, -1, data.size(-1)])  # (B, N + Q, D)\n        data_support = data[:, :num_support_samples]  # (B, N, D)\n        data_query = data[:, num_support_samples:]  # (B, Q, D)\n        distances = self.classifier(data_query, data_support, way, shot)\n        return distances\n\n    def train_forward(self, batch, batch_size, way, shot):\n        return self(batch, batch_size, way, shot)\n\n    def val_test_forward(self, batch, batch_size, way, shot):\n        return self(batch, batch_size, way, shot)","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:58:55.399199Z","iopub.execute_input":"2022-04-19T01:58:55.399721Z","iopub.status.idle":"2022-04-19T01:58:55.499802Z","shell.execute_reply.started":"2022-04-19T01:58:55.399685Z","shell.execute_reply":"2022-04-19T01:58:55.499070Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_batch_size = 2\nval_batch_size = 4\n\nmodel = ProtoNet(\n    train_batch_size_per_gpu=train_batch_size, val_batch_size_per_gpu=val_batch_size\n)\n\n# trainer = Trainer(precision=16)  # need torch >= 1.10\n# trainer = Trainer()\nwandb_logger = WandbLogger(\n    project=\"BirdCLEF22\",\n    name=\"5-way-5-shot-prototypical-conv4w-flatten-7s\",\n    log_model=\"all\",\n)\ncallbacks = [\n    LearningRateMonitor(logging_interval=\"step\"),\n    ModelCheckpoint(verbose=True, save_last=True, monitor=\"val/acc\", mode=\"max\"),\n]\n\ntrainer = Trainer(\n    accelerator=\"gpu\",\n    devices=1,\n    max_epochs=cfg[\"max_epochs\"],\n    logger=wandb_logger,\n    callbacks=callbacks,\n    # precision=16,\n)\ndatamodule = FewShotDataModule(\n    train,\n    val,\n    img_path,\n    cfg,\n    train_transform,\n    val_transform,\n    train_batch_size=train_batch_size,\n    val_batch_size=val_batch_size,\n)\n\nseed_everything(cfg[\"random_seed\"])\ntry:\n    trainer.fit(model, datamodule=datamodule)\nfinally:\n    wandb.finish()","metadata":{"execution":{"iopub.status.busy":"2022-04-19T01:59:05.506155Z","iopub.execute_input":"2022-04-19T01:59:05.506912Z","iopub.status.idle":"2022-04-19T02:02:00.875163Z","shell.execute_reply.started":"2022-04-19T01:59:05.506861Z","shell.execute_reply":"2022-04-19T02:02:00.874452Z"},"trusted":true},"execution_count":null,"outputs":[]}]}