{"cells":[{"metadata":{},"cell_type":"markdown","source":"In this kernel, I show how I changed the backbone of the Faster-R-CNN model from ResNet50 to ResNet152. To achieve that, I used some of the source code of the torchvision and changed it manually.","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"The following points are covered:\n* Create dataset\n* Create dataloader\n* Prepare the model\n* Training the model","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## Load Libraries","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nfrom pathlib import Path\nfrom dataclasses import dataclass\nfrom enum import Enum\nfrom tqdm import tqdm\n\nimport pandas as pd\nimport numpy as np\n\nimport torch\nfrom torch.utils.data import Dataset\nfrom torch.utils.data import DataLoader\nimport torch.optim as optim\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor, FasterRCNN\nfrom torchvision.models.detection.backbone_utils import resnet_fpn_backbone\n\nfrom albumentations import *\nfrom albumentations.pytorch import ToTensorV2\nimport cv2","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Prepare Input variables","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class Device(Enum):\n    GPU = \"GPU\"\n    TPU = \"TPU\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"root_data_dir = Path(\"/kaggle/input/global-wheat-detection/\")\ntest_train_ratio = 0.1\nbatch_size=8\nseed = 0\ntrain_device = Device.GPU\nnumber_of_epochs = 2\nlearning_rate = 0.0001\nweight_decay = 1e-5","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@dataclass\nclass DatasetArguments:\n    data_dir: Path\n    images_lists_dict: dict\n    labels_csv_file_name: str\n\n@dataclass\nclass DataLoaderArguments:\n    batch_size: int\n    num_workers: int\n    dataset_arguments: DatasetArguments","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"## Split Data to train and val datasets","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def _get_images_file_names_from_csv(directory):\n    dataframe = pd.read_csv(os.path.join(directory, \"train.csv\"))\n    files = dataframe[\"image_id\"].unique().tolist()\n    return files","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _choose_train_valid_file_names(file_names, valid_numbers, seed):\n    np.random.seed(seed)\n    valid_file_names = np.random.choice(file_names, valid_numbers, replace=False).tolist()\n    train_file_names = [file_name_i for file_name_i in file_names if file_name_i not in valid_file_names]\n    return train_file_names, valid_file_names","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#split data\nfile_names = _get_images_file_names_from_csv(root_data_dir)\nvalid_numbers = round(len(file_names) * test_train_ratio)\ntrain_file_names, valid_file_names = _choose_train_valid_file_names(file_names, valid_numbers, seed)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"images_lists_dict = {\n    \"train\": train_file_names,\n    \"val\": valid_file_names,\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset_arguments = DatasetArguments(\n    data_dir=root_data_dir,\n    images_lists_dict=images_lists_dict,\n    labels_csv_file_name=\"train.csv\",\n)\n\ndataloaders_arguments = DataLoaderArguments(\n    batch_size=batch_size,\n    num_workers=1,\n    dataset_arguments=dataset_arguments\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Prepare the transforms:\nI chose some of the transforms from this [notebook](https://www.kaggle.com/shonenkov/training-efficientdet)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def transform_set():\n    transforms_dict = {\n        'train': get_train_transforms(),\n        'val': get_valid_transforms()\n    }\n    return transforms_dict\n\n\ndef get_train_transforms():\n    return Compose(\n        [OneOf([HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2,\n                                   val_shift_limit=0.2, p=0.9),\n                RandomBrightnessContrast(brightness_limit=0.2,\n                                         contrast_limit=0.2, p=0.9)],\n               p=0.9),\n            ToGray(p=0.01),\n            HorizontalFlip(p=0.5),\n            VerticalFlip(p=0.5),\n            Cutout(num_holes=8, max_h_size=64, max_w_size=64, fill_value=0, p=0.5),\n            ToTensorV2(p=1.0)],\n        p=1.0,\n        bbox_params=BboxParams(\n            format='pascal_voc',\n            min_area=0,\n            min_visibility=0,\n            label_fields=['labels']\n        )\n    )\n\n\ndef get_valid_transforms():\n    return Compose(\n        [\n            ToTensorV2(p=1.0),\n        ],\n        p=1.0,\n        bbox_params=BboxParams(\n            format='pascal_voc',\n            min_area=0,\n            min_visibility=0,\n            label_fields=['labels']\n        )\n    )\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Create the pytorch dataset","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def _adjust_boxes_format(boxes):\n    # original format [xmin, ymin, width, height]\n    # new format [xmin, ymin, xmax, ymax]\n    adjusted_boxes = []\n    for box_i in boxes:\n        adjusted_box_i = [0, 0, 0, 0]\n        adjusted_box_i[0] = box_i[0]\n        adjusted_box_i[1] = box_i[1]\n        adjusted_box_i[2] = box_i[0] + box_i[2]\n        adjusted_box_i[3] = box_i[1] + box_i[3]\n        adjusted_boxes.append(adjusted_box_i)\n    return adjusted_boxes","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _areas(boxes):\n    # original format [xmin, ymin, width, height]\n    areas = []\n    for box_i in boxes:\n        areas.append(box_i[2] * box_i[3])\n    return areas","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# dataset\nclass ObjectDetectionDataset(Dataset):\n    def __init__(self, images_root_directory,\n                 images_list,\n                 labels_csv_file_name,\n                 phase,\n                 transforms):\n        super(ObjectDetectionDataset).__init__()\n        self.images_root_directory = images_root_directory\n        self.phase = phase\n        self.transforms = transforms\n        self.images_list = images_list\n        if self.phase in [\"train\", \"val\"]:\n            self.labels_dataframe = pd.read_csv(os.path.join(images_root_directory, labels_csv_file_name))\n\n    def __getitem__(self, item):\n        sample = {\n            \"local_image_id\": None,\n            \"image_id\": None,\n            \"labels\": None,\n            \"boxes\": None,\n            \"area\": None,\n            \"iscrowd\": None\n        }\n\n        image_id = self.images_list[item]\n        image_path = os.path.join(self.images_root_directory,\n                                  \"train\" if self.phase in [\"train\", \"val\"] else \"test\",\n                                  image_id + \".jpg\")\n        image = cv2.imread(image_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n        sample[\"local_image_id\"] = image_id\n        sample[\"image_id\"] = torch.tensor([item])\n        if self.phase in [\"train\", \"val\"]:\n            boxes = self.labels_dataframe[self.labels_dataframe.image_id == image_id].bbox.values.tolist()\n            boxes = [eval(box_i) for box_i in boxes]\n            areas = _areas(boxes)\n            boxes = _adjust_boxes_format(boxes)\n\n            sample[\"labels\"] = torch.ones((len(boxes),), dtype=torch.int64)\n            sample[\"boxes\"] = torch.as_tensor(boxes, dtype=torch.float32)\n            sample[\"area\"] = torch.as_tensor(areas, dtype=torch.float32)\n            sample[\"iscrowd\"] = torch.zeros((len(boxes),), dtype=torch.int64)\n        if self.transforms is not None:\n            sample[\"image\"] = image\n            transformed_sample = self.transforms(image=sample[\"image\"],\n                                                 bboxes=sample[\"boxes\"],\n                                                 labels=sample[\"labels\"])\n            image = transformed_sample[\"image\"]\n            sample[\"boxes\"] = torch.as_tensor(transformed_sample[\"bboxes\"], dtype=torch.float32)\n            del sample[\"image\"]\n        return image, sample\n\n    def __len__(self):\n        return len(self.images_list)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_dataset(arguments):\n    dataset = ObjectDetectionDataset(arguments.data_dir,\n                                     arguments.images_lists_dict[arguments.phase],\n                                     arguments.labels_csv_file_name,\n                                     arguments.phase,\n                                     arguments.transforms)\n    return dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_datasets_dictionary(arguments, input_size):\n    data_transforms = transform_set()\n    image_datasets = {\n        'train': None,\n        'val': None\n    }\n    for phase in ['train', 'val']:\n        arguments.phase = phase\n        arguments.transforms = data_transforms[phase]\n        image_datasets[phase] = create_dataset(arguments)\n    return image_datasets","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"## Create the pytorch dataloaders","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_dataloaders_dictionary(arguments, input_size):\n    batch_size = arguments.batch_size\n    num_workers = arguments.num_workers\n    image_datasets = create_datasets_dictionary(arguments.dataset_arguments, input_size)\n    dataloaders_dict = {x: DataLoader(image_datasets[x],\n                                      batch_size=batch_size,\n                                      shuffle=True,\n                                      pin_memory=True,\n                                      num_workers=num_workers,\n                                      collate_fn=collate_fn) for x in ['train', 'val']}\n    return dataloaders_dict","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Prepare the model\nI used the code here:\n[torchvision source code](https://github.com/pytorch/vision/blob/3d65fc6723f1e0709916f24d819d6e17a925b394/torchvision/models/detection/backbone_utils.py#L44)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def fasterrcnn_resnet101_fpn(pretrained=False, progress=True,\n                            num_classes=91, pretrained_backbone=True,\n                             trainable_backbone_layers=3, **kwargs):\n    assert trainable_backbone_layers <= 5 and trainable_backbone_layers >= 0\n    # dont freeze any layers if pretrained model or backbone is not used\n    if not (pretrained or pretrained_backbone):\n        trainable_backbone_layers = 5\n    if pretrained:\n        # no need to download the backbone if pretrained is set\n        pretrained_backbone = False\n    backbone = resnet_fpn_backbone('resnet152', pretrained_backbone)\n    model = FasterRCNN(backbone, num_classes, **kwargs)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def initialize_model():\n    model = fasterrcnn_resnet101_fpn(pretrained=False)\n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_training_device(train_device):\n    if train_device == Device.GPU:\n        if torch.cuda.is_available():\n            device = torch.device(\"cuda:0\")\n        else:\n            raise ValueError(\"No GPU was found\")\n    else:\n        device = torch.device(\"cpu\")\n    return device","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = get_training_device(train_device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = initialize_model()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = model.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataloaders = create_dataloaders_dictionary(dataloaders_arguments,input_size=None)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Some basic calculations that could be useful later","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset_size = len(dataloaders[\"train\"].dataset)\nnumber_of_iteration_per_epoch = int(train_dataset_size / dataloaders_arguments.batch_size)\ntotal_number_of_iteration = number_of_epochs * number_of_iteration_per_epoch\nlearning_rate_step_size = 2 * number_of_iteration_per_epoch","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Prepare for training","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_learnable_parameters(model, feature_extract):\n    params_to_update = model.parameters()\n\n    if feature_extract:\n        params_to_update = []\n        for name, param in model.named_parameters():\n            if param.requires_grad:\n                params_to_update.append(param)\n                print(\"\\t\", name)\n    return params_to_update","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"params_to_update = get_learnable_parameters(model, feature_extract=False)\noptimizer = optim.Adam(params_to_update, lr=learning_rate, weight_decay=weight_decay)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"lr_scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer,\n                                                              T_0=learning_rate_step_size,\n                                                              T_mult=1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _save_model(model, model_path):\n    torch.save(model, model_path)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Saving the checkpoint helps to starting training from a certain point.\ndef _save_checkpoint(epoch, model, optimizer, checkpoint_path):\n    torch.save({\n        'epoch': epoch,\n        'model_state_dict': model.state_dict(),\n        'optimizer_state_dict': optimizer.state_dict(),\n    }, checkpoint_path)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def save_model(epoch,model, optimizer):\n    model_path = f\"best_model_epoch_{epoch}.pth\"\n    _save_model(model.state_dict(), model_path)\n    checkpoint_path = f\"checkpoint_{epoch}.pth\"\n    _save_checkpoint(epoch, model, optimizer, checkpoint_path)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Detector:\n    def fit_model(self):\n        start_epoch = 0\n        iteration_i = 0\n        for current_epoch in range(start_epoch, number_of_epochs):\n            running_loss = 0\n            print(f\"Starting Epoch: {current_epoch}\")\n            progress_bar = tqdm(dataloaders[\"train\"])\n            for inputs, labels in  progress_bar:\n                running_loss_i = self.training_round(inputs, labels)\n                running_loss += running_loss_i\n                current_running_error = running_loss/((iteration_i - \n                                                      current_epoch * \n                                                      number_of_iteration_per_epoch + 1)*batch_size)\n                progress_bar.set_description(f\"Running train loss: {current_running_error}\")\n                iteration_i += 1\n            epoch_loss = running_loss / len(dataloaders[\"train\"].dataset)\n            print(f\"Finishing Current epoch: {current_epoch} ... training loss: {epoch_loss}\")\n            print(\"saving the model and checkpoint: \")\n            save_model(current_epoch, model, optimizer)\n            for inputs, labels in tqdm(dataloaders[\"val\"]):\n                self.validation_round(inputs, labels)\n\n    def training_round(self, inputs, labels):\n        inputs = list(image.to(device) for image in inputs)\n        inputs = torch.stack(inputs)\n        labels = [{k: v.to(device) for k, v in t.items() if not isinstance(v, str)} for t in labels]\n        model.train()\n        loss_dict = model(inputs, labels)\n        loss = sum(loss for loss in loss_dict.values())\n        loss.backward()\n        optimizer.step()\n        optimizer.zero_grad()\n        lr_scheduler.step()\n        running_loss_i = loss.item() * inputs.size(0) \n        return running_loss_i\n\n    def validation_round(self, inputs, labels):\n        model.eval()\n        inputs = list(image.to(device) for image in inputs)\n        inputs = torch.stack(inputs)\n        labels = [{k: v.to(device) for k, v in t.items() if not isinstance(v, str)} for t in labels]\n        outputs = model(inputs)\n        outputs = [{k: v.to(\"cpu\") for k, v in t.items()} for t in outputs]\n        # Note: I used here MSCOCO evaluation metric locally. Unfortunately, I could not run in this kernel.\n        # I appreciate it if you can help here\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Start training","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"detector  =  Detector()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"## The model will be saved for each epoch\ndetector.fit_model()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**I appreciate your feedback and upvote if you think it was useful**","execution_count":null},{"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}