{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Really good training pipeline for pytorch EfficientDet \n\nHi everyone!\n\nMy name is Hieu, I am AI engineer. Especially I am in Love with CV & DL.\n\n\nI hope it is useful for you, my friends! If you didn't read this kernel, don't forget to do it! :)\n\n\nToday I would like to share really good training pipeline for this competition using SOTA [EfficientDet: Scalable and Efficient Object Detection](https://arxiv.org/pdf/1911.09070.pdf)"},{"metadata":{},"cell_type":"markdown","source":"## Main Idea\n \nI read [all public kernels about EfficientDet in kaggle community](https://www.kaggle.com/search?q=efficientdet+in%3Anotebooks) and understand that kaggle don't have really good working public kernels with good score. Why? You can see below picture about COCO AP for different architectures, I think everyone should be able to use such a strong tools EfficientDet for own research, lets do it!\n\n<img src='https://miro.medium.com/max/2400/0*ApAKUWtseHcvRV2U.png' width=400>   "},{"metadata":{},"cell_type":"markdown","source":"**The model only trains with 1 epoch to save GPU resources, I will update the weight to the newest one after I finish training**"},{"metadata":{},"cell_type":"markdown","source":"## Code inference here \nhttps://www.kaggle.com/backtracking/efficientdet-inference"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install timm==0.1.26","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install omegaconf","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import sys\nsys.path.insert(0, \"../input/efficientdet\")\nsys.path.insert(0, \"../input/timm-efficientdet-pytorch\")\nimport re\nimport os\nfrom datetime import datetime\nimport random\nimport time\nimport torch\nimport numpy as np\nimport pandas as pd\nfrom glob import glob\nfrom torch.utils.data import Dataset\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\nfrom torch.utils.data.sampler import SequentialSampler, RandomSampler\nimport cv2\nfrom matplotlib import pyplot as plt\nfrom effdet import get_efficientdet_config, EfficientDet, DetBenchTrain\nfrom effdet.efficientdet import HeadNet\nfrom sklearn.model_selection import StratifiedKFold\nimport pydicom\nimport warnings\nfrom matplotlib.pyplot import figure\nwarnings.filterwarnings('ignore')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"'''https://www.kaggle.com/backtracking/smart-data-split-train-eval-for-object-detection'''\n\nN_CLASSES = 14\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\ndf = pd.read_csv(r'../input/vinbigdata-chest-xray-abnormalities-detection/train.csv')\ndf = pd.DataFrame(df)\ndf = df[df['class_name'] != 'No finding']\n\nskf = StratifiedKFold(n_splits=5, shuffle=True, random_state=101)\ndf_folds = df[['image_id']].copy()\n\ndf_folds.loc[:, 'bbox_count'] = 1\ndf_folds = df_folds.groupby('image_id').count()\ndf_folds.loc[:, 'object_count'] = df.groupby('image_id')['class_id'].nunique()\n\ndf_folds.loc[:, 'stratify_group'] = np.char.add(\n    df_folds['object_count'].values.astype(str),\n    df_folds['bbox_count'].apply(lambda x: f'_{x // 15}').values.astype(str)\n)\n\ndf_folds.loc[:, 'fold'] = 0\nfor fold_number, (train_index, val_index) in enumerate(skf.split(X=df_folds.index, y=df_folds['stratify_group'])):\n    df_folds.loc[df_folds.iloc[val_index].index, 'fold'] = fold_number\n\ndf_folds.reset_index(inplace=True)\nvalid_df = pd.merge(df, df_folds[df_folds['fold'] == 0], on='image_id')\ntrain_df = pd.merge(df, df_folds[df_folds['fold'] != 0], on='image_id')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Data Spliting**"},{"metadata":{"trusted":true},"cell_type":"code","source":"figure(num=None, figsize=(30, 8))\ntrain_df['class_name'].hist()\nvalid_df['class_name'].hist()\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Albumentations"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_transforms():\n    return A.Compose(\n        [\n            A.ToGray(p=0.01),\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.Resize(height=512, width=512, p=1),\n            A.Cutout(num_holes=8, max_h_size=64, max_w_size=64, fill_value=0, p=0.5),\n            ToTensorV2(p=1.0),\n        ],\n        p=1.0,\n        bbox_params=A.BboxParams(\n            format='pascal_voc',\n            min_area=0,\n            min_visibility=0,\n            label_fields=['labels']\n        )\n    )\n\n\ndef get_valid_transforms():\n    return A.Compose(\n        [\n            A.Resize(height=512, width=512, p=1.0),\n            ToTensorV2(p=1.0),\n        ],\n        p=1.0,\n        bbox_params=A.BboxParams(\n            format='pascal_voc',\n            min_area=0,\n            min_visibility=0,\n            label_fields=['labels']\n        )\n    )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"class DatasetRetriever(Dataset):\n\n    def __init__(self, marking, image_dir=None, transforms=None):\n        super().__init__()\n\n        self.image_ids = marking[\"image_id\"].unique()\n        self.marking = marking\n        self.image_dir = image_dir\n        self.transforms = transforms\n\n    def __getitem__(self, index: int):\n        image_id = self.image_ids[index]\n\n        records = self.marking[(self.marking['image_id'] == image_id)]\n        records = records.reset_index(drop=True)\n        image, boxes = self.load_image_and_boxes(index)\n\n        labels = records[\"class_id\"].values\n        target = {'boxes': torch.tensor(boxes), 'labels': torch.tensor(labels), 'image_id': torch.tensor([index])}\n\n        if self.transforms:\n            sample = self.transforms(**{\n                'image': image,\n                'bboxes': target['boxes'],\n                'labels': labels\n            })\n            if len(sample['bboxes']) > 0:\n                image = sample['image']\n                target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)\n                target['boxes'][:, [0, 1, 2, 3]] = target['boxes'][:, [1, 0, 3, 2]]\n        return image, target, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n\n    def load_image_and_boxes(self, index):\n        image_id = self.image_ids[index]\n        dicom = pydicom.dcmread(f\"{self.image_dir}/{image_id}.dicom\")\n        image = dicom.pixel_array\n        if \"PhotometricInterpretation\" in dicom:\n            if dicom.PhotometricInterpretation == \"MONOCHROME1\":\n                image = np.amax(image) - image\n\n        image = np.stack([image, image, image])\n        image = image.astype('float32')\n        image = image - image.min()\n        image = image / image.max()\n        image = image.transpose(1, 2, 0)\n\n        records = self.marking[self.marking['image_id'] == image_id]\n        boxes = records[['x_min', 'y_min', 'x_max', 'y_max']].values\n\n        return image, boxes","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"image_paths = '../input/vinbigdata-chest-xray-abnormalities-detection/train'\ntrain_dataset = DatasetRetriever(\n    marking=train_df,\n    image_dir=image_paths,\n    transforms=get_train_transforms(),\n)\n\nvalidation_dataset = DatasetRetriever(\n    marking=valid_df,\n    image_dir=image_paths,\n    transforms=get_valid_transforms(),\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Visualization**"},{"metadata":{"trusted":true},"cell_type":"code","source":"image, target, image_id = train_dataset[0]\nboxes = target['boxes'].cpu().numpy().astype(np.int64)\n\nnumpy_image = image.permute(1, 2, 0).cpu().numpy()\n\nfig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\nfor box in boxes:\n    cv2.rectangle(numpy_image, (box[1], box[0]), (box[3], box[2]), (0, 1, 0), 2)\n\nax.set_axis_off()\nax.imshow(numpy_image)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## AverageMeter"},{"metadata":{"trusted":true},"cell_type":"code","source":"class AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Filter"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Fitter:\n\n    def __init__(self, model, device, config):\n        self.config = config\n        self.epoch = 0\n\n        self.base_dir = f'./{config.folder}'\n        if not os.path.exists(self.base_dir):\n            os.makedirs(self.base_dir)\n\n        self.log_path = f'{self.base_dir}/log.txt'\n        self.best_summary_loss = 10 ** 5\n\n        self.model = model\n        self.device = device\n\n        param_optimizer = list(self.model.named_parameters())\n        no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n        optimizer_grouped_parameters = [\n            {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.001},\n            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}\n        ]\n\n        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=config.lr)\n        self.scheduler = config.SchedulerClass(self.optimizer, **config.scheduler_params)\n        self.log(f'Fitter prepared. Device is {self.device}')\n\n    def fit(self, train_loader, validation_loader):\n        for e in range(self.config.n_epochs):\n            if self.config.verbose:\n                lr = self.optimizer.param_groups[0]['lr']\n                timestamp = datetime.utcnow().isoformat()\n                self.log(f'\\n{timestamp}\\nLR: {lr}')\n\n            t = time.time()\n            summary_loss = self.train_one_epoch(train_loader)\n\n            self.log(\n                f'[RESULT]: Train. Epoch: {self.epoch}, summary_loss: {summary_loss.avg:.5f}, time: {(time.time() - t):.5f}')\n            self.save(f'{self.base_dir}/last-checkpoint.bin')\n\n            t = time.time()\n            summary_loss = self.validation(validation_loader)\n\n            self.log(\n                f'[RESULT]: Val. Epoch: {self.epoch}, summary_loss: {summary_loss.avg:.5f}, time: {(time.time() - t):.5f}')\n            if summary_loss.avg < self.best_summary_loss:\n                self.best_summary_loss = summary_loss.avg\n                self.model.eval()\n                self.save(f'{self.base_dir}/best-checkpoint-{str(self.epoch).zfill(3)}epoch.bin')\n                for path in sorted(glob(f'{self.base_dir}/best-checkpoint-*epoch.bin'))[:-3]:\n                    os.remove(path)\n\n            if self.config.validation_scheduler:\n                self.scheduler.step(metrics=summary_loss.avg)\n\n            self.epoch += 1\n\n    def validation(self, val_loader):\n        self.model.eval()\n        summary_loss = AverageMeter()\n        t = time.time()\n        for step, (images, targets, image_ids) in enumerate(val_loader):\n            if self.config.verbose:\n                if step % self.config.verbose_step == 0:\n                    print(\n                        f'Val Step {step}/{len(val_loader)}, ' + \\\n                        f'summary_loss: {summary_loss.avg:.5f}, ' + \\\n                        f'time: {(time.time() - t):.5f}', end='\\r'\n                    )\n            with torch.no_grad():\n                images = torch.stack(images)\n                batch_size = images.shape[0]\n                images = images.to(self.device).float()\n                boxes = [target['boxes'].to(self.device).float() for target in targets]\n                labels = [target['labels'].to(self.device).float() for target in targets]\n\n                loss, _, _ = self.model(images, boxes, labels)\n                summary_loss.update(loss.detach().item(), batch_size)\n\n        return summary_loss\n\n    def train_one_epoch(self, train_loader):\n        self.model.train()\n        summary_loss = AverageMeter()\n        t = time.time()\n        for step, (images, targets, image_ids) in enumerate(train_loader):\n            if self.config.verbose:\n                if step % self.config.verbose_step == 0:\n                    print(\n                        f'Train Step {step}/{len(train_loader)}, ' + \\\n                        f'summary_loss: {summary_loss.avg:.5f}, ' + \\\n                        f'time: {(time.time() - t):.5f}', end='\\r'\n                    )\n            images = torch.stack(images)\n            images = images.to(self.device).float()\n            batch_size = images.shape[0]\n            boxes = [target['boxes'].to(self.device).float() for target in targets]\n            labels = [target['labels'].to(self.device).float() for target in targets]\n            self.optimizer.zero_grad()\n            loss, _, _ = self.model(images, boxes, labels)\n\n            loss.backward()\n\n            summary_loss.update(loss.detach().item(), batch_size)\n\n            self.optimizer.step()\n\n            if self.config.step_scheduler:\n                self.scheduler.step()\n\n        return summary_loss\n\n    def save(self, path):\n        self.model.eval()\n        torch.save({\n            'model_state_dict': self.model.model.state_dict(),\n            'optimizer_state_dict': self.optimizer.state_dict(),\n            'scheduler_state_dict': self.scheduler.state_dict(),\n            'best_summary_loss': self.best_summary_loss,\n            'epoch': self.epoch,\n        }, path)\n\n    def load(self, path):\n        checkpoint = torch.load(path)\n        self.model.model.load_state_dict(checkpoint['model_state_dict'])\n        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n        self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])\n        self.best_summary_loss = checkpoint['best_summary_loss']\n        self.epoch = checkpoint['epoch'] + 1\n\n    def log(self, message):\n        if self.config.verbose:\n            print(message)\n        with open(self.log_path, 'a+') as logger:\n            logger.write(f'{message}\\n')\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Config for Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"class TrainGlobalConfig:\n    num_workers = 1\n    batch_size = 4\n    folder = 'logging_path'\n    n_epochs = 1\n    lr = 0.0002\n    verbose = True\n    verbose_step = 1\n    step_scheduler = False\n    validation_scheduler = True\n\n    SchedulerClass = torch.optim.lr_scheduler.ReduceLROnPlateau\n    scheduler_params = dict(\n        mode='min',\n        factor=0.5,\n        patience=1,\n        verbose=False,\n        threshold=0.0001,\n        threshold_mode='abs',\n        cooldown=0,\n        min_lr=1e-8,\n        eps=1e-08\n    )\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def run_training():\n    net.to(device)\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=TrainGlobalConfig.batch_size,\n        sampler=RandomSampler(train_dataset),\n        pin_memory=False,\n        drop_last=True,\n        num_workers=TrainGlobalConfig.num_workers,\n        collate_fn=collate_fn,\n    )\n    val_loader = torch.utils.data.DataLoader(\n        validation_dataset,\n        batch_size=TrainGlobalConfig.batch_size,\n        num_workers=TrainGlobalConfig.num_workers,\n        shuffle=False,\n        sampler=SequentialSampler(validation_dataset),\n        pin_memory=False,\n        collate_fn=collate_fn,\n    )\n\n    fitter = Fitter(model=net, device=device, config=TrainGlobalConfig)\n    fitter.fit(train_loader, val_loader)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_net():\n    config = get_efficientdet_config(model_name='tf_efficientdet_d4')\n    net = EfficientDet(config, pretrained_backbone=False)\n    checkpoint = torch.load('../input/efficientdet-weights/efficientdet_d4-5b370b7a.pth')\n    net.load_state_dict(checkpoint)\n    config.num_classes = 14\n    config.image_size = 512\n    net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))\n    return DetBenchTrain(net, config)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"net = get_net()\nrun_training()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Upvote for me if it helpful :))"}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}