{"cells":[{"metadata":{},"cell_type":"markdown","source":"This is the notebook I used to train my CenterNet using Google Colab Pro. You can reach 0.6964 on private LB without pseudo label with 3-fold ensemble and TTA. \n\nThe inference notebook is here: https://www.kaggle.com/jihangz/wheat-centernet-infer-bifpn-nfolds-tta-v2?scriptVersionId=40044689","execution_count":null},{"metadata":{"id":"SXxrhGRJ9Hwm","outputId":"2325cfd0-a0c4-406c-ced1-277a8e994b9a","trusted":false},"cell_type":"code","source":"!pip install -U git+https://github.com/albumentations-team/albumentations > /dev/null # no output","execution_count":null,"outputs":[]},{"metadata":{"id":"gCECzUKq9Nrm","outputId":"38de0030-2bff-41b2-ef94-28573da8a6b8","trusted":false},"cell_type":"code","source":"!pip install git+https://github.com/qubvel/segmentation_models.pytorch > /dev/null # no output","execution_count":null,"outputs":[]},{"metadata":{"id":"DHvpRz1c9bQX","trusted":false},"cell_type":"code","source":"import os\nimport re\nimport time\nimport logging\nimport math\nimport random\nimport json\nimport copy\nfrom collections import defaultdict\nfrom tqdm.auto import tqdm\n\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.utils.data as data\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import LambdaLR\n\nimport segmentation_models_pytorch as smp\nimport albumentations as A\n\nfrom pynvml import *\nnvmlInit()\n\ntry:\n    from torch.cuda.amp import autocast, GradScaler\nexcept:\n    print('Current PyTorch does NOT support AMP training')","execution_count":null,"outputs":[]},{"metadata":{"id":"qCaFgmcrGNho","trusted":false},"cell_type":"code","source":"class Config:\n    seed = 5468\n    arch = 'timm-efficientnet-b5'\n    heads = {\n        'hm': 1,\n        'wh': 2,\n        'reg': 2}\n    head_conv = 64\n    reg_offset = True\n    \n    # Image\n    data_root = '../input/global-wheat-detection/train'\n    crop_size = 512\n    scale = 0.\n    shift = 0.\n    rotate = 15.\n    shear = 5.\n    down_ratio = 4\n\n    debug = False\n\n    # loss\n    hm_weight = 1\n    off_weight = 1\n    wh_weight = 0.1\n\n    # train\n    batch_size = 12\n    base_lr = 0.75e-4\n    warmup_iters = 1000\n    total_epochs = 100\n    stage_epochs = 100\n    freeze_bn = False\n    accumulate = 1\n    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n    ema = False\n    amp = True\n\n    # logging\n    output_dir = '../input/output/centernet_effnet-b5_bifpn_fold0'\n    logs_dir = os.path.join(output_dir, 'logs')\n    log_interval = 10\n\n    # saving\n    checkpoint = 5\n    load_model = ''\n    \nopt = Config()","execution_count":null,"outputs":[]},{"metadata":{"id":"nG5OuvH6KOgs"},"cell_type":"markdown","source":"# Define Dataset","execution_count":null},{"metadata":{"id":"GElOJfca9soA","trusted":false},"cell_type":"code","source":"with open('../input/wheat-splits/wheat_train_0.json', 'r') as f:\n    data_dict_train = json.load(f)\n\nwith open('../input/wheat-splits/wheat_valid_0.json', 'r') as f:\n    data_dict_valid = json.load(f)","execution_count":null,"outputs":[]},{"metadata":{"id":"T2_zgV4T93uh","trusted":false},"cell_type":"code","source":"def random_affine(img, targets=(), degrees=10, translate=.1, scale=.1, shear=10, border=(0, 0)):\n    # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10))\n    # https://medium.com/uruvideo/dataset-augmentation-with-random-homographies-a8f4b44830d4\n    # targets = [cls, xyxy]\n\n    height = img.shape[0] + border[0] * 2  # shape(h,w,c)\n    width = img.shape[1] + border[1] * 2\n\n    # Rotation and Scale\n    R = np.eye(3)\n    a = random.uniform(-degrees, degrees)\n    # a += random.choice([-180, -90, 0, 90])  # add 90deg rotations to small rotations\n    s = random.uniform(1 - scale, 1 + scale)\n    # s = 2 ** random.uniform(-scale, scale)\n    R[:2] = cv2.getRotationMatrix2D(angle=a, center=(img.shape[1] / 2, img.shape[0] / 2), scale=s)\n\n    # Translation\n    T = np.eye(3)\n    T[0, 2] = random.uniform(-translate, translate) * img.shape[1] + border[1]  # x translation (pixels)\n    T[1, 2] = random.uniform(-translate, translate) * img.shape[0] + border[0]  # y translation (pixels)\n\n    # Shear\n    S = np.eye(3)\n    S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # x shear (deg)\n    S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # y shear (deg)\n\n    # Combined rotation matrix\n    M = S @ T @ R  # ORDER IS IMPORTANT HERE!!\n    if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any():  # image changed\n        img = cv2.warpAffine(img, M[:2], dsize=(width, height), flags=cv2.INTER_LINEAR, borderValue=(114, 114, 114))\n\n    # Transform label coordinates\n    n = len(targets)\n    if n:\n        # warp points\n        xy = np.ones((n * 4, 3))\n        xy[:, :2] = targets[:, [1, 2, 3, 4, 1, 4, 3, 2]].reshape(n * 4, 2)  # x1y1, x2y2, x1y2, x2y1\n        xy = (xy @ M.T)[:, :2].reshape(n, 8)\n\n        # create new boxes\n        x = xy[:, [0, 2, 4, 6]]\n        y = xy[:, [1, 3, 5, 7]]\n        xy = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T\n\n        # reject warped points outside of image\n        xy[:, [0, 2]] = xy[:, [0, 2]].clip(0, width)\n        xy[:, [1, 3]] = xy[:, [1, 3]].clip(0, height)\n        w = xy[:, 2] - xy[:, 0]\n        h = xy[:, 3] - xy[:, 1]\n        area = w * h\n        area0 = (targets[:, 3] - targets[:, 1]) * (targets[:, 4] - targets[:, 2])\n        ar = np.maximum(w / (h + 1e-16), h / (w + 1e-16))  # aspect ratio\n        i = (w > 2) & (h > 2) & (area / (area0 * s + 1e-16) > 0.3) & (ar < 10)\n\n        targets = targets[i]\n        targets[:, 1:5] = xy[i]\n\n    return img, targets","execution_count":null,"outputs":[]},{"metadata":{"id":"RM9T0xuu-Ljz","trusted":false},"cell_type":"code","source":"def multi_scale_transforms(img_size, output_sizes=[512, 640, 768, 896, 1024, 1280, 1536, 1792, 2048]):\n    size = random.choice(output_sizes)\n    scale = size / img_size - 1\n\n    return A.Compose(\n        [   \n            A.RandomScale(scale_limit=(scale, scale), p=1)\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=['category_id']\n        )\n    )\n\n\ndef get_train_transforms(output_size=512):\n    return A.Compose(\n        [   \n            A.HueSaturationValue(hue_shift_limit=5, sat_shift_limit=30, val_shift_limit=20, p=0.8),\n            A.RandomBrightnessContrast(brightness_limit=0.1, contrast_limit=0.2, p=0.8),\n            A.OneOf([\n               A.GaussNoise(p=0.9),\n               A.ISONoise(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.Cutout(num_holes=5, max_h_size=54, max_w_size=54, fill_value=114, p=0.7),\n            A.RandomCrop(height=output_size, width=output_size, p=1),\n        ], \n        p=1.0, \n        bbox_params=A.BboxParams(\n            format='pascal_voc',\n            min_area=0, \n            min_visibility=0.2,\n            label_fields=['category_id']\n        )\n    )\n\ndef get_valid_transforms(output_size=1024):\n    return A.Compose(\n        [\n            A.Resize(height=output_size, width=output_size, 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=['category_id']\n        )\n    )","execution_count":null,"outputs":[]},{"metadata":{"id":"-CoFCZtr-gaB","trusted":false},"cell_type":"code","source":"def xyxy2xywh(x):\n    # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right\n    y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x)\n    y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center\n    y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center\n    y[:, 2] = x[:, 2] - x[:, 0]  # width\n    y[:, 3] = x[:, 3] - x[:, 1]  # height\n    return y\n\n\ndef letterbox(img, new_shape=(512, 512), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True):\n    # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232\n    shape = img.shape[:2]  # current shape [height, width]\n    if isinstance(new_shape, int):\n        new_shape = (new_shape, new_shape)\n\n    # Scale ratio (new / old)\n    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])\n    if not scaleup:  # only scale down, do not scale up (for better test mAP)\n        r = min(r, 1.0)\n\n    # Compute padding\n    ratio = r, r  # width, height ratios\n    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))\n    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding\n    if auto:  # minimum rectangle\n        dw, dh = np.mod(dw, 64), np.mod(dh, 64)  # wh padding\n    elif scaleFill:  # stretch\n        dw, dh = 0.0, 0.0\n        new_unpad = new_shape\n        ratio = new_shape[0] / shape[1], new_shape[1] / shape[0]  # width, height ratios\n\n    dw /= 2  # divide padding into 2 sides\n    dh /= 2\n\n    if shape[::-1] != new_unpad:  # resize\n        img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)\n    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))\n    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))\n    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border\n    return img, ratio, (dw, dh)\n\n\ndef load_image(self, index):\n    # loads 1 image from dataset, returns img, original hw, resized hw\n    img_dict = self.data_dict[index]\n    if self.load_to_ram:\n      img = img_dict['image']\n    else:\n      file_name = img_dict['file_name']\n      img_path = os.path.join(self.data_root, file_name)\n      img = cv2.imread(img_path)\n      assert img is not None, 'Image Not Found ' + img_path\n      img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # RGB\n\n    labels = img_dict['labels']\n    h0, w0 = img.shape[:2]  # orig hw\n    r = self.img_size / max(h0, w0)  # resize image to img_size\n    if r != 1:  # always resize down, only resize up if training with augmentation\n        interp = cv2.INTER_AREA if r < 1 else cv2.INTER_LINEAR\n        img = cv2.resize(img, (int(w0 * r), int(h0 * r)), interpolation=interp)\n    return img, labels, (h0, w0), img.shape[:2]\n\n\ndef load_mosaic(self, index):\n    # loads images in a mosaic\n\n    labels4 = []\n    s = self.img_size\n    yc, xc = [int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border]  # mosaic center x, y\n    indices = [index] + [random.randint(0, len(self.data_dict) - 1) for _ in range(3)]  # 3 additional image indices\n    for i, index in enumerate(indices):\n        # Load image\n        img, labels0, _, (h, w) = load_image(self, index)\n\n        # place img in img4\n        if i == 0:  # top left\n            img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)  # base image with 4 tiles\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\n        img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b]  # img4[ymin:ymax, xmin:xmax]\n        padw = x1a - x1b\n        padh = y1a - y1b\n\n        # Labels\n        if labels0.size > 0:  # Normalized xywh to pixel xyxy format\n          labels = labels0.copy()\n          labels[:, 1] = w * (labels0[:, 1] - labels0[:, 3] / 2) + padw\n          labels[:, 2] = h * (labels0[:, 2] - labels0[:, 4] / 2) + padh\n          labels[:, 3] = w * (labels0[:, 1] + labels0[:, 3] / 2) + padw\n          labels[:, 4] = h * (labels0[:, 2] + labels0[:, 4] / 2) + padh\n        labels4.append(labels)\n\n    # Concat/clip labels\n    if len(labels4):\n        labels4 = np.concatenate(labels4, 0)\n        # np.clip(labels4[:, 1:] - s / 2, 0, s, out=labels4[:, 1:])  # use with center crop\n        np.clip(labels4[:, 1:], 0, 2 * s, out=labels4[:, 1:])  # use with random_affine\n\n        # Replicate\n        # img4, labels4 = replicate(img4, labels4)\n\n    # Augment\n    # img4 = img4[s // 2: int(s * 1.5), s // 2:int(s * 1.5)]  # center crop (WARNING, requires box pruning)\n    img4, labels4 = random_affine(img4, labels4,\n                                  degrees=0.,\n                                  translate=0.,\n                                  scale=0.,\n                                  shear=0.,\n                                  border=self.mosaic_border)  # border to remove\n\n    return img4, labels4\n\n\ndef gaussian_radius_wh(det_size, alpha):\n    height, width = det_size\n    h_radiuses_alpha = int(height / 2. * alpha)\n    w_radiuses_alpha = int(width / 2. * alpha)\n    return h_radiuses_alpha, w_radiuses_alpha\n\ndef gaussian_2d(shape, sigma_x=1, sigma_y=1):\n    m, n = [(ss - 1.) / 2. for ss in shape]\n    y, x = np.ogrid[-m:m + 1, -n:n + 1]\n\n    h = np.exp(-(x * x / (2 * sigma_x * sigma_x) + y * y / (2 * sigma_y * sigma_y)))\n    h[h < np.finfo(h.dtype).eps * h.max()] = 0\n    return h\n\ndef draw_truncate_gaussian(heatmap, center, h_radius, w_radius, k=1):\n    h, w = 2 * h_radius + 1, 2 * w_radius + 1\n    sigma_x = w / 6\n    sigma_y = h / 6\n    gaussian = gaussian_2d((h, w), sigma_x=sigma_x, sigma_y=sigma_y)\n\n    x, y = int(center[0]), int(center[1])\n\n    height, width = heatmap.shape[0:2]\n\n    left, right = min(x, w_radius), min(width - x, w_radius + 1)\n    top, bottom = min(y, h_radius), min(height - y, h_radius + 1)\n\n    masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right]\n    masked_gaussian = gaussian[h_radius - top:h_radius + bottom,\n                      w_radius - left:w_radius + right]\n    if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0:\n        np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap)\n    return heatmap","execution_count":null,"outputs":[]},{"metadata":{"id":"zpZZTbnUED2i","trusted":false},"cell_type":"code","source":"class WheatDataset(data.Dataset):\n\n    def __init__(self, opt, data_root, data_dict, img_size=1024, transforms=None, is_train=True, load_to_ram=False):\n        super().__init__()\n\n        self.num_classes = 1\n        self.mean = np.array([0.315290, 0.317253, 0.214556], dtype=np.float32).reshape(1, 1, 3)\n        self.std = np.array([0.245211, 0.238036, 0.193879], dtype=np.float32).reshape(1, 1, 3)\n        self.max_objs = 256\n\n        self.opt = opt\n\n        self.data_root = data_root\n        self.data_dict = []\n        for _d in tqdm(copy.deepcopy(data_dict)):\n            img_dict = {\n                'file_name': os.path.basename(_d['file_name']),\n                'image_id': _d['id'],\n                'height': _d['height'],\n                'width': _d['width'],\n                'labels': []\n            }\n            if load_to_ram:\n                file_path = os.path.join(self.data_root, img_dict['file_name'])\n                img_dict['image'] = cv2.cvtColor(cv2.imread(file_path), cv2.COLOR_BGR2RGB)\n            for annot in _d['annotations']:\n                bbox = annot['bbox']\n                labels = [\n                    annot['category_id'],\n                    (bbox[0]+bbox[2]/2)/img_dict['width'],\n                    (bbox[1]+bbox[3]/2)/img_dict['height'],\n                    bbox[2]/img_dict['width'],\n                    bbox[3]/img_dict['height']\n                ] # normalized xywh\n                img_dict['labels'].append(labels)\n            img_dict['labels'] = np.array(img_dict['labels'], dtype=np.float32)\n            \n            self.data_dict.append(img_dict)\n\n        self.img_size = img_size\n        self.mosaic_border = [-self.img_size // 2, -self.img_size // 2]\n        self.transforms = transforms\n        self.is_train = is_train\n        self.load_to_ram = load_to_ram\n\n    def __getitem__(self, index):\n        if (not self.is_train) or random.random() > 0.5:\n            img, labels0, (h0, w0), (h, w) = load_image(self, index)\n\n            # Letterbox\n            img, ratio, pad = letterbox(img, self.img_size, auto=False, scaleup=False)\n            shapes = (h0, w0), ((h / h0, w / w0), pad)  # for COCO mAP rescaling\n\n            if labels0.size > 0:\n                # Normalized xywh to pixel xyxy format\n                labels = labels0.copy()\n                labels[:, 1] = ratio[0] * w * (labels0[:, 1] - labels0[:, 3] / 2) + pad[0]  # pad width\n                labels[:, 2] = ratio[1] * h * (labels0[:, 2] - labels0[:, 4] / 2) + pad[1]  # pad height\n                labels[:, 3] = ratio[0] * w * (labels0[:, 1] + labels0[:, 3] / 2) + pad[0]\n                labels[:, 4] = ratio[1] * h * (labels0[:, 2] + labels0[:, 4] / 2) + pad[1]\n        else:\n            # Load mosaic\n            img, labels = load_mosaic(self, index)\n            shapes = None\n\n        if self.is_train:\n            scaler = multi_scale_transforms(self.img_size)\n            scaled = scaler(image=img, bboxes=labels[:, 1:5], category_id=labels[:, 0])\n            img = scaled['image']\n            boxes = scaled['bboxes']\n            cat = scaled['category_id']\n\n            augment_func = get_train_transforms(output_size=self.opt.crop_size)\n            augmented = augment_func(**{\n                'image': img,\n                'bboxes': boxes,\n                'category_id': cat\n                })\n            img = augmented['image']\n            labels = np.zeros((len(augmented['category_id']), 5), dtype=np.float32)\n            if len(labels) > 0:\n                labels[:, 1:5] = augmented['bboxes']\n            \n            img, labels = random_affine(\n                img,\n                labels,\n                degrees=self.opt.rotate,\n                translate=self.opt.shift,\n                scale=self.opt.scale,\n                shear=self.opt.shear)\n            \n        num_objs = len(labels)  # number of labels\n        if num_objs > 0:\n            # convert xyxy to xywh\n            labels[:, 1:5] = xyxy2xywh(labels[:, 1:5])\n\n            # Normalize coordinates 0 - 1\n            labels[:, [2, 4]] /= img.shape[0]  # height\n            labels[:, [1, 3]] /= img.shape[1]  # width\n\n        img = (img.astype(np.float32) / 255.)\n        img = (img - self.mean) / self.std\n        img = img.transpose(2, 0, 1)\n\n        output_h = img.shape[1] // self.opt.down_ratio\n        output_w = img.shape[2] // self.opt.down_ratio\n        num_classes = self.num_classes\n\n        # generate targets\n        hm = np.zeros((num_classes, output_h, output_w), dtype=np.float32)\n        wh = np.zeros((self.max_objs, 2), dtype=np.float32)\n        reg = np.zeros((self.max_objs, 2), dtype=np.float32)\n        ind = np.zeros((self.max_objs), dtype=np.int64)\n        reg_mask = np.zeros((self.max_objs), dtype=np.uint8)\n\n        \n        gt_det = []\n        for k in range(min(num_objs, self.max_objs)):\n            label = labels[k]\n            bbox = label[1:]\n            cls_id = int(label[0])\n            bbox[[0, 2]] = bbox[[0, 2]] * output_w\n            bbox[[1, 3]] = bbox[[1, 3]] * output_h\n            bbox[0] = np.clip(bbox[0], 0, output_w - 1)\n            bbox[1] = np.clip(bbox[1], 0, output_h - 1)\n            h = bbox[3]\n            w = bbox[2]\n\n            if h > 0 and w > 0:\n                ct = np.array([bbox[0], bbox[1]], dtype=np.float32)\n                ct_int = ct.astype(np.int32)\n                h_radius, w_radius = gaussian_radius_wh((math.ceil(h), math.ceil(w)), 0.54)\n                draw_truncate_gaussian(hm[cls_id], ct_int, h_radius, w_radius)\n\n                wh[k] = 1. * w, 1. * h\n                ind[k] = ct_int[1] * output_w + ct_int[0]\n                reg[k] = ct - ct_int\n                reg_mask[k] = 1\n\n                gt_det.append([ct[0] - w / 2, ct[1] - h / 2, \n                        ct[0] + w / 2, ct[1] + h / 2, 1, cls_id])\n        \n        ret = {'input': img, 'hm': hm, 'reg_mask': reg_mask, 'ind': ind, 'wh': wh, 'reg': reg}\n\n        if self.opt.debug:\n            gt_det = np.array(gt_det, dtype=np.float32) if len(gt_det) > 0 else np.zeros((1, 6), dtype=np.float32)\n            meta = {'gt_det': gt_det, 'labels': labels}\n            ret['meta'] = meta\n        return ret\n    \n    def __len__(self):\n        return len(self.data_dict)","execution_count":null,"outputs":[]},{"metadata":{"id":"agoLMgfRQwVS"},"cell_type":"markdown","source":"# Define Dataloader","execution_count":null},{"metadata":{"id":"t0WS4BDZQfsk","trusted":false},"cell_type":"code","source":"from torch._six import container_abcs, string_classes, int_classes\n\ndef collate(batch):\n    r\"\"\"Puts each data field into a tensor with outer dimension batch size\"\"\"\n\n    np_str_obj_array_pattern = re.compile(r'[SaUO]')\n    error_msg = \"batch must contain tensors, numbers, dicts or lists; found {}\"\n    elem = batch[0]\n    elem_type = type(elem)\n    if isinstance(elem, torch.Tensor):\n        out = None\n        if torch.utils.data.get_worker_info() is not None:\n            # If we're in a background process, concatenate directly into a\n            # shared memory tensor to avoid an extra copy\n            numel = sum([x.numel() for x in batch])\n            storage = elem.storage()._new_shared(numel)\n            out = elem.new(storage)\n        return torch.stack(batch, 0, out=out)\n    elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \\\n            and elem_type.__name__ != 'string_':\n        elem = batch[0]\n        if elem_type.__name__ == 'ndarray':\n            # array of string classes and object\n            if np_str_obj_array_pattern.search(elem.dtype.str) is not None:\n                raise TypeError(error_msg.format(elem.dtype))\n\n            return collate([torch.as_tensor(b) for b in batch])\n        elif elem.shape == ():  # scalars\n            return torch.as_tensor(batch)\n    elif isinstance(elem, float):\n        return torch.tensor(batch, dtype=torch.float64)\n    elif isinstance(elem, int_classes):\n        return torch.tensor(batch)\n    elif isinstance(elem, string_classes):\n        return batch\n    elif isinstance(elem, container_abcs.Mapping):\n        res = {key: collate([d[key] for d in batch]) for key in batch[0] if key!='instance_mask'}\n        if 'instance_mask' in batch[0]:\n            max_obj = max([d['instance_mask'].shape[0] for d in batch])\n            instance_mask = torch.zeros(len(batch),max_obj,*(batch[0]['instance_mask'].shape[1:]))\n            for i in range(len(batch)):\n                num_obj = batch[i]['instance_mask'].shape[0]\n                instance_mask[i,:num_obj] = torch.as_tensor(batch[i]['instance_mask'])\n            res.update({'instance_mask':instance_mask})\n        return res\n    elif isinstance(elem, tuple) and hasattr(elem, '_fields'):  # namedtuple\n        return elem_type(*(collate(samples) for samples in zip(*batch)))\n    elif isinstance(elem, container_abcs.Sequence):\n        # check to make sure that the elements in batch have consistent size\n        it = iter(batch)\n        elem_size = len(next(it))\n        if not all(len(elem) == elem_size for elem in it):\n            raise RuntimeError('each element in list of batch should be of equal size')\n        transposed = zip(*batch)\n        return [collate(samples) for samples in transposed]\n\n    raise TypeError(error_msg.format(elem_type))","execution_count":null,"outputs":[]},{"metadata":{"id":"Wal1g50hQyRY","trusted":false},"cell_type":"code","source":"class _RepeatSampler(object):\n    \"\"\" Sampler that repeats forever.\n\n    Args:\n        sampler (Sampler)\n    \"\"\"\n\n    def __init__(self, sampler):\n        self.sampler = sampler\n\n    def __iter__(self):\n        while True:\n            yield from iter(self.sampler)\n\n\nclass FastDataLoader(data.DataLoader):\n\n    def __init__(self, *args, **kwargs):\n        super().__init__(*args, **kwargs)\n        object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler))\n        self.iterator = super().__iter__()\n\n    def __len__(self):\n        return len(self.batch_sampler.sampler)\n\n    def __iter__(self):\n        for i in range(len(self)):\n            yield next(self.iterator)","execution_count":null,"outputs":[]},{"metadata":{"id":"FaZv2sVFKLZg"},"cell_type":"markdown","source":"# Define Model","execution_count":null},{"metadata":{"id":"7-D8M39wgY2R","trusted":false},"cell_type":"code","source":"class Conv2d(torch.nn.Conv2d):\n    \"\"\"\n    A wrapper around :class:`torch.nn.Conv2d` to support empty inputs and more features.\n    \"\"\"\n\n    def __init__(self, *args, **kwargs):\n        \"\"\"\n        Extra keyword arguments supported in addition to those in `torch.nn.Conv2d`:\n\n        Args:\n            norm (nn.Module, optional): a normalization layer\n            activation (callable(Tensor) -> Tensor): a callable activation function\n\n        It assumes that norm layer is used before activation.\n        \"\"\"\n        norm = kwargs.pop(\"norm\", None)\n        activation = kwargs.pop(\"activation\", None)\n        super().__init__(*args, **kwargs)\n\n        self.norm = norm\n        self.activation = activation\n\n\n    def forward(self, x):\n        if x.numel() == 0 and self.training:\n            # https://github.com/pytorch/pytorch/issues/12013\n            assert not isinstance(\n                self.norm, torch.nn.SyncBatchNorm\n            ), \"SyncBatchNorm does not support empty inputs!\"\n\n        if x.numel() == 0 and TORCH_VERSION <= (1, 4):\n            assert not isinstance(\n                self.norm, torch.nn.GroupNorm\n            ), \"GroupNorm does not support empty inputs in PyTorch <=1.4!\"\n            # When input is empty, we want to return a empty tensor with \"correct\" shape,\n            # So that the following operations will not panic\n            # if they check for the shape of the tensor.\n            # This computes the height and width of the output tensor\n            output_shape = [\n                (i + 2 * p - (di * (k - 1) + 1)) // s + 1\n                for i, p, di, k, s in zip(\n                    x.shape[-2:], self.padding, self.dilation, self.kernel_size, self.stride\n                )\n            ]\n            output_shape = [x.shape[0], self.weight.shape[0]] + output_shape\n            empty = _NewEmptyTensorOp.apply(x, output_shape)\n            if self.training:\n                # This is to make DDP happy.\n                # DDP expects all workers to have gradient w.r.t the same set of parameters.\n                _dummy = sum(x.view(-1)[0] for x in self.parameters()) * 0.0\n                return empty + _dummy\n            else:\n                return empty\n\n        x = super().forward(x)\n        if self.norm is not None:\n            x = self.norm(x)\n        if self.activation is not None:\n            x = self.activation(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{"id":"EuBtn82JgfkC","trusted":false},"cell_type":"code","source":"class FeatureMapResampler(nn.Module):\n    def __init__(self, in_channels, out_channels, stride):\n        super(FeatureMapResampler, self).__init__()\n        if in_channels != out_channels:\n            self.reduction = Conv2d(\n                in_channels, out_channels, kernel_size=1,\n                bias=False,\n                norm=nn.BatchNorm2d(out_channels),\n                activation=None\n            )\n        else:\n            self.reduction = None\n\n        assert stride <= 2\n        self.stride = stride\n\n    def forward(self, x):\n        if self.reduction is not None:\n            x = self.reduction(x)\n\n        if self.stride == 2:\n            x = F.max_pool2d(\n                x, kernel_size=self.stride + 1,\n                stride=self.stride, padding=1\n            )\n        elif self.stride == 1:\n            pass\n        else:\n            raise NotImplementedError()\n        return x\n\n\nclass EncoderWithC6(nn.Module):\n    def __init__(self, encoder, out_channels):\n        super(EncoderWithC6, self).__init__()\n        self.encoder = encoder\n        self.sampler = FeatureMapResampler(encoder.out_channels[-1], out_channels, 2)\n        self.out_channels = encoder.out_channels + (out_channels, )\n\n    def forward(self, x):\n        feats = self.encoder(x)\n        x = feats[-1]\n        feats.append(self.sampler(x))\n\n        return feats","execution_count":null,"outputs":[]},{"metadata":{"id":"MnWcgc0Fggpj","trusted":false},"cell_type":"code","source":"class SwishImplementation(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, x):\n        ctx.save_for_backward(x)\n        return x * torch.sigmoid(x)\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        x = ctx.saved_tensors[0]\n        sx = torch.sigmoid(x)\n        return grad_output * (sx * (1 + x * (1 - sx)))\n\n\nclass MemoryEfficientSwish(nn.Module):\n    @staticmethod\n    def forward(x):\n        return SwishImplementation.apply(x)\n    \n\ndef swish(x, inplace: bool = False):\n    \"\"\"Swish - Described in: https://arxiv.org/abs/1710.05941\n    \"\"\"\n    return x.mul_(x.sigmoid()) if inplace else x.mul(x.sigmoid())\n\n\nclass Swish(nn.Module):\n    def __init__(self, inplace: bool = False):\n        super(Swish, self).__init__()\n        self.inplace = inplace\n\n    def forward(self, x):\n        return swish(x, self.inplace)","execution_count":null,"outputs":[]},{"metadata":{"id":"vbpfkft8gj0k","trusted":false},"cell_type":"code","source":"class SingleBiFPN(nn.Module):\n    \"\"\"\n    This module implements Feature Pyramid Network.\n    It creates pyramid features built on top of some input feature maps.\n    \"\"\"\n\n    def __init__(self, in_channels_list, out_channels):\n        super(SingleBiFPN, self).__init__()\n        \n        self.swish = MemoryEfficientSwish()\n        self.out_channels = out_channels\n        # build 5-levels bifpn\n        if len(in_channels_list) == 5:\n            self.nodes = [\n                {'feat_level': 3, 'inputs_offsets': [3, 4]},\n                {'feat_level': 2, 'inputs_offsets': [2, 5]},\n                {'feat_level': 1, 'inputs_offsets': [1, 6]},\n                {'feat_level': 0, 'inputs_offsets': [0, 7]},\n                {'feat_level': 1, 'inputs_offsets': [1, 7, 8]},\n                {'feat_level': 2, 'inputs_offsets': [2, 6, 9]},\n                {'feat_level': 3, 'inputs_offsets': [3, 5, 10]},\n                {'feat_level': 4, 'inputs_offsets': [4, 11]},\n            ]\n        elif len(in_channels_list) == 3:\n            self.nodes = [\n                {'feat_level': 1, 'inputs_offsets': [1, 2]},\n                {'feat_level': 0, 'inputs_offsets': [0, 3]},\n                {'feat_level': 1, 'inputs_offsets': [1, 3, 4]},\n                {'feat_level': 2, 'inputs_offsets': [2, 5]},\n            ]\n        else:\n            raise NotImplementedError\n\n        node_info = [_ for _ in in_channels_list]\n\n        num_output_connections = [0 for _ in in_channels_list]\n        for fnode in self.nodes:\n            feat_level = fnode[\"feat_level\"]\n            inputs_offsets = fnode[\"inputs_offsets\"]\n            inputs_offsets_str = \"_\".join(map(str, inputs_offsets))\n            for input_offset in inputs_offsets:\n                num_output_connections[input_offset] += 1\n\n                in_channels = node_info[input_offset]\n                if in_channels != out_channels:\n                    lateral_conv = Conv2d(\n                        in_channels,\n                        out_channels,\n                        kernel_size=1,\n                        norm=nn.BatchNorm2d(out_channels)\n                    )\n                    self.add_module(\n                        \"lateral_{}_f{}\".format(input_offset, feat_level), lateral_conv\n                    )\n            node_info.append(out_channels)\n            num_output_connections.append(0)\n\n            # generate attention weights\n            name = \"weights_f{}_{}\".format(feat_level, inputs_offsets_str)\n            self.__setattr__(name, nn.Parameter(\n                    torch.ones(len(inputs_offsets), dtype=torch.float32),\n                    requires_grad=True\n                ))\n\n            # generate convolutions after combination\n            name = \"outputs_f{}_{}\".format(feat_level, inputs_offsets_str)\n            self.add_module(name, Conv2d(\n                out_channels,\n                out_channels,\n                kernel_size=3,\n                padding=1,\n                norm=nn.BatchNorm2d(out_channels),\n                bias=False\n            ))\n\n    def forward(self, feats):\n        feats = [_ for _ in feats]\n        num_levels = len(feats)\n        num_output_connections = [0 for _ in feats]\n        for fnode in self.nodes:\n            feat_level = fnode[\"feat_level\"]\n            inputs_offsets = fnode[\"inputs_offsets\"]\n            inputs_offsets_str = \"_\".join(map(str, inputs_offsets))\n            input_nodes = []\n            _, _, target_h, target_w = feats[feat_level].size()\n            for input_offset in inputs_offsets:\n                num_output_connections[input_offset] += 1\n                input_node = feats[input_offset]\n\n                # reduction\n                if input_node.size(1) != self.out_channels:\n                    name = \"lateral_{}_f{}\".format(input_offset, feat_level)\n                    input_node = self.__getattr__(name)(input_node)\n\n                # maybe downsample\n                _, _, h, w = input_node.size()\n                if h > target_h and w > target_w:\n                    height_stride_size = int((h - 1) // target_h + 1)\n                    width_stride_size = int((w - 1) // target_w + 1)\n                    assert height_stride_size == width_stride_size == 2\n                    input_node = F.max_pool2d(\n                        input_node, kernel_size=(height_stride_size + 1, width_stride_size + 1),\n                        stride=(height_stride_size, width_stride_size), padding=1\n                    )\n                elif h <= target_h and w <= target_w:\n                    if h < target_h or w < target_w:\n                        input_node = F.interpolate(\n                            input_node,\n                            size=(target_h, target_w),\n                            mode=\"nearest\"\n                        )\n                else:\n                    raise NotImplementedError()\n                input_nodes.append(input_node)\n\n            # attention\n            name = \"weights_f{}_{}\".format(feat_level, inputs_offsets_str)\n            weights = F.relu(self.__getattr__(name))\n            norm_weights = weights / (weights.sum() + 0.0001)\n\n            new_node = torch.stack(input_nodes, dim=-1)\n            new_node = (norm_weights * new_node).sum(dim=-1)\n            new_node = self.swish(new_node)\n\n            name = \"outputs_f{}_{}\".format(feat_level, inputs_offsets_str)\n            feats.append(self.__getattr__(name)(new_node))\n\n            num_output_connections.append(0)\n\n        output_feats = []\n        for idx in range(num_levels):\n            for i, fnode in enumerate(reversed(self.nodes)):\n                if fnode['feat_level'] == idx:\n                    output_feats.append(feats[-1 - i])\n                    break\n            else:\n                raise ValueError()\n        return output_feats\n\nclass BiFPN(nn.Module):\n    \"\"\"\n    This module implements Feature Pyramid Network.\n    It creates pyramid features built on top of some input feature maps.\n    \"\"\"\n\n    def __init__(self, in_channels, out_channels, num_repeats):\n        super(BiFPN, self).__init__()\n        \n        self.out_channels = out_channels\n\n        # build bifpn\n        self.repeated_bifpn = nn.ModuleList()\n        for i in range(num_repeats):\n            if i == 0:\n                in_channels_list = in_channels\n            else:\n                in_channels_list = [out_channels] * len(in_channels)\n            self.repeated_bifpn.append(SingleBiFPN(\n                in_channels_list, out_channels\n            ))\n\n    def forward(self, feats):\n        for bifpn in self.repeated_bifpn:\n             feats = bifpn(feats)\n        return feats","execution_count":null,"outputs":[]},{"metadata":{"id":"KfpzbEFzgnbn","trusted":false},"cell_type":"code","source":"class Conv3x3GNReLU(nn.Module):\n    def __init__(self, in_channels, out_channels, upsample=False):\n        super().__init__()\n        self.upsample = upsample\n        self.block = nn.Sequential(\n            nn.Conv2d(\n                in_channels, out_channels, (3, 3), stride=1, padding=1, bias=False\n            ),\n            nn.GroupNorm(32, out_channels),\n            nn.ReLU(inplace=True),\n        )\n\n    def forward(self, x):\n        x = self.block(x)\n        if self.upsample:\n            x = F.interpolate(x, scale_factor=2, mode=\"bilinear\", align_corners=True)\n        return x\n    \nclass SegmentationBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, n_upsamples=0):\n        super().__init__()\n\n        blocks = [Conv3x3GNReLU(in_channels, out_channels, upsample=bool(n_upsamples))]\n\n        if n_upsamples > 1:\n            for _ in range(1, n_upsamples):\n                blocks.append(Conv3x3GNReLU(out_channels, out_channels, upsample=True))\n\n        self.block = nn.Sequential(*blocks)\n\n    def forward(self, x):\n        return self.block(x)\n    \nclass MergeBlock(nn.Module):\n    def __init__(self, policy):\n        super().__init__()\n        if policy not in [\"add\", \"cat\"]:\n            raise ValueError(\n                \"`merge_policy` must be one of: ['add', 'cat'], got {}\".format(\n                    policy\n                )\n            )\n        self.policy = policy\n\n    def forward(self, x):\n        if self.policy == 'add':\n            return sum(x)\n        elif self.policy == 'cat':\n            return torch.cat(x, dim=1)\n        else:\n            raise ValueError(\n                \"`merge_policy` must be one of: ['add', 'cat'], got {}\".format(self.policy)\n            )\n            \nclass BiFPNDecoder(nn.Module):\n    def __init__(\n        self,\n        encoder_channels,\n        pyramid_channels=160,\n        num_repeats=3,\n        segmentation_channels=64,\n        merge_policy='add',\n    ):\n        super().__init__()\n        \n        self.out_channels = segmentation_channels if merge_policy == \"add\" else segmentation_channels * 5\n        encoder_channels = encoder_channels[-5:]\n        \n        self.bifpn = BiFPN(encoder_channels, pyramid_channels, num_repeats)\n        \n        self.seg_blocks = nn.ModuleList([\n            SegmentationBlock(pyramid_channels, segmentation_channels, n_upsamples=n_upsamples)\n            for n_upsamples in [0, 1, 2, 3, 4]\n        ])\n        \n        self.merge = MergeBlock(merge_policy)\n    \n    def forward(self, feats):\n        ps = self.bifpn(feats)\n        \n        feature_pyramid = [seg_block(p) for seg_block, p in zip(self.seg_blocks, ps)]\n        x = self.merge(feature_pyramid)\n        \n        return x","execution_count":null,"outputs":[]},{"metadata":{"id":"jP3rqK7kKEiZ","trusted":false},"cell_type":"code","source":"class PoseFPNNet(nn.Module):\n    def __init__(self, base_name, heads, head_conv=256, pyramid_channels=160, num_repeats=3):\n        super(PoseFPNNet, self).__init__()\n\n        source = 'noisy-student' if base_name[:4] == 'timm' else 'imagenet'\n        print('Pretrained source {}'.format(source))\n        backbone = smp.encoders.get_encoder(base_name, weights=source)\n        self.encoder = EncoderWithC6(backbone, pyramid_channels)\n        self.decoder = BiFPNDecoder(self.encoder.out_channels, pyramid_channels=pyramid_channels, num_repeats=num_repeats)\n\n        self.heads = heads\n        for head in self.heads:\n            classes = self.heads[head]\n            fc = nn.Sequential(\n                nn.Conv2d(64, head_conv,\n                          kernel_size=3, padding=1, bias=True),\n                nn.ReLU(inplace=True),\n                nn.Conv2d(head_conv, classes,\n                          kernel_size=1, stride=1,\n                          padding=0, bias=True))\n            if 'hm' in head:\n                fc[-1].bias.data.fill_(-2.19)\n            else:\n                fill_fc_weights(fc)\n            self.__setattr__(head, fc)\n\n        del backbone\n\n    def forward(self, x):\n        features = self.encoder(x)\n        x = self.decoder(features[-5:])\n\n        z = {}\n        for head in self.heads:\n            z[head] = self.__getattr__(head)(x)\n        return [z]\n\n    def freeze_backbone(self):\n        for p in self.encoder.parameters():\n            p.requires_grad = False\n        for p in self.decoder.parameters():\n            p.requires_grad = False\n\n    def freeze_head(self, heads):\n        for head in heads:\n            for p in self.__getattr__(head).parameters():\n                p.requires_grad = False\n\n    def set_mode(self, mode, is_freeze_bn=False):\n        self.mode = mode\n        if mode in ['eval', 'valid', 'test']:\n            self.eval()\n        elif mode in ['train']:\n            self.train()\n            if is_freeze_bn==True: ##freeze\n                for m in self.modules():\n                    if isinstance(m, nn.BatchNorm2d):\n                        m.eval()\n                        # m.weight.requires_grad = False\n                        # m.bias.requires_grad   = False\n\ndef fill_fc_weights(layers):\n    for m in layers.modules():\n        if isinstance(m, nn.Conv2d):\n            nn.init.normal_(m.weight, std=0.001)\n            # torch.nn.init.kaiming_normal_(m.weight.data, nonlinearity='relu')\n            # torch.nn.init.xavier_normal_(m.weight.data)\n            if m.bias is not None:\n                nn.init.constant_(m.bias, 0)\n\n\ndef get_pose_net(base_name, heads, head_conv):\n    model = PoseFPNNet(base_name, heads, head_conv)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"id":"jreThbDBPk91","trusted":false},"cell_type":"code","source":"# Helper functions\n\ndef _sigmoid(x):\n  y = torch.clamp(x.sigmoid_(), min=1e-4, max=1-1e-4)\n  return y\n\ndef _gather_feat(feat, ind, mask=None):\n    dim  = feat.size(2)\n    ind  = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)\n    feat = feat.gather(1, ind)\n    if mask is not None:\n        mask = mask.unsqueeze(2).expand_as(feat)\n        feat = feat[mask]\n        feat = feat.view(-1, dim)\n    return feat\n\ndef _tranpose_and_gather_feat(feat, ind):\n    feat = feat.permute(0, 2, 3, 1).contiguous()\n    feat = feat.view(feat.size(0), -1, feat.size(3))\n    feat = _gather_feat(feat, ind)\n    return feat","execution_count":null,"outputs":[]},{"metadata":{"id":"7-mCCnBIP78Y","trusted":false},"cell_type":"code","source":"def _neg_loss(pred, gt):\n  ''' Modified focal loss. Exactly the same as CornerNet.\n      Runs faster and costs a little bit more memory\n    Arguments:\n      pred (batch x c x h x w)\n      gt_regr (batch x c x h x w)\n  '''\n  pos_inds = gt.eq(1).float()\n  neg_inds = gt.lt(1).float()\n\n  neg_weights = torch.pow(1 - gt, 4)\n\n  loss = 0\n\n  pos_loss = torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds\n  neg_loss = torch.log(1 - pred) * torch.pow(pred, 2) * neg_weights * neg_inds\n\n  num_pos  = pos_inds.float().sum()\n  pos_loss = pos_loss.sum()\n  neg_loss = neg_loss.sum()\n\n  if num_pos == 0:\n    loss = loss - neg_loss\n  else:\n    loss = loss - (pos_loss + neg_loss) / num_pos\n  return loss\n\n\nclass FocalLoss(nn.Module):\n  '''nn.Module warpper for focal loss'''\n  def __init__(self):\n    super(FocalLoss, self).__init__()\n    self.neg_loss = _neg_loss\n\n  def forward(self, out, target):\n    return self.neg_loss(out, target)\n\n\nclass RegL1Loss(nn.Module):\n  def __init__(self):\n    super(RegL1Loss, self).__init__()\n  \n  def forward(self, output, mask, ind, target):\n    pred = _tranpose_and_gather_feat(output, ind)\n    mask = mask.unsqueeze(2).expand_as(pred).float()\n    # loss = F.l1_loss(pred * mask, target * mask, reduction='elementwise_mean')\n    loss = F.l1_loss(pred * mask, target * mask, reduction='sum')\n    loss = loss / (mask.sum() + 1e-4)\n    return loss","execution_count":null,"outputs":[]},{"metadata":{"id":"V-hgqZgIL2gM","trusted":false},"cell_type":"code","source":"class CtdetLoss(nn.Module):\n  def __init__(self, opt):\n    super(CtdetLoss, self).__init__()\n    self.crit = FocalLoss()\n    self.crit_reg = RegL1Loss()\n    self.crit_wh = self.crit_reg\n    self.opt = opt\n\n  def forward(self, outputs, batch):\n    opt = self.opt\n    hm_loss, wh_loss, off_loss = 0, 0, 0\n\n    output = outputs[-1]\n    output['hm'] = _sigmoid(output['hm'])\n\n    hm_loss += self.crit(output['hm'], batch['hm'])\n    if opt.wh_weight > 0:\n        wh_loss += self.crit_reg(\n            output['wh'], batch['reg_mask'],\n            batch['ind'], batch['wh'])\n      \n    if opt.off_weight > 0:\n        off_loss += self.crit_reg(\n            output['reg'], batch['reg_mask'],\n            batch['ind'], batch['reg'])\n        \n    loss = opt.hm_weight * hm_loss + opt.wh_weight * wh_loss + opt.off_weight * off_loss\n    loss_stats = {'loss': loss, 'hm_loss': hm_loss, 'wh_loss': wh_loss, 'off_loss': off_loss}\n    return loss, loss_stats\n\n\nclass ModleWithLoss(nn.Module):\n  def __init__(self, model, loss):\n    super(ModleWithLoss, self).__init__()\n    self.model = model\n    self.loss = loss\n  \n  def forward(self, batch):\n    outputs = self.model(batch['input'])\n    loss, loss_stats = self.loss(outputs, batch)\n    return outputs[-1], loss, loss_stats","execution_count":null,"outputs":[]},{"metadata":{"id":"diJM1xD-T9fe"},"cell_type":"markdown","source":"# Logger helpers","execution_count":null},{"metadata":{"id":"2MuMg3NyTnj9","trusted":false},"cell_type":"code","source":"def create_folder(fd):\n    if not os.path.exists(fd):\n        os.makedirs(fd)\n        \n\ndef create_logging(log_dir, filemode):\n    create_folder(log_dir)\n    i1 = 0\n\n    while os.path.isfile(os.path.join(log_dir, '{:04d}.log'.format(i1))):\n        i1 += 1\n        \n    log_path = os.path.join(log_dir, '{:04d}.log'.format(i1))\n    logging.basicConfig(\n        level=logging.DEBUG,\n        format='%(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s %(message)s',\n        datefmt='%a, %d %b %Y %H:%M:%S',\n        filename=log_path,\n        filemode=filemode)\n\n    # Print to console\n    console = logging.StreamHandler()\n    console.setLevel(logging.INFO)\n    formatter = logging.Formatter('%(name)-12s: %(levelname)-8s %(message)s')\n    console.setFormatter(formatter)\n    logging.getLogger('').addHandler(console)\n    \n    return logging","execution_count":null,"outputs":[]},{"metadata":{"id":"NIO9QtwoUCVA"},"cell_type":"markdown","source":"# Training","execution_count":null},{"metadata":{"id":"5_8WKJON3xvQ","trusted":false},"cell_type":"code","source":"class ModelEMA:\n    \"\"\" Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models\n    Keep a moving average of everything in the model state_dict (parameters and buffers).\n    This is intended to allow functionality like\n    https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage\n    A smoothed version of the weights is necessary for some training schemes to perform well.\n    This class is sensitive where it is initialized in the sequence of model init,\n    GPU assignment and distributed training wrappers.\n    \"\"\"\n\n    def __init__(self, model, decay=0.9999, updates=0):\n        # Create EMA\n        self.ema = copy.deepcopy(model).eval()  # FP32 EMA\n        # if next(model.parameters()).device.type != 'cpu':\n        #     self.ema.half()  # FP16 EMA\n        self.updates = updates  # number of EMA updates\n        self.decay = lambda x: decay * (1 - math.exp(-x / 2000))  # decay exponential ramp (to help early epochs)\n        for p in self.ema.parameters():\n            p.requires_grad_(False)\n\n    def update(self, model):\n        # Update EMA parameters\n        with torch.no_grad():\n            self.updates += 1\n            d = self.decay(self.updates)\n\n            msd = model.state_dict()  # model state_dict\n            for k, v in self.ema.state_dict().items():\n                if v.dtype.is_floating_point:\n                    v *= d\n                    v += (1. - d) * msd[k].detach()\n\n    def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):\n        # Update EMA attributes\n        copy_attr(self.ema, model, include, exclude)","execution_count":null,"outputs":[]},{"metadata":{"id":"tE9VHaNmehES","trusted":false},"cell_type":"code","source":"def get_cosine_schedule_with_warmup(\n    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1\n):\n    \"\"\"\n    Create a schedule with a learning rate that decreases following the values of the cosine function between the\n    initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the\n    initial lr set in the optimizer.\n    Args:\n        optimizer (:class:`~torch.optim.Optimizer`):\n            The optimizer for which to schedule the learning rate.\n        num_warmup_steps (:obj:`int`):\n            The number of steps for the warmup phase.\n        num_training_steps (:obj:`int`):\n            The total number of training steps.\n        num_cycles (:obj:`float`, `optional`, defaults to 0.5):\n            The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0\n            following a half-cosine).\n        last_epoch (:obj:`int`, `optional`, defaults to -1):\n            The index of the last epoch when resuming training.\n    Return:\n        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.\n    \"\"\"\n\n    def lr_lambda(current_step):\n        if current_step < num_warmup_steps:\n            return float(current_step) / float(max(1, num_warmup_steps))\n        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))\n        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))\n\n    return LambdaLR(optimizer, lr_lambda, last_epoch)","execution_count":null,"outputs":[]},{"metadata":{"id":"Q2PLzki3UDbk","trusted":false},"cell_type":"code","source":"def set_seed(seed):\n    torch.manual_seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed(seed)\n\n    random.seed(seed)\n    np.random.seed(seed)\n\ndef freeze_bn(module):\n    if isinstance(module, nn.BatchNorm1d) or isinstance(module, nn.BatchNorm2d) or isinstance(module, nn.SyncBatchNorm):\n        module.eval()\n\ndef train_one_epoch(opt, model, optimizer, scheduler, data_loader, epoch, ema=None, scaler=None):\n    start_time = time.time()\n\n    tq = enumerate(data_loader)\n\n    total_loss = defaultdict(float)\n\n    model.train()\n    if opt.freeze_bn:\n        model.apply(freeze_bn) # freeze bn\n    optimizer.zero_grad()\n\n    for batch_idx, inputs in tq:\n        for k in inputs:\n            if k != 'meta':\n                inputs[k] = inputs[k].to(device=opt.device, non_blocking=True)\n    \n        if opt.amp:\n            with autocast():\n                output, loss, loss_stats = model(inputs)\n            scaler.scale(loss).backward()\n            \n            # update weights\n            if (batch_idx % opt.accumulate == 0) and (batch_idx > 0):\n                scaler.step(optimizer)\n                scaler.update()\n                scheduler.step()\n                optimizer.zero_grad()\n                if ema is not None:\n                    ema.update(model)\n        \n        else:\n            output, loss, loss_stats = model(inputs)\n            loss.backward()\n\n            # update weights\n            if (batch_idx % opt.accumulate == 0) and (batch_idx > 0):\n                optimizer.step()\n                scheduler.step()\n                optimizer.zero_grad()\n                if ema is not None:\n                    ema.update(model)\n\n        # logging\n        for k, v in loss_stats.items():\n            total_loss[k] += v.item()\n        if (batch_idx % opt.log_interval == 0) and (batch_idx > 0):\n            elapsed = time.time() - start_time\n            logging.info('| epoch {:3d} | {:3d}/{:3d} batches | ms/batch {:5.2f} | lr {:5.3e} | loss {:5.4f} | hm loss {:5.4f} | wh loss {:5.4f} | reg loss {:5.4f} | memory {:5.2f} MB'.format(\n                epoch, batch_idx, len(data_loader),\n                elapsed * 1000 / opt.log_interval, optimizer.param_groups[0]['lr'],\n                total_loss['loss'] / opt.log_interval, total_loss['hm_loss'] / opt.log_interval, total_loss['wh_loss'] / opt.log_interval, total_loss['off_loss'] / opt.log_interval,\n                nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(0)).used / 1e6))\n            \n            total_loss = defaultdict(float)\n            start_time = time.time()\n\n\ndef do_evaluate(opt, model, data_loader, scaler=None):\n    model.eval()\n    tq = enumerate(data_loader)\n    total_loss = defaultdict(float)\n\n    with torch.no_grad():\n        for batch_idx, inputs in tq:\n            for k in inputs:\n                if k != 'meta':\n                    inputs[k] = inputs[k].to(device=opt.device, non_blocking=True)\n\n            if opt.amp:\n                with autocast():\n                    output, loss, loss_stats = model(inputs)\n            else:\n                output, loss, loss_stats = model(inputs)\n\n            for k, v in loss_stats.items():\n                total_loss[k] += v.item() * data_loader.batch_size\n    \n    for k, v in total_loss.items():\n        total_loss[k] = v / len(data_loader.dataset)\n    \n    return total_loss","execution_count":null,"outputs":[]},{"metadata":{"id":"2K1lRmyXcimR"},"cell_type":"markdown","source":"# Main process","execution_count":null},{"metadata":{"id":"Y9ta9qBScdWD","trusted":false},"cell_type":"code","source":"def main(opt):\n    set_seed(opt.seed)\n    torch.backends.cudnn.benchmark = True\n    \n    create_logging(opt.logs_dir, 'w')\n\n    train_dataset = WheatDataset(opt, opt.data_root, data_dict_train, img_size=1024, transforms=get_train_transforms(opt.crop_size), is_train=True, load_to_ram=True)\n    train_loader = FastDataLoader(\n        train_dataset, \n        opt.batch_size,\n        collate_fn=collate,\n        shuffle=True, \n        drop_last=True,\n        pin_memory=True,\n        num_workers=4)\n\n    valid_dataset = WheatDataset(opt, opt.data_root, data_dict_valid, img_size=1024, transforms=get_valid_transforms(1024), is_train=False, load_to_ram=False)\n    valid_loader = FastDataLoader(\n        valid_dataset, \n        opt.batch_size,\n        collate_fn=collate,\n        shuffle=False, \n        drop_last=False,\n        pin_memory=True,\n        num_workers=4)\n    \n    model = ModleWithLoss(PoseFPNNet(opt.arch, opt.heads, opt.head_conv), CtdetLoss(opt)).to(opt.device)\n    if opt.ema:\n        print('Training with EMA')\n        ema = ModelEMA(model)\n    else:\n        ema = None\n    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=opt.base_lr)\n    num_training_steps = int(opt.total_epochs * len(train_dataset) / opt.batch_size / opt.accumulate)\n    scheduler = get_cosine_schedule_with_warmup(optimizer, opt.warmup_iters, num_training_steps)\n    current_epoch = 0\n\n    if opt.load_model != '':\n        # Load model weights\n        checkpoint = torch.load(opt.load_model)\n        model.load_state_dict(checkpoint['model'])\n        optimizer.load_state_dict(checkpoint['optimizer'])\n        scheduler.load_state_dict(checkpoint['scheduler'])\n        current_epoch = checkpoint['epoch'] + 1\n\n    scaler = GradScaler() if opt.amp else None\n    for epoch in tqdm(range(current_epoch, current_epoch + opt.stage_epochs)):\n\n        epoch_start_time = time.time()\n        train_one_epoch(opt, model, optimizer, scheduler, train_loader, epoch, ema=ema, scaler=scaler)\n\n        if ((epoch + 1) % opt.checkpoint == 0) and (epoch > 0):\n            logging.info('Saving checkpoint...')\n            torch.save({\n                'epoch': epoch,\n                'model': model.state_dict(),\n                'ema': ema.ema.state_dict() if opt.ema else None,\n                'ema_updates': ema.updates if opt.ema else None,\n                'optimizer': optimizer.state_dict(),\n                'scheduler': scheduler.state_dict()\n            }, os.path.join(opt.output_dir, '{:05d}.pth'.format(epoch)))\n\n            eval_loss = do_evaluate(opt, ema.ema if opt.ema else model, valid_loader, scaler=scaler)\n            logging.info('-' * 89)\n            logging.info('end of epoch {:4d} | time: {:5.2f}s | val loss {:5.4f} | val hm loss {:5.4f} | val wh loss {:5.4f} | val reg loss {:5.4f} |'.format(\n                epoch, (time.time() - epoch_start_time), eval_loss['loss'], eval_loss['hm_loss'], eval_loss['wh_loss'], eval_loss['off_loss']))\n            logging.info('-' * 89)\n\n        torch.save({\n            'epoch': epoch,\n            'model': model.state_dict(),\n            'optimizer': optimizer.state_dict(),\n            'scheduler': scheduler.state_dict()\n        }, os.path.join(opt.output_dir, 'last.pth'))","execution_count":null,"outputs":[]},{"metadata":{"id":"p_2Fy6wzFjYv","outputId":"3dc2aa55-8cb0-4550-e33b-6d5665439f5a","trusted":false},"cell_type":"code","source":"main(opt)","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}