{"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":"# PyTorch Lightning Faster R-CNN Finetuning LB=0.358\n\n\nFinetune TorchVision's [Faster R-CNN](https://pytorch.org/vision/stable/models.html#id57) with [PyTorch Lightninng](https://www.pytorchlightning.ai/).\n\nTracks the training with [Weights & Biases](https://wandb.ai/site).\n\nCorresponding Inference kernel: [Starter PyTorch Lightning Faster R-CNN Inference](https://www.kaggle.com/clemchris/starter-pytorch-lightning-faster-r-cnn-inference)\n\n## Lightninng Features Used in this Kernel:\n- Easily switch between CPU and GPU training (`gpus=[0/1]` Trainer flag)\n- Quickly check if complete training is running without errors with `fast_dev_run=True` Trainer flag\n- Use half precision training on GPU with `precision=16` Trainer flag\n- Log learning rate with `LearningRateMonitor` callback\n- Log losses and metrics to Weights & Biases with `WandbLogger`\n\n## Ideas for next steps:\n- Add data augmentation\n- Use also those image without annotations for training\n- Try other TorchVision models\n- Try TPU training (made easy with PyTorch Lightning's [TPU Support](https://pytorch-lightning.readthedocs.io/en/latest/advanced/tpu.html))\n- Tune hyperparameters of model\n- Visualize inputs and predictions with Weights & Biases\n\n## Sources and Inspirations:\n- [TorchVision Object Detection Finetuning Tutorial](https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html)\n- [Reef- Starter Torch FasterRCNN Train [LB=0.416]](https://www.kaggle.com/julian3833/reef-starter-torch-fasterrcnn-train-lb-0-416)\n\n","metadata":{}},{"cell_type":"markdown","source":"## Installs\n\nLet's use the latest versions of PyTorch Lightning and [TorchMetrics](https://torchmetrics.readthedocs.io/en/latest/) to be able to use TorchMetric's [MAP](https://torchmetrics.readthedocs.io/en/latest/references/modules.html#map) metric, which needs [pycocotools](https://github.com/cocodataset/cocoapi/tree/master/PythonAPI/pycocotools).","metadata":{}},{"cell_type":"code","source":"%%capture\n!pip install pytorch-lightning==1.5.3 torchmetrics==0.6.0 pycocotools","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-01-03T18:26:59.406336Z","iopub.execute_input":"2022-01-03T18:26:59.406660Z","iopub.status.idle":"2022-01-03T18:27:17.443615Z","shell.execute_reply.started":"2022-01-03T18:26:59.406575Z","shell.execute_reply":"2022-01-03T18:27:17.442731Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Imports","metadata":{}},{"cell_type":"code","source":"import ast\nimport math\nimport multiprocessing as mp\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\nimport pytorch_lightning as pl\nimport torch\nimport torchmetrics\nimport torchvision\nimport wandb\nfrom PIL import Image\nfrom pytorch_lightning.callbacks import LearningRateMonitor\nfrom pytorch_lightning.loggers import WandbLogger\nfrom torchmetrics.metric import Metric\nfrom torchvision.datasets import VisionDataset\nfrom torchvision.models.detection import fasterrcnn_resnet50_fpn\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torch.utils.data.dataloader import DataLoader","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-01-03T18:27:17.445815Z","iopub.execute_input":"2022-01-03T18:27:17.446077Z","iopub.status.idle":"2022-01-03T18:27:24.582398Z","shell.execute_reply.started":"2022-01-03T18:27:17.446040Z","shell.execute_reply":"2022-01-03T18:27:24.581598Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load Weights & Biases API Key from secrets\n- Copy key from [W&B settings page](https://wandb.ai/settings)\n- Add key to Kaggle through [Add-ons](https://www.kaggle.com/product-feedback/114053) ","metadata":{}},{"cell_type":"code","source":"from kaggle_secrets import UserSecretsClient\nsecret_label = \"wb_api\"\nwb_api_key = UserSecretsClient().get_secret(secret_label)\nwandb.login(key=wb_api_key)","metadata":{"execution":{"iopub.status.busy":"2022-01-03T18:27:24.584020Z","iopub.execute_input":"2022-01-03T18:27:24.584269Z","iopub.status.idle":"2022-01-03T18:27:26.054087Z","shell.execute_reply.started":"2022-01-03T18:27:24.584232Z","shell.execute_reply":"2022-01-03T18:27:26.053365Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Paths","metadata":{}},{"cell_type":"code","source":"INPUT_DIR = Path(\"../input\")\nDATA_DIR = INPUT_DIR / \"tensorflow-great-barrier-reef\"\nTRAIN_CSV_PATH = DATA_DIR / \"train.csv\"","metadata":{"execution":{"iopub.status.busy":"2022-01-03T18:27:26.055211Z","iopub.execute_input":"2022-01-03T18:27:26.055453Z","iopub.status.idle":"2022-01-03T18:27:26.060198Z","shell.execute_reply.started":"2022-01-03T18:27:26.055420Z","shell.execute_reply":"2022-01-03T18:27:26.059242Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Settings","metadata":{}},{"cell_type":"code","source":"# Data Module Args\nNUM_WORKERS = mp.cpu_count()\nBATCH_SIZE = 8\n\nWANDB_PROJECT = \"kaggle-great-barrier-reef\"\n\n# Trainer Args\nGPUS = 1             # Set to 1 if GPU is enabled for notebook\nFAST_DEV_RUN = True  # Set to False to properly train\nMAX_EPOCHS = 5","metadata":{"execution":{"iopub.status.busy":"2022-01-03T18:27:26.062668Z","iopub.execute_input":"2022-01-03T18:27:26.063053Z","iopub.status.idle":"2022-01-03T18:27:26.070144Z","shell.execute_reply.started":"2022-01-03T18:27:26.063016Z","shell.execute_reply":"2022-01-03T18:27:26.069374Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Competition Metric implemented with TorchMetrics\n\nMostly copied from https://www.kaggle.com/bamps53/competition-metric-implementation","metadata":{}},{"cell_type":"code","source":"def f_beta(tp, fp, fn, beta=2):\n    return (1+beta**2)*tp / ((1+beta**2)*tp + beta**2*fn+fp)","metadata":{"execution":{"iopub.status.busy":"2022-01-03T18:27:26.071399Z","iopub.execute_input":"2022-01-03T18:27:26.071845Z","iopub.status.idle":"2022-01-03T18:27:26.080200Z","shell.execute_reply.started":"2022-01-03T18:27:26.071809Z","shell.execute_reply":"2022-01-03T18:27:26.079560Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class KaggleF2(Metric):\n    def __init__(\n        self,\n        compute_on_step=True,\n        dist_sync_on_step=False,\n        process_group=None,\n        dist_sync_fn=None,\n    ) -> None:\n        super().__init__(\n            compute_on_step=compute_on_step,\n            dist_sync_on_step=dist_sync_on_step,\n            process_group=process_group,\n            dist_sync_fn=dist_sync_fn,\n        )\n\n        self.add_state(\"detection_boxes\", default=[], dist_reduce_fx=None)\n        self.add_state(\"detection_scores\", default=[], dist_reduce_fx=None)\n        self.add_state(\"groundtruth_boxes\", default=[], dist_reduce_fx=None)\n\n    def update(self, preds, target):\n        for item in preds:\n            self.detection_boxes.append(\n                torchvision.ops.box_convert(item[\"boxes\"], in_fmt=\"xywh\", out_fmt=\"xyxy\")\n                if len(item[\"boxes\"]) > 0\n                else item[\"boxes\"]\n            )\n            self.detection_scores.append(item[\"scores\"])\n\n        for item in target:\n            self.groundtruth_boxes.append(\n                torchvision.ops.box_convert(item[\"boxes\"], in_fmt=\"xywh\", out_fmt=\"xyxy\")\n                if len(item[\"boxes\"]) > 0\n                else item[\"boxes\"]\n            )\n\n    def compute(self):\n        tps, fps, fns = 0, 0, 0\n        for gt_boxes, pred_boxes, pred_scores in zip(\n            self.groundtruth_boxes, self.detection_boxes, self.detection_scores\n        ):\n            tp, fp, fn = self._compute_stat_scores(gt_boxes, pred_boxes, pred_scores)\n            tps += tp\n            fps += fp\n            fns += fn\n\n        return f_beta(tps, fps, fns, beta=2)\n\n    def _compute_stat_scores(self, gt_boxes, pred_boxes, pred_scores):\n        if len(gt_boxes) == 0 and len(pred_boxes) == 0:\n            tps, fps, fns = 0, 0, 0\n            return tps, fps, fns\n\n        elif len(gt_boxes) == 0:\n            tps, fps, fns = 0, len(pred_boxes), 0\n            return tps, fps, fns\n\n        elif len(pred_boxes) == 0:\n            tps, fps, fns = 0, 0, len(gt_boxes)\n            return tps, fps, fns\n\n        # sort by conf\n        _, indices = torch.sort(pred_scores, descending=True)\n        pred_boxes = pred_boxes[indices]\n\n        tps, fps, fns = 0, 0, 0\n        for iou_th in np.arange(0.3, 0.85, 0.05):\n            tp, fp, fn = self._compute_stat_scores_at_iou_th(gt_boxes, pred_boxes, iou_th)\n            tps += tp\n            fps += fp\n            fns += fn\n\n        return tps, fps, fns\n\n    def _compute_stat_scores_at_iou_th(self, gt_boxes, pred_boxes, iou_th):\n        gt_boxes = gt_boxes.clone()\n        pred_boxes = pred_boxes.clone()\n\n        tp = 0\n        fp = 0\n        for k, pred_bbox in enumerate(pred_boxes):\n            ious = torchvision.ops.box_iou(gt_boxes, pred_bbox[None, ...])\n\n            max_iou = ious.max()\n            if max_iou > iou_th:\n                tp += 1\n                \n                # Delete max_iou box\n                argmax_iou = ious.argmax()\n                gt_boxes = torch.cat([gt_boxes[0:argmax_iou], gt_boxes[argmax_iou+1:]])\n            else:\n                fp += 1\n            if len(gt_boxes) == 0:\n                fp += len(pred_boxes) - (k + 1)\n                break\n\n        fn = len(gt_boxes)\n\n        return tp, fp, fn\n","metadata":{"execution":{"iopub.status.busy":"2022-01-03T18:27:26.081669Z","iopub.execute_input":"2022-01-03T18:27:26.082145Z","iopub.status.idle":"2022-01-03T18:27:26.102162Z","shell.execute_reply.started":"2022-01-03T18:27:26.082111Z","shell.execute_reply":"2022-01-03T18:27:26.101500Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## PyTorch Dataset Class","metadata":{}},{"cell_type":"code","source":"class GBRDataset(VisionDataset):\n    \"\"\"Custom VisionDataset class that creates a dataset from the train DataFrame. \n       \n       Uses only those images with annotations.\n    \"\"\"\n    \n    def __init__(self, csv_path):\n        super().__init__(csv_path.parent)\n\n        self.train_df = pd.read_csv(csv_path)\n        self.image_paths_annotations = self._create_image_paths_annotations()\n\n    def _create_image_paths_annotations(self):\n        \"\"\"Iterate over train DataFrame and extract image paths and annotations.\"\"\"\n        \n        image_paths_annotations = []\n\n        for _, row in self.train_df.iterrows():\n            image_path = Path(self.root) / \"train_images\" / f\"video_{row['video_id']}\" / f\"{row['video_frame']}.jpg\"\n            annotations = ast.literal_eval(row[\"annotations\"])\n\n            # Use only those images with annotations\n            if annotations:\n                image_paths_annotations.append((image_path, annotations))\n\n        return image_paths_annotations\n\n    def __getitem__(self, index):\n        image_path, annotations = self.image_paths_annotations[index]\n\n        # Image\n        image = Image.open(image_path).convert(\"RGB\")\n        image = np.array(image, dtype=np.float32) / 255.0\n\n        # Convert to pascal_voc boxes\n        boxes = []\n        for annotation in annotations:\n            x_min, y_min, width, height = (\n                annotation[\"x\"],\n                annotation[\"y\"],\n                annotation[\"width\"],\n                annotation[\"height\"],\n            )\n            \n            x_max = x_min + width\n            y_max = y_min + height\n\n            box = (x_min, y_min, x_max, y_max)\n            boxes.append(box)\n\n        labels = [1] * len(boxes)\n\n        # Convert to tensors\n        image = torch.from_numpy(image.transpose(2, 0, 1))\n\n        target = {}\n        target[\"boxes\"] = torch.as_tensor(boxes, dtype=torch.float32)\n        target[\"labels\"] = torch.as_tensor(labels, dtype=torch.int64)\n\n        return image, target\n\n    def __len__(self) -> int:\n        return len(self.image_paths_annotations)","metadata":{"execution":{"iopub.status.busy":"2022-01-03T18:29:10.392159Z","iopub.execute_input":"2022-01-03T18:29:10.392422Z","iopub.status.idle":"2022-01-03T18:29:10.406768Z","shell.execute_reply.started":"2022-01-03T18:29:10.392393Z","shell.execute_reply":"2022-01-03T18:29:10.403708Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Lightning Module Class","metadata":{}},{"cell_type":"code","source":"class GBRModule(pl.LightningModule):\n    \"\"\"LightningModule class to finetune torchvision's Faster R-CNN model.\"\"\"\n    \n    def __init__(self, pretrained_weights_path=None):\n        super().__init__()\n\n        self.model = self._create_model(pretrained_weights_path)\n\n        self.val_map = torchmetrics.MAP()\n        self.val_f2 = KaggleF2()\n\n    def _create_model(self, pretrained_weights_path):\n        \"\"\"Creates finetunable Faster R-CNN model.\n        \n        In the finetuning notebook, the internet can be used and the weights are downloaded.\n        In the inference notebook, there is no internet access and the weights are provided\n        with the pretrained_weights_path arg.\n        \"\"\"\n        if pretrained_weights_path is None:\n            model = fasterrcnn_resnet50_fpn(pretrained=True)\n        else:\n            model = fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False)\n            model.load_state_dict(torch.load(pretrained_weights_path))\n\n        in_features = model.roi_heads.box_predictor.cls_score.in_features\n        model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes=2)\n\n        return model\n\n    def forward(self, image):\n        \"\"\"Runs inference.\"\"\"\n        self.model.eval()\n        output = self.model(image)\n\n        return output\n\n    def training_step(self, batch, batch_idx):\n        image, target = batch\n        loss_dict = self.model(image, target)\n        losses = sum(loss for loss in loss_dict.values())\n\n        batch_size = len(batch[0])\n        self.log_dict(loss_dict, batch_size=batch_size)\n        self.log(\"train_loss\", losses, batch_size=batch_size)\n\n        return losses\n\n    def validation_step(self, batch, batch_idx):\n        image, target = batch\n        output = self.model(image)\n\n        val_map = self.val_map(output, target)\n        val_f2 = self.val_f2(output, target)\n\n        self.log(\"val_map\", val_map[\"map\"])\n        self.log(\"val_f2\", val_f2)\n\n\n    def configure_optimizers(self):\n        params = [p for p in self.model.parameters() if p.requires_grad]\n        optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)\n\n        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)\n\n        return [optimizer], [lr_scheduler]\n","metadata":{"execution":{"iopub.status.busy":"2022-01-03T18:27:26.125784Z","iopub.execute_input":"2022-01-03T18:27:26.126038Z","iopub.status.idle":"2022-01-03T18:27:26.141523Z","shell.execute_reply.started":"2022-01-03T18:27:26.126006Z","shell.execute_reply":"2022-01-03T18:27:26.140839Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Lightning DataModule Class","metadata":{}},{"cell_type":"code","source":"class GBRDataModule(pl.LightningDataModule):\n    \"\"\"LightningDataModule class to split dataset and create dataloaders.\"\"\"\n    def __init__(self, csv_path, batch_size, num_workers):\n        super().__init__()\n\n        self.save_hyperparameters()\n\n        self.csv_path = csv_path\n        self.batch_size = batch_size\n        self.num_workers = num_workers\n\n    def setup(self, stage=None):\n        train_dataset = GBRDataset(self.csv_path)\n        val_dataset = GBRDataset(self.csv_path)\n\n        # Split\n        len_total = len(train_dataset)\n        len_train = int(0.8 * len_total)\n        indices = torch.randperm(len_total).tolist()\n        train_dataset = torch.utils.data.Subset(train_dataset, indices[:len_train])\n        val_dataset = torch.utils.data.Subset(val_dataset, indices[len_train:])\n\n        self.train_dataset, self.val_dataset = train_dataset, val_dataset\n\n    def train_dataloader(self):\n        return self._dataloader(self.train_dataset, shuffle=True)\n\n    def val_dataloader(self):\n        return self._dataloader(self.val_dataset)\n\n    def _dataloader(self, dataset, shuffle=False):\n        return DataLoader(\n            dataset,\n            batch_size=self.batch_size,\n            shuffle=shuffle,\n            num_workers=self.num_workers,\n            collate_fn=collate_fn,\n            pin_memory=True,\n            drop_last=True,\n        )\n\n\ndef collate_fn(batch):\n    return tuple(zip(*batch))\n","metadata":{"execution":{"iopub.status.busy":"2022-01-03T18:27:26.142629Z","iopub.execute_input":"2022-01-03T18:27:26.142884Z","iopub.status.idle":"2022-01-03T18:27:26.154814Z","shell.execute_reply.started":"2022-01-03T18:27:26.142844Z","shell.execute_reply":"2022-01-03T18:27:26.154111Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train Function","metadata":{}},{"cell_type":"code","source":"def train():\n    pl.seed_everything(42, workers=True)\n\n    gbr_module = GBRModule()\n\n    gbr_data_module = GBRDataModule(TRAIN_CSV_PATH, BATCH_SIZE, NUM_WORKERS)\n\n    trainer = pl.Trainer(\n        fast_dev_run=FAST_DEV_RUN,\n        gpus=GPUS,\n        logger=WandbLogger(project=WANDB_PROJECT, log_model=True),\n        max_epochs=MAX_EPOCHS,\n        precision=16 if GPUS else 32,\n    )\n\n    trainer.fit(gbr_module, gbr_data_module)","metadata":{"execution":{"iopub.status.busy":"2022-01-03T18:27:26.155958Z","iopub.execute_input":"2022-01-03T18:27:26.156596Z","iopub.status.idle":"2022-01-03T18:27:26.166976Z","shell.execute_reply.started":"2022-01-03T18:27:26.156560Z","shell.execute_reply":"2022-01-03T18:27:26.166244Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Run the Training","metadata":{}},{"cell_type":"code","source":"train()","metadata":{"execution":{"iopub.status.busy":"2022-01-03T18:29:17.372018Z","iopub.execute_input":"2022-01-03T18:29:17.372269Z","iopub.status.idle":"2022-01-03T18:29:40.924899Z","shell.execute_reply.started":"2022-01-03T18:29:17.372241Z","shell.execute_reply":"2022-01-03T18:29:40.924119Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}