{"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":"# A clean, fast, simple bird identifier training pipeline in pytorch-lightning\n\nThis notebook is based on [kenroma's baseline](https://www.kaggle.com/kneroma/clean-fast-simple-bird-identifier-training-colab), and I use pytorch-lightning to control my workflow.\n\nI'm new to deep learning, and I want to learn some skills from kaggle, any suggestions are welcome.\n\nI will try to use bigger models, augmentation tircks, and so on.\n\nYou can find external dataset from this [notebook](https://www.kaggle.com/kneroma/clean-fast-simple-bird-identifier-training-colab), enjoy your kaggle journey.","metadata":{}},{"cell_type":"markdown","source":"## Prepare for the environment\n\n### import packages","metadata":{}},{"cell_type":"code","source":"!pip install timm\n!pip install librosa","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport librosa as lb\nimport librosa.display as lbd\nimport soundfile as sf\nfrom  soundfile import SoundFile\nimport pandas as pd\nfrom  IPython.display import Audio\nfrom pathlib import Path\n\nimport torch\nfrom torch import nn, optim\nfrom  torch.utils.data import Dataset, DataLoader\n\nimport timm\nimport pytorch_lightning as pl\n\nfrom matplotlib import pyplot as plt\n\nimport os, random, gc\nimport re, time, json\nfrom  ast import literal_eval\n\n\nfrom IPython.display import Audio\nfrom sklearn.metrics import label_ranking_average_precision_score\nfrom sklearn.model_selection import train_test_split\n\nfrom tqdm.notebook import tqdm\nimport joblib","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Set seed for everything","metadata":{}},{"cell_type":"code","source":"def seed_everything(seed=42):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\nseed_everything()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### define a config class for experiement\n\nThis part is used to control our whole training pipeline.\n\nYou can change model by reset model_name, such model list can check [here](https://rwightman.github.io/pytorch-image-models/)","metadata":{}},{"cell_type":"code","source":"class config:\n    debug          = True\n    num_workers    = 8\n    epochs = 12\n    model_name     = 'resnest50d_4s2x40d'\n    pretrained_path = '../input/resnest50-fast-package/resnest50_fast_4s2x40d-41d14ed0.pth'\n    num_classes = 397\n    sr = 32_000\n    duration = 7\n    max_read_samples = 5 # Each record will have 10 melspecs at most, you can increase this on Colab with High Memory Enabled\n    data_root = Path(\"\")\n    mel_paths = sorted(Path(\"../input/\").glob(\"kkiller-birdclef-mels-computer-d7-part?/rich_train_metadata.csv\"))\n    train_label_paths = sorted(Path(\"../input/\").glob(\"kkiller-birdclef-mels-computer-d7-part?/LABEL_IDS.json\"))\n#     model_root = Path(\"../input/\")\n    batch_size = 128\n    num_workers = 8\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    lr = 1e-3\n    weight_decay = 0\n    save_path = './pl_pipeline'\n    precision = 16\n    gradient_clip_val = 0.1\n    seed = 42","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Prepare data for training\n\nThis part is used to prepare data for training. We will do a quick look at this dataset use dataframe","metadata":{}},{"cell_type":"code","source":"def get_df(mel_paths, train_label_paths):\n    df = None\n    label_ids = {}\n    \n    for file_path in mel_paths:\n        temp = pd.read_csv(str(file_path), index_col=0)\n        temp[\"impath\"] = temp.apply(lambda row: file_path.parent/\"audio_images/{}/{}.npy\".format(row.primary_label, row.filename), axis=1) \n        df = temp if df is None else df.append(temp)\n    \n    df[\"secondary_labels\"] = df[\"secondary_labels\"].apply(literal_eval)\n\n    for file_path in train_label_paths:\n        with open(str(file_path)) as f:\n            label_ids.update(json.load(f))\n\n    return label_ids, df","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"label_ids, df = get_df(config.mel_paths, config.train_label_paths)\nprint(df.shape)\ndf.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df[\"primary_label\"].value_counts()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df[\"label_id\"].min(), df[\"label_id\"].max()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Show the pictures\n\nLoad all data and make a quick look at the picture format sound","metadata":{}},{"cell_type":"code","source":"def load_data(config, df):\n    def load_row(row):\n        # impath = TRAIN_IMAGES_ROOT/f\"{row.primary_label}/{row.filename}.npy\"\n        return row.filename, np.load(str(row.impath))[:config.max_read_samples]\n    pool = joblib.Parallel(4)\n    mapper = joblib.delayed(load_row)\n    tasks = [mapper(row) for row in df.itertuples(False)]\n    res = pool(tqdm(tasks))\n    res = dict(res)\n    return res","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# We cache the train set to reduce training time\naudio_image_store = load_data(config, df)\nlen(audio_image_store)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"shape:\", next(iter(audio_image_store.values())).shape)\nlbd.specshow(next(iter(audio_image_store.values()))[0])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pd.Series([len(x) for x in audio_image_store.values()]).value_counts()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Define a dataset class\n\nThis part is used to define a dataset class for our dataloader. It organize the data into three channel images","metadata":{}},{"cell_type":"code","source":"class BirdClefDataset(Dataset):\n\n    def __init__(self, audio_image_store, meta, config=config, is_train=True):\n        \n        self.audio_image_store = audio_image_store\n        self.meta = meta.copy().reset_index(drop=True)\n        self.sr = config.sr\n        self.is_train = is_train\n        self.num_classes = config.num_classes\n        self.duration = config.duration\n        self.audio_length = self.duration*self.sr\n    \n    @staticmethod\n    def normalize(image):\n        image = image.astype(\"float32\", copy=False) / 255.0\n        image = np.stack([image, image, image])\n        return image\n\n    def __len__(self):\n        return len(self.meta)\n    \n    def __getitem__(self, idx):\n        row = self.meta.iloc[idx]\n        image = self.audio_image_store[row.filename]\n\n        image = image[np.random.choice(len(image))]\n        image = self.normalize(image)\n        \n        t = np.zeros(self.num_classes, dtype=np.float32) + 0.0025 # Label smoothing\n        t[row.label_id] = 0.995\n        \n        return image, t","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ds = BirdClefDataset(audio_image_store, meta=df, config=config, is_train=True)\nlen(df)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"x, y = ds[np.random.choice(len(ds))]\nx.shape, y.shape, np.where(y >= 0.5)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"lbd.specshow(x[0])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"y[:5]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define LitModule","metadata":{}},{"cell_type":"code","source":"class LitModel(pl.LightningModule):\n\n    def __init__(self, config, audio_image_store, df):\n        super().__init__()\n        \"\"\"\n        Loads a pretrained model from timm (https://github.com/rwightman/pytorch-image-models).\n\n        Arguments:\n            config {object} -- config of the whole notebook\n        \"\"\"\n        self.config = config\n        self.model = timm.create_model(config.model_name, pretrained = False)\n        self.model.load_state_dict(torch.load(config.pretrained_path))\n        \n        if hasattr(self.model, \"fc\"):\n            nb_ft = self.model.fc.in_features\n            self.model.fc = nn.Linear(nb_ft, config.num_classes)\n        elif hasattr(self.model, \"_fc\"):\n            nb_ft = self.model._fc.in_features\n            self.model._fc = nn.Linear(nb_ft, config.num_classes)\n        elif hasattr(self.model, \"classifier\"):\n            nb_ft = model.classifier.in_features\n            self.model.classifier = nn.Linear(nb_ft, config.num_classes)\n        elif hasattr(self.model, \"last_linear\"):\n            nb_ft = self.model.last_linear.in_features\n            self.model.last_linear = nn.Linear(nb_ft, config.num_classes)\n            \n        train_idx, val_idx = train_test_split(np.arange(len(df.index)), test_size=0.2, random_state=config.seed)\n        self.trainset = BirdClefDataset(audio_image_store, meta=df.iloc[train_idx].reset_index(drop=True),\n                                        config=config, is_train=True)\n        \n        self.valset = BirdClefDataset(audio_image_store, meta=df.iloc[val_idx].reset_index(drop=True),\n                                      config=config, is_train=False)\n\n    def forward(self, x):\n        # in lightning, forward defines the prediction/inference actions\n        output = self.model(x)\n        return output\n    \n    def train_dataloader(self):\n        # Simply define a pytorch dataloader here that will take care of batching. Note it works well with dictionnaries !\n        train_dl = DataLoader(self.trainset, batch_size=self.config.batch_size, shuffle=True,\n                                    num_workers=self.config.num_workers)\n        return train_dl\n    \n    def val_dataloader(self):\n        # Simply define a pytorch dataloader here that will take care of batching. Note it works well with dictionnaries !\n        val_dl = DataLoader(self.valset, batch_size=self.config.batch_size, shuffle=False,\n                                    num_workers=self.config.num_workers)\n        return val_dl\n    \n    def loss_function(self, preds, labels):\n        # How to calculate the loss. Note this method is actually not a part of pytorch lightning ! It's only good practice\n        loss_fn = nn.BCEWithLogitsLoss()  # Let's rebalance the weights for each class here.\n#         loss_fn = FocalLoss(logits=True)\n        loss = loss_fn(preds, labels)\n        return loss\n    \n    def train_dataloader(self):\n        # Simply define a pytorch dataloader here that will take care of batching. Note it works well with dictionnaries !\n        train_dl = DataLoader(self.trainset, batch_size=self.config.batch_size, shuffle=True,\n                                    num_workers=self.config.num_workers)\n        return train_dl\n    \n    def training_step(self, batch, batch_idx):\n        # training_step defined the train loop.\n        # It is independent of forward\n        x, y = batch\n        x, y = x.to(self.config.device), y.to(self.config.device)\n        \n        output = self(x)\n        loss = self.loss_function(output, y)\n        \n        with torch.no_grad():\n            output = output.sigmoid()\n            y = (y > 0.5 )*1.0\n            label_rank_avg_precision_score = label_ranking_average_precision_score(y.cpu().numpy(), output.cpu().numpy())\n\n            output = (output > 0.5)*1.0\n\n            precision = (output*y).sum()/(1e-6 + output.sum())\n            recall = (output*y).sum()/(1e-6 + y.sum())\n            f1 = 2*precision*recall/(1e-6+precision+recall)\n        \n        self.log('train_loss', loss, on_step=True, on_epoch=True, \n                 prog_bar=True, logger=True) \n        self.log('train_label_rank_avg_precision_score', label_rank_avg_precision_score,\n                 on_step=True, on_epoch=True, \n                 prog_bar=True, logger=True)\n        self.log('train_f1', f1, on_step=True, on_epoch=True, \n                 prog_bar=True, logger=True)\n        self.log('train_recall', recall, on_step=True, on_epoch=True, \n                 prog_bar=True, logger=True)\n        self.log('train_precision',precision, on_step=True, on_epoch=True, \n                 prog_bar=True, logger=True)\n        return loss\n    \n    def validation_step(self, batch, batch_idx):\n        # This is where you must define what happens during a validation step (per batch)\n        x, y = batch\n        x, y = x.to(self.config.device), y.to(self.config.device)\n        output = self(x)\n        \n        return {'output': output, 'y': y}\n    \n    def validation_epoch_end(self, outputs):\n        # This is what happens at the end of validation epoch. Usually gathering all predictions\n        # outputs is a list of dictionary from each step.\n        outputs_, y_ = [], []\n        for output in outputs:\n            y_.append(output['y'])\n            outputs_.append(output['output'])\n        \n        y_ = torch.cat(y_)\n        outputs_ = torch.cat(outputs_)\n        loss = self.loss_function(outputs_, y_)\n        outputs_ = outputs_.sigmoid()\n        y_ = (y_ > 0.5 )*1.0\n        label_rank_avg_precision_score = label_ranking_average_precision_score(y_.cpu().numpy(), outputs_.cpu().numpy())\n\n        outputs_ = (outputs_ > 0.5)*1.0\n\n        precision = (outputs_*y_).sum()/(1e-6 + outputs_.sum())\n        recall = (outputs_*y_).sum()/(1e-6 + y_.sum())\n        f1 = 2*precision*recall/(1e-6+precision+recall)\n        self.log('val_loss', loss, on_epoch=True, \n                 prog_bar=True, logger=True) \n        self.log('val_label_rank_avg_precision_score', label_rank_avg_precision_score,\n                 on_epoch=True, prog_bar=True, logger=True)\n        self.log('val_f1', f1, on_epoch=True, \n                 prog_bar=True, logger=True)\n        self.log('val_recall', recall, on_epoch=True, \n                 prog_bar=True, logger=True)\n        self.log('val_precision',precision, on_epoch=True, \n                 prog_bar=True, logger=True)\n    \n    def configure_optimizers(self):\n        # Optimizers and schedulers. Note that each are in lists of equal length to allow multiple optimizers (for GAN for example)\n        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config.lr, weight_decay=self.config.weight_decay)\n        scheduler = scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, eta_min=1e-5, T_max=config.epochs)\n        \n        return [optimizer], [scheduler]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = LitModel(config, audio_image_store, df)\n\ncheckpoint_callback = pl.callbacks.ModelCheckpoint(\n                                        dirpath=config.save_path,\n                                        filename=f'{config.model_name}'+'-{epoch}-{val_loss:.3f}-{val_f1:.3f}',\n                                        save_weights_only=True,\n                                        monitor=\"val_loss\",\n                                        mode=\"min\",\n                                        save_last=True,\n                                    )\n\nearly_stop_callback = pl.callbacks.EarlyStopping(monitor=\"val_loss\", mode=\"min\")\n# Define trainer\n# Here you can\ntrainer = pl.Trainer(\n                    gpus=1,\n                    callbacks=[checkpoint_callback, early_stop_callback],\n                    max_epochs=1 if config.debug else config.epochs,\n                    gradient_clip_val=config.gradient_clip_val,\n                    precision=config.precision,\n                   )","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer.fit(model)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"checkpoint_callback.best_model_path","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"best_model = torch.load(checkpoint_callback.best_model_path)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.save(best_model['state_dict'], f'{config.save_path}/best_model.pth')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}