{"cells":[{"metadata":{},"cell_type":"markdown","source":"# EfficientDet - Training","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"This notebook is based on Alex Shonenkov's [Training EfficientDet](https://www.kaggle.com/shonenkov/training-efficientdet) notebook.  \n**EfficientDet** is an object detection model developed at Google.  Paper: [EfficientDet: Scalable and Efficient Object Detection](https://arxiv.org/pdf/1911.09070.pdf).\n\nRequired Data for this Notebook:\n1. `efficientdet-pytorch` - EfficientDet, object detection model.\n2. `timm`, or `pytorch-image-models` - required by `efficientdet-pytorch`.\n3. `omegaconf` - required by `efficientdet-pytorch`","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"! pip install --no-deps '/kaggle/input/timm-package/timm-0.1.26-py3-none-any.whl'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import sys\nsys.path.insert(0, \"/kaggle/input/timm-efficientdet-pytorch\")\nsys.path.insert(0, \"/kaggle/input/omegaconf\")\nimport os, time, random\nfrom pathlib import Path\nfrom glob import glob\nfrom datetime import datetime\nimport numpy as np, pandas as pd\nimport PIL\nimport matplotlib.pyplot as plt\nimport numba\nfrom numba import jit\nfrom sklearn.model_selection import StratifiedKFold\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.utils.data.sampler import SequentialSampler, RandomSampler\nimport cv2\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\nfrom effdet import get_efficientdet_config, EfficientDet, DetBenchTrain, DetBenchEval\nfrom effdet.efficientdet import HeadNet\n\nPath.ls = lambda x: list(x.iterdir())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"SZ = 512","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Data","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"marking = pd.read_csv('/kaggle/input/global-wheat-detection/train.csv')\nbboxs = np.stack(marking.bbox.apply(lambda r: np.fromstring(r[1:-1], sep=',')))\nfor i, column in enumerate(['x', 'y', 'w', 'h']): marking[column] = bboxs[:,i]\nmarking.drop(columns='bbox', inplace=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)\ndf_folds = marking[['image_id']].copy()\ndf_folds.loc[:,'bbox_count'] = 1\ndf_folds = df_folds.groupby('image_id').count()\ndf_folds.loc[:,'source'] = marking[['image_id', 'source']].groupby('image_id').min()['source']\ndf_folds.loc[:,'stratify_group'] = np.char.add(\n    df_folds['source'].values.astype(str),\n    df_folds['bbox_count'].apply(lambda x: f'_{x // 15}').values.astype(str))\ndf_folds.loc[:,'fold'] = 0\nfor fold_number, (train_index, valid_index) in enumerate(skf.split(X=df_folds.index, y=df_folds['stratify_group'])):\n    df_folds.loc[df_folds.iloc[valid_index].index,'fold'] = fold_number","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_transforms():\n    return A.Compose(\n        [A.RandomSizedCrop(min_max_height=(800, 800), height=1024, width=1024, p=0.5),\n         A.OneOf([A.HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.9),\n                  A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.9)],\n                 p=0.9),\n         A.ToGray(p=0.01),\n         A.HorizontalFlip(p=0.5),\n         A.VerticalFlip(p=0.5),\n         A.Resize(height=SZ, width=SZ, p=1),\n#          A.Cutout(num_holes=8, max_h_size=64, max_w_size=64, fill_value=0, p=.5), \n         A.Cutout(num_holes=80, max_h_size=SZ//20, max_w_size=SZ//20, fill_value=0, p=.5),\n#          A.Blur(p=.5),\n         ToTensorV2(p=1)],\n        p=1,\n        bbox_params=A.BboxParams(format='pascal_voc', min_area=0, min_visibility=0, label_fields=['labels'])\n    )\n\ndef get_valid_transforms():\n    return A.Compose(\n        [A.Resize(height=SZ, width=SZ, p=1), ToTensorV2(p=1)], \n        p=1,\n        bbox_params=A.BboxParams(format='pascal_voc', min_area=0, min_visibility=0, label_fields=['labels']))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TRAIN_ROOT_PATH = '../input/global-wheat-detection/train/'\n\nclass DatasetRetriever(Dataset):\n    def __init__(self, marking, image_ids, transforms=None, test=False):\n        super().__init__()\n        self.marking, self.image_ids = marking, image_ids\n        self.transforms, self.test = transforms, test\n        \n    def __getitem__(self, index:int): \n        image_id = self.image_ids[index]\n        if self.test or random.random() > 0.5: \n            image, boxes = self.load_image_and_boxes(index)\n        else: \n            image, boxes = self.load_cutmix_image_and_boxes(index)\n        labels = torch.ones((len(boxes),), dtype=torch.int64)\n        if self.transforms:\n            for i in range(10):\n                sample = self.transforms(image=image, bboxes=boxes, labels=labels)\n                if len(sample['bboxes']) > 0:\n                    image = sample['image']\n                    boxes = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)\n                    boxes[:,[0,1,2,3]] = boxes[:,[1,0,3,2]]\n                    break\n        labels = torch.ones((len(boxes),), dtype=torch.int64)\n        target = {}\n        target['boxes'] = boxes\n        target['labels'] = labels\n        target['image_id'] = torch.tensor([index])\n        return image, target, image_id\n        \n    def __len__(self) -> int: return len(self.image_ids)\n    \n    def load_image_and_boxes(self, index):\n        image_id = self.image_ids[index]\n        image = cv2.imread(f'{TRAIN_ROOT_PATH}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.\n        records = self.marking[self.marking['image_id']==image_id]\n        boxes = records[['x', 'y', 'w', 'h']].values\n        boxes[:,2] = boxes[:,0] + boxes[:,2]\n        boxes[:,3] = boxes[:,1] + boxes[:,3]\n        return image, boxes\n        \n    def load_cutmix_image_and_boxes(self, index, imsize=1024):\n        w, h = imsize, imsize\n        s = imsize // 2\n        xc, yc = [int(random.uniform(.25*imsize, .75*imsize)) for _ in range(2)]\n        indexes = [index] + [random.randint(0, len(self.image_ids)-1) for _ in range(3)]\n        result_image = np.full((imsize, imsize, 3), 1, dtype=np.float32)\n        result_boxes = []\n        for i, index in enumerate(indexes):\n            image, boxes = self.load_image_and_boxes(index)\n            if i == 0:\n                x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (large image)\n                x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h  # xmin, ymin, xmax, ymax (small image)\n            elif i == 1:  # top right\n                x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc\n                x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h\n            elif i == 2:  # bottom left\n                x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)\n                x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, max(xc, w), min(y2a - y1a, h)\n            elif i == 3:  # bottom right\n                x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)\n                x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)\n            result_image[y1a:y2a,x1a:x2a] = image[y1b:y2b,x1b:x2b]\n            padw, padh = x1a - x1b, y1a - y1b\n            boxes[:,0] += padw; boxes[:,1] += padh\n            boxes[:,2] += padw; boxes[:,3] += padh\n            result_boxes.append(boxes)\n        result_boxes = np.concatenate(result_boxes, 0)\n        np.clip(result_boxes[:,0:], 0, 2*s, out=result_boxes[:,0:])\n        result_boxes = result_boxes.astype(np.int32)\n        result_boxes = result_boxes[np.where((result_boxes[:,2]-result_boxes[:,0]) * (result_boxes[:,3]-result_boxes[:,1]) > 0)]\n        return result_image, result_boxes            ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fold_number = 3\n\ntrain_dataset = DatasetRetriever(\n    image_ids=df_folds[df_folds.fold!=fold_number].index.values, \n    marking=marking, transforms=get_train_transforms(), test=False)\n\nvalidation_dataset = DatasetRetriever(\n    image_ids = df_folds[df_folds.fold==fold_number].index.values, \n    marking=marking, transforms=get_valid_transforms(), test=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image, target, image_id = train_dataset[8]\nboxes = target['boxes'].cpu().numpy().astype(np.int32)\nassert len(target['boxes']) == len(target['labels'])\nnumpy_image = image.permute(1, 2, 0).cpu().numpy()\n\nfig, ax = plt.subplots(1, 1, figsize=(16, 8))\nfor box in boxes:\n    cv2.rectangle(numpy_image, (box[1], box[0]), (box[3], box[2]), (0, 1, 0), 2)\nax.set_axis_off()\nax.imshow(numpy_image);","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Data:\n    def __init__(self, train_dl, valid_dl):\n        self.train_dl, self.valid_dl = train_dl, valid_dl","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"bs = 4","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch): return tuple(zip(*batch))\n\ntrain_loader = torch.utils.data.DataLoader(\n    train_dataset, \n    batch_size=bs,\n    sampler=RandomSampler(train_dataset),\n    pin_memory=False,\n    drop_last=True,\n    num_workers=2,\n    collate_fn=collate_fn)\n\nval_loader = torch.utils.data.DataLoader(\n    validation_dataset,\n    batch_size=bs,\n    num_workers=2,\n    shuffle=False,\n    sampler=SequentialSampler(validation_dataset),\n    pin_memory=False,\n    collate_fn=collate_fn)\n\ndata = Data(train_loader, val_loader)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_benches():\n    config = get_efficientdet_config('tf_efficientdet_d7')\n    net = EfficientDet(config, pretrained_backbone=False)\n    checkpoint = torch.load('/kaggle/input/efficientdet/efficientdet_d7-f05bf714.pth')\n#     checkpoint = torch.load('/kaggle/input/efficientdet/efficientdet_d6-51cb0132.pth')\n#     checkpoint = torch.load('/kaggle/input/efficientdet/efficientdet_d5-ef44aea8.pth')\n    net.load_state_dict(checkpoint)\n    config.num_classes = 1\n    config.image_size = SZ\n    net.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=dict(eps=.001, momentum=.01))\n    return DetBenchTrain(net, config), DetBenchEval(net, config)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Metric: precision","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"@jit(nopython=True)\ndef calculate_iou(gt, pr, form='pascal_voc') -> float:\n    \"\"\"Calculates the Intersection over Union.\n\n    Args:\n        gt: (np.ndarray[Union[int, float]]) coordinates of the ground-truth box\n        pr: (np.ndarray[Union[int, float]]) coordinates of the prdected box\n        form: (str) gt/pred coordinates format\n            - pascal_voc: [xmin, ymin, xmax, ymax]\n            - coco: [xmin, ymin, w, h]\n    Returns:\n        (float) Intersection over union (0.0 <= iou <= 1.0)\n    \"\"\"\n    if form == 'coco':\n        gt = gt.copy()\n        pr = pr.copy()\n\n        gt[2] = gt[0] + gt[2]\n        gt[3] = gt[1] + gt[3]\n        pr[2] = pr[0] + pr[2]\n        pr[3] = pr[1] + pr[3]\n\n    # Calculate overlap area\n    dx = min(gt[2], pr[2]) - max(gt[0], pr[0]) + 1\n    if dx < 0: return 0.0\n\n    dy = min(gt[3], pr[3]) - max(gt[1], pr[1]) + 1\n    if dy < 0: return 0.0\n\n    overlap_area = dx * dy\n    # Calculate union area\n    union_area = (\n            (gt[2] - gt[0] + 1) * (gt[3] - gt[1] + 1) +\n            (pr[2] - pr[0] + 1) * (pr[3] - pr[1] + 1) -\n            overlap_area\n    )\n    return overlap_area / union_area\n\n@jit(nopython=True)\ndef find_best_match(gts, pred, pred_idx, threshold = 0.5, form = 'pascal_voc', ious=None) -> int:\n    \"\"\"Returns the index of the 'best match' between the\n    ground-truth boxes and the prediction. The 'best match'\n    is the highest IoU. (0.0 IoUs are ignored).\n\n    Args:\n        gts: (List[List[Union[int, float]]]) Coordinates of the available ground-truth boxes\n        pred: (List[Union[int, float]]) Coordinates of the predicted box\n        pred_idx: (int) Index of the current predicted box\n        threshold: (float) Threshold\n        form: (str) Format of the coordinates\n        ious: (np.ndarray) len(gts) x len(preds) matrix for storing calculated ious.\n\n    Return:\n        (int) Index of the best match GT box (-1 if no match above threshold)\n    \"\"\"\n    best_match_iou = -np.inf\n    best_match_idx = -1\n\n    for gt_idx in range(len(gts)):\n        \n        if gts[gt_idx][0] < 0:\n            # Already matched GT-box\n            continue\n        \n        iou = -1 if ious is None else ious[gt_idx][pred_idx]\n\n        if iou < 0:\n            iou = calculate_iou(gts[gt_idx], pred, form=form)\n            \n            if ious is not None:\n                ious[gt_idx][pred_idx] = iou\n\n        if iou < threshold:\n            continue\n\n        if iou > best_match_iou:\n            best_match_iou = iou\n            best_match_idx = gt_idx\n    return best_match_idx\n\n@jit(nopython=True)\ndef calculate_precision(gts, preds, threshold = 0.5, form = 'coco', ious=None) -> float:\n    \"\"\"Calculates precision for GT - prediction pairs at one threshold.\n\n    Args:\n        gts: (List[List[Union[int, float]]]) Coordinates of the available ground-truth boxes\n        preds: (List[List[Union[int, float]]]) Coordinates of the predicted boxes,\n               sorted by confidence value (descending)\n        threshold: (float) Threshold\n        form: (str) Format of the coordinates\n        ious: (np.ndarray) len(gts) x len(preds) matrix for storing calculated ious.\n\n    Return:\n        (float) Precision\n    \"\"\"\n    n = len(preds)\n    tp = 0\n    fp = 0\n    \n    # for pred_idx, pred in enumerate(preds_sorted):\n    for pred_idx in range(n):\n        best_match_gt_idx = find_best_match(gts, preds[pred_idx], pred_idx,\n                                            threshold=threshold, form=form, ious=ious)\n\n        if best_match_gt_idx >= 0:\n            # True positive: The predicted box matches a gt box with an IoU above the threshold.\n            tp += 1\n            # Remove the matched GT box\n            gts[best_match_gt_idx] = -1\n\n        else:\n            # No match\n            # False positive: indicates a predicted box had no associated gt box.\n            fp += 1\n\n    # False negative: indicates a gt box had no associated predicted box.\n    fn = (gts.sum(axis=1) > 0).sum()\n    return tp / (tp + fp + fn)\n\n\n@jit(nopython=True)\ndef calculate_image_precision(gts, preds, thresholds = (0.5, ), form = 'coco') -> float:\n    \"\"\"Calculates image precision.\n\n    Args:\n        gts: (List[List[Union[int, float]]]) Coordinates of the available ground-truth boxes\n        preds: (List[List[Union[int, float]]]) Coordinates of the predicted boxes,\n               sorted by confidence value (descending)\n        thresholds: (float) Different thresholds\n        form: (str) Format of the coordinates\n\n    Return:\n        (float) Precision\n    \"\"\"\n    n_threshold = len(thresholds)\n    image_precision = 0.0\n    \n    ious = np.ones((len(gts), len(preds))) * -1\n    # ious = None\n\n    for threshold in thresholds:\n        precision_at_threshold = calculate_precision(gts.copy(), preds, threshold=threshold,\n                                                     form=form, ious=ious)\n        image_precision += precision_at_threshold / n_threshold\n\n    return image_precision\n\n# Numba typed list!\niou_thresholds = numba.typed.List()\nfor x in [0.5, 0.55, 0.6, 0.65, 0.7, 0.75]:\n    iou_thresholds.append(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def single_image_score(prediction, score_threshold):\n    '''\n    Returns the precision for a single image.\n    '''\n    gt_boxes = prediction['gt_boxes'].copy()\n    pred_boxes = prediction['pred_boxes'].copy()\n    scores = prediction['scores'].copy()\n    image_id = prediction['image_id']\n    indexes = np.where(scores > score_threshold)\n    pred_boxes = pred_boxes[indexes]\n    scores = scores[indexes]\n    image_precision = calculate_image_precision(gt_boxes, pred_boxes, thresholds=iou_thresholds, form='pascal_voc')\n    return image_precision\n\ndef calculate_final_score(all_predictions, score_threshold):\n    '''\n    Returns the average precision for multiple images.\n    '''\n    final_scores = []\n    for i in range(len(all_predictions)):\n        image_precision = single_image_score(all_predictions[i], score_threshold)\n        final_scores.append(image_precision)\n    return np.mean(final_scores)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Performance recorder","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class AverageMeter:\n    def __init__(self): self.reset()\n        \n    def reset(self): \n        self.val, self.avg, self.sum, self.count = 0, 0, 0, 0\n        \n    def update(self, val, n=1):\n        self.val = val\n        self.sum += n*val\n        self.count += n\n        self.avg = self.sum / self.count","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Fitter","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def mixup_things(imgs, boxss, labelss, alpha=3.):\n    '''\n    Take a batch of samples and convert it into a batch of mixup samples.\n    '''\n    assert len(imgs) == len(boxss) == len(labelss)\n    for i in range(len(imgs)):\n        assert len(boxss[i]) == len(labelss[i])\n\n    beta_dist = torch.distributions.Beta(alpha, alpha)\n    \n    batch_size = len(imgs)\n    idxs = torch.arange(batch_size)\n    idxs_shifted = torch.cat([torch.arange(1, batch_size), torch.arange(1)])\n\n    imgs_mixup, boxss_mixup, labelss_mixup = [], [], []\n    for i, j in zip(idxs, idxs_shifted):\n        if i == j:\n            img, boxs, labels = imgs[i], boxss[i], labelss[i]\n        else:\n            w = beta_dist.sample()\n            img = torch.stack([w * imgs[i], (1 - w) * imgs[j]]).sum(dim=0)\n            boxs = torch.cat([boxss[i], boxss[j]], dim=0)\n            labels = torch.cat([labelss[i], labelss[j]], dim=0)\n        imgs_mixup.append(img)\n        boxss_mixup.append(boxs)\n        labelss_mixup.append(labels)\n    imgs_mixup = torch.stack(imgs_mixup, dim=0)\n\n    assert len(imgs_mixup) == len(boxss_mixup) == len(labelss_mixup)\n    for i in range(len(imgs_mixup)):\n        assert len(boxss_mixup[i]) == len(labelss_mixup[i])\n        \n    return imgs_mixup, boxss_mixup, labelss_mixup","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import warnings\nwarnings.filterwarnings('ignore')\n\nclass StopTrainException(Exception): pass\n\ndef organize_prediction(image_id, target, det):\n    boxes = det.detach().cpu().numpy()[:,:4]\n    scores = det.detach().cpu().numpy()[:,4]\n    boxes[:,2] = boxes[:,2] + boxes[:,0]\n    boxes[:,3] = boxes[:,3] + boxes[:,1]\n    pred = {'pred_boxes':(2 * boxes).clip(min=0, max=1023).astype(int),\n            'scores':scores,\n            'gt_boxes':(2 * target['boxes'].cpu().numpy()).clip(min=0, max=1023).astype(int)[:,[1,0,3,2]],  # Need to change yxyx to xyxy for gt_boxes for metric calculation\n            'image_id':image_id}\n    return pred\n\n\nclass Fitter:\n    def __init__(self, data, train_bench, eval_bench, device, base_dir='./', \n                 opt_func=torch.optim.AdamW, scheduler_func=torch.optim.lr_scheduler.OneCycleLR,\n                 mixup = False,\n                 verbose=True, verbose_step=True): \n        self.data, self.train_bench, self.eval_bench = data, train_bench, eval_bench\n        self.opt_func, self.scheduler_func = opt_func, scheduler_func\n        self.base_dir, self.device = base_dir, device\n        self.mixup = mixup\n        self.verbose, self.verbose_step = verbose, verbose_step\n        if not os.path.exists(self.base_dir): os.makedirs(self.base_dir)\n        self.log_path = f'{self.base_dir}/log.txt'\n        self.best_summary_score = 0.\n        self.freeze(-1)\n        self.log(f'Fitter prepared.  Device is {self.device}.')\n    \n    def _get_param_groups(self):\n        return [{'params':self.train_bench.model.backbone.parameters()},\n                {'params':self.train_bench.model.class_net.parameters()}]\n\n    def freeze(self, idx_grp=None):\n        if idx_grp is None:\n            for p in self.train_bench.model.parameters(): \n                p.requires_grad = False\n            return\n        gs = self._get_param_groups()\n        for g in gs[:idx_grp]: \n            for p in g['params']: \n                p.requires_grad = False\n        for g in gs[idx_grp:]:\n            for p in g['params']: \n                p.requires_grad = True\n    \n    def unfreeze(self):\n        for p in self.train_bench.model.parameters(): \n            p.requires_grad = True\n    \n    def _get_batch_model_inputs(self, images, targets):\n        images = torch.stack(images)\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        return images, boxes, labels\n    \n    def lr_find(self, min_lr=1e-7, max_lr=1e-2):\n        self.train_bench.eval()\n        torch.save(self.train_bench.model.state_dict(), f'{self.base_dir}/tmp_lr_find.pth')\n        self.train_bench.train()\n        \n        opt = self.opt_func(self._get_param_groups(), lr=min_lr)\n        n_iter = len(self.data.train_dl)\n        loss_min = 1e9\n        lrs, losses = [], []\n        for i, (images, targets, image_ids) in enumerate(self.data.train_dl):\n            pos = i / n_iter\n            lr = min_lr * (max_lr / min_lr)**pos\n            print(f'lr = {lr:.5e}', end='\\r')\n            for g in opt.param_groups: g['lr'] += lr\n            images, boxes, labels = self._get_batch_model_inputs(images, targets)\n            opt.zero_grad()\n            loss, _, _ = self.train_bench(images, boxes, labels)\n            loss.backward()\n            opt.step()\n            loss = loss.item() / len(images)\n            lrs.append(opt.param_groups[0]['lr']); losses.append(loss)\n            if loss > 10 * loss_min: \n                break\n            else: \n                loss_min = loss\n        self.train_bench.model.load_state_dict(torch.load(f'{self.base_dir}/tmp_lr_find.pth'))\n        self.lrs, self.losses = np.array(lrs), np.array(losses)\n        return self.lrs, self.losses\n    \n    def plt_lr_find_results(self, skip_end=1):\n        assert skip_end >=0\n        if skip_end==0: lrs, losses = self.lrs[:], self.losses[:]\n        else: lrs, losses = self.lrs[:-skip_end], self.losses[:-skip_end]\n        plt.semilogx(lrs, losses);\n\n    def _display_progress(self, dl, step=0, process_name='train', start_time=0, \n                          meter=None, metric_name='loss'):\n        process_name = process_name[0].upper() + process_name[1:]\n        if self.verbose:\n            if step % self.verbose_step == 0:\n                print(f'{process_name} Step {step}/{len(dl)}, ' + \n                      (f'{metric_name}: {meter.avg:.5f}, ' if (meter and metric_name) else '') + \n                      f'time: {time.time() - start_time:.5f}', end='\\r')\n        \n    def fit(self, n_epochs=1, max_lr=1e-3):\n        self.epoch, self.train_losses, self.valid_losses, self.valid_scores = 0, [], [], []\n        self.optimizer = self.opt_func(self._get_param_groups())\n        if isinstance(max_lr, int): max_lr = [max_lr / 10, max_lr]\n        self.scheduler = self.scheduler_func(\n            self.optimizer, max_lr=max_lr, epochs=n_epochs, steps_per_epoch=len(self.data.train_dl))\n        train_loader, validation_loader = self.data.train_dl, self.data.valid_dl\n        for e in range(n_epochs):\n            if self.verbose:\n                lr = self.optimizer.param_groups[0]['lr']\n                timestamp = datetime.utcnow().isoformat()\n                self.log(f'\\n{timestamp}\\nLR: {lr}')\n            t = time.time()\n            summary_loss = self.train_one_epoch(train_loader)\n            self.train_losses.append(summary_loss.avg)\n            self.log(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            summary_score = self.calculate_metric(validation_loader)\n            self.valid_losses.append(summary_loss.avg)\n            self.valid_scores.append(summary_score.avg)\n            self.log(f'[RESULT]: Val. Epoch: {self.epoch}, summary_loss: {summary_loss.avg:.5f}, summary_score: {summary_score.avg:.5f} time: {time.time() - t:.5f}')\n            if summary_score.avg > self.best_summary_score:\n                self.best_summary_score = summary_score.avg\n                self.train_bench.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]: os.remove(path)\n            self.epoch += 1\n    \n    def single_image_metrics(self, val_loader, score_threshold=0.37):\n        self.eval_bench.eval()\n        preds = []\n        t = time.time()\n        for step, (images, targets, image_ids) in enumerate(val_loader):\n            self._display_progress(dl=val_loader, step=step, process_name='metric', start_time=t)\n            with torch.no_grad():\n                images, boxes, labels = self._get_batch_model_inputs(images, targets)\n                dets = self.eval_bench(images, torch.tensor(len(images) * [1]).float().to(self.device))\n                for i in range(len(images)):\n                    pred = organize_prediction(image_ids[i], targets[i], dets[i])\n                    pred.update({'metric':single_image_score(pred, score_threshold)})\n                    preds.append(pred)\n        return preds\n    \n    def calculate_metric(self, val_loader, score_threshold=0.37):\n        self.eval_bench.eval()\n        summary_score = AverageMeter()\n        t = time.time()\n        for step, (images, targets, image_ids) in enumerate(val_loader):\n            self._display_progress(dl=val_loader, step=step, process_name='metric', start_time=t,\n                                   meter=summary_score, metric_name='score')\n            with torch.no_grad():\n                images, boxes, labels = self._get_batch_model_inputs(images, targets)\n                dets = self.eval_bench(images, torch.tensor(len(images) * [1]).float().to(self.device))\n                preds = []\n                for i in range(len(images)):\n                    preds.append(organize_prediction(image_ids[i], targets[i], dets[i]))\n                score = calculate_final_score(preds, score_threshold)\n                summary_score.update(score, len(images))\n        return summary_score\n    \n    def validation(self, val_loader):\n        self.train_bench.eval()\n        summary_loss = AverageMeter()    \n        t = time.time()\n        for step, (images, targets, image_ids) in enumerate(val_loader):\n            self._display_progress(dl=val_loader, step=step, process_name='Val', start_time=t,\n                                   meter=summary_loss, metric_name='valid_loss')\n            with torch.no_grad():\n                images, boxes, labels = self._get_batch_model_inputs(images, targets)\n                loss, _, _ = self.train_bench(images, boxes, labels)\n                summary_loss.update(loss.detach().item(), len(images))\n        return summary_loss\n    \n    def train_one_epoch(self, train_loader):\n        self.train_bench.train()\n        summary_loss = AverageMeter()\n        t = time.time()\n        for step, (images, targets, image_ids) in enumerate(train_loader):\n            self._display_progress(dl=train_loader, step=step, meter=summary_loss, start_time=t, \n                                   process_name='train', metric_name='train_loss')\n            images, boxes, labels = self._get_batch_model_inputs(images, targets)\n            if self.mixup:\n                images, boxes, labels = mixup_things(images, boxes, labels)\n            self.optimizer.zero_grad()\n            loss, _, _ = self.train_bench(images, boxes, labels)\n            loss.backward()\n            summary_loss.update(loss.detach().item(), len(images))\n            self.optimizer.step()\n            self.scheduler.step()\n        return summary_loss\n        \n    def save(self, path):\n        self.train_bench.eval()\n        torch.save({'model_state_dict':self.train_bench.model.state_dict(),\n                    'optimizer_state_dict':self.optimizer.state_dict(),\n                    'scheduler_state_dict':self.scheduler.state_dict(),\n                    'best_summary_score':self.best_summary_score,\n                    'epoch':self.epoch}, path)\n        \n    def load(self, path):\n        checkpoint = torch.load(path)\n        self.train_bench.model.load_state_dict(checkpoint['model_state_dict'])\n        self.optimizer = self.opt_func(self._get_param_groups())\n        self.scheduler = self.scheduler_func(\n            self.optimizer, max_lr=[1e-4, 1e-3], epochs=1, steps_per_epoch=len(self.data.train_dl))\n        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n        self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])\n        self.best_summary_score = checkpoint['best_summary_score']\n        self.epoch = checkpoint['epoch'] + 1\n        \n    def log(self, message):\n        if self.verbose: print(message)\n        with open(self.log_path, 'a+') as logger:\n            logger.write(f'{message}\\n')\n            \n    def plt_losses(self):\n        _, ax = plt.subplots()\n        ax.plot(self.train_losses, label='train')\n        ax.plot(self.valid_losses, label='valid')\n        ax.legend()\n        ax.set_xlabel('epoch')\n        ax.set_ylabel('loss')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_fitter():\n    device = torch.device('cuda:0') if torch.cuda.is_available() else 'cpu'\n    train_bench, eval_bench = get_benches()\n    train_bench.to(device); eval_bench.to(device)\n    return Fitter(data, train_bench, eval_bench, device=device, mixup=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fitter = get_fitter()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def plt_img_boxs(img, boxs=None, ax=None, figsize=(6, 6)):\n    if ax is None:\n        fig, ax = plt.subplots(figsize=figsize)\n    if boxs is not None:\n        for b in boxs:\n            cv2.rectangle(img, (b[1], b[0]),(b[3], b[2]), (0, 1, 0), 2)\n    ax.imshow(img)\n    ax.axis('off')\n    return fig, ax","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# del imgs, boxss, labelss, targs, imgids","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# imgs, targs, imgids = next(iter(fitter.data.train_dl))\n# imgs, boxss, labelss = fitter._get_batch_model_inputs(imgs, targs)\n# imgs, boxss, labelss = mixup_things(imgs, boxss, labelss, alpha=1.5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# i = 3\n# img, boxs, labels = imgs[i], boxss[i], labelss[i]\n# img_np = img.permute(1, 2, 0).cpu().numpy().copy()\n# boxs = boxs.cpu().numpy().astype(np.int32)\n# fig, ax = plt_img_boxs(img_np, boxs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nlrs, losses = fitter.lr_find(max_lr=1e2)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fitter.plt_lr_find_results(skip_end=1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nfitter.fit(n_epochs=1, max_lr=[1e-4, 5e-4])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fitter.unfreeze()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nlrs, losses = fitter.lr_find(max_lr=1e2)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fitter.plt_lr_find_results(skip_end=10)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nfitter.fit(36, max_lr=[1e-4, 4e-4])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fitter.plt_losses()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"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}