{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Pytorch starter - MaskRCNN Train\nIn this notebook I enabled the GPU and the Internet access (needed for the pre-trained weights). We can not use Internet during inference, so I'll create another notebook for commiting. Stay tuned!\n\nYou can find the [inference notebook here](https://www.kaggle.com/pestipeti/pytorch-starter-fasterrcnn-inference)\n\n- FasterRCNN from torchvision\n- Use Resnet50 backbone\n- Albumentation enabled (simple flip for now)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport cv2\nimport os\nimport re\n\nfrom PIL import Image\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nimport torch\nimport torchvision\n\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN\nfrom torchvision.models.detection.rpn import AnchorGenerator\n\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.sampler import SequentialSampler\n\nfrom matplotlib import pyplot as plt\n\nDIR_INPUT = '/kaggle/input/global-wheat-detection'\nDIR_TRAIN = f'{DIR_INPUT}/train'\nDIR_TEST = f'{DIR_INPUT}/test'\n\nprint(\"Done with initial imports\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import numba\nfrom numba import jit\n\n@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    \n    if dx < 0:\n        return 0.0\n    \n    dy = min(gt[3], pr[3]) - max(gt[1], pr[1]) + 1\n\n    if dy < 0:\n        return 0.0\n\n    overlap_area = dx * dy\n\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\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\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\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\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\nprint(\"Done with mAP definition\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(f'{DIR_INPUT}/train.csv')\nprint(train_df.shape)\nprint(train_df)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df['x'] = -1\ntrain_df['y'] = -1\ntrain_df['w'] = -1\ntrain_df['h'] = -1\n\nprint(train_df)\n\ndef expand_bbox(x):\n    r = np.array(re.findall(\"([0-9]+[.]?[0-9]*)\", x))\n    if len(r) == 0:\n        r = [-1, -1, -1, -1]\n    return r\n\ntrain_df[['x', 'y', 'w', 'h']] = np.stack(train_df['bbox'].apply(lambda x: expand_bbox(x)))\ntrain_df.drop(columns=['bbox'], inplace=True)\ntrain_df['x'] = train_df['x'].astype(np.float)\ntrain_df['y'] = train_df['y'].astype(np.float)\ntrain_df['w'] = train_df['w'].astype(np.float)\ntrain_df['h'] = train_df['h'].astype(np.float)\n\nprint(train_df)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_ids = train_df['image_id'].unique()\nvalid_ids = image_ids[-665:]\ntrain_ids = image_ids[:-665]\n# We pick 665 images for validation I think\nprint(valid_ids[:20])\nprint(train_ids)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"valid_df = train_df[train_df['image_id'].isin(valid_ids)]\ntrain_df = train_df[train_df['image_id'].isin(train_ids)]\nprint(valid_df.shape, train_df.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Added extra stuff here\nclass WheatDataset(Dataset):\n\n    def __init__(self, dataframe, image_dir, transforms=None):\n        super().__init__()\n\n        self.image_ids = dataframe['image_id'].unique()\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.transforms = transforms\n\n    def __getitem__(self, index: int):\n\n        image_id = self.image_ids[index]\n        records = self.df[self.df['image_id'] == image_id]\n\n        image = cv2.imread(f'{self.image_dir}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n\n        # Define bbox coordinates\n        boxes = records[['x', 'y', 'w', 'h']].values\n        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n        \n        # Define areas of bboxes\n        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n        area = torch.as_tensor(area, dtype=torch.float32)\n\n        # there is only one class\n        labels = torch.ones((records.shape[0],), dtype=torch.int64)\n        \n        # suppose all instances are not crowd\n        iscrowd = torch.zeros((records.shape[0],), dtype=torch.int64)\n        \n        target = {}\n        target['boxes'] = boxes\n        target['labels'] = labels\n        target['image_id'] = torch.tensor([index])\n        target['area'] = area\n        target['iscrowd'] = iscrowd\n\n        if self.transforms:\n            sample = {\n                'image': image,\n                'bboxes': target['boxes'],\n                'labels': labels\n            }\n            sample = self.transforms(**sample)\n            image = sample['image']\n            \n            target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)\n\n#        # Define elliptic mask based in boxes AFTER augmentation\n#         im_H, im_W = image.shape[1], image.shape[2] ##\n#         bcoords = target['boxes'] ##\n#         mask = np.zeros((bcoords.shape[0], im_H, im_W)) ##\n#         \n#         for i in range(bcoords.shape[0]):\n#             mask[i,:,:] = get_mask(im_H, im_W, bcoords[i, :]) ##\n#         mask = torch.tensor(mask, dtype = torch.uint8) ##\n#        # target['masks'] needs to be a UInt8Tensor[N, H, W] with N=no of bboxes.\n#        # The uint8 values need to be binary, 1 or 0.\n#        \n#        target['masks'] = mask ##\n            \n        return image, target, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n    \nprint('Done with dataset loader...')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Albumentations\ndef get_train_transform():\n    return A.Compose([\n        A.Flip(0.5),\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\ndef get_valid_transform():\n    return A.Compose([\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\nprint('Done with augmentation definitions...')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Call the Mask R-CNN model from torchvision with pretrained weights","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# load a model; pre-trained on COCO\nmodel = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"num_classes = 2  # 1 class (wheat) + background\n\n# get number of input channels for the final linear classifier\nin_features = model.roi_heads.box_predictor.cls_score.in_features\n\n# replace the pre-trained head with a new one\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n\nprint(in_features)\nprint('Done with initializing model...')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(model)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Averager:\n    def __init__(self):\n        self.current_total = 0.0\n        self.iterations = 0.0\n\n    def send(self, value):\n        self.current_total += value\n        self.iterations += 1\n\n    @property\n    def value(self):\n        if self.iterations == 0:\n            return 0\n        else:\n            return 1.0 * self.current_total / self.iterations\n\n    def reset(self):\n        self.current_total = 0.0\n        self.iterations = 0.0\n\nprint('Done with useful functions...')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))\n\ntrain_dataset = WheatDataset(train_df, DIR_TRAIN, get_train_transform())\nvalid_dataset = WheatDataset(valid_df, DIR_TRAIN, get_valid_transform())\n\n\n# split the dataset in train and test set\nindices = torch.randperm(len(train_dataset)).tolist()\n\ntrain_data_loader = DataLoader(\n    train_dataset,\n    batch_size=12,\n    shuffle=False,\n    num_workers=0,\n    collate_fn=collate_fn\n) ## changed batch size from 16\n\nvalid_data_loader = DataLoader(\n    valid_dataset,\n    batch_size=8,\n    shuffle=False,\n    num_workers=4,\n    collate_fn=collate_fn\n)\nprint('Done with data split...')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\nprint('GPU?: '+str(torch.cuda.is_available()))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Load and get fft masks\nall_ffts=np.load('../input/wheat-mean-ffts-200-images/wheat_detection_mean_ffts_200_images.npz')\nwheat_freq=all_ffts['wheat_freq']\nim_freq=all_ffts['im_freq']\npad2=1024\n\nthr_list=[0.2, 0.4, 0.6, 0.8]\nn_thr=len(thr_list)\nmask=np.zeros((pad2,pad2,3,n_thr))\nfor i in range(3):\n    plot_wheat=np.log(wheat_freq[:,:,i])\n    plot_im=np.log(im_freq[:,:,i])\n\n#     print(np.min(plot_wheat))\n#     print(np.min(plot_im))\n\n    # if printed values all positive\n    if np.min(plot_wheat)>0 and np.min(plot_im)>0:\n        plot_wheat[0,:]=0\n        plot_wheat[:,0]=0\n        plot_im[0,:]=0\n        plot_im[:,0]=0\n\n    plot_wheat=plot_wheat/np.sum(np.abs(plot_wheat))\n    plot_im=plot_im/np.sum(np.abs(plot_im))\n\n    fft_diff=plot_wheat-plot_im\n    fft_diff=np.fft.fftshift(fft_diff)\n\n    for kt in range(n_thr):\n        f_thr=thr_list[kt]\n        mask[:,:,i,kt]=fft_diff>f_thr*1e-7\n        exclude=120\n        mask[:exclude,:,:,kt],mask[-exclude:,:,:,kt]=0,0\n        mask[:,:exclude,:,kt],mask[:,-exclude:,:,kt]=0,0\n        \n        plt.figure()\n        plt.imshow(mask[:,:,i,kt],vmin=-0.2e-7,vmax=2e-7)\n        plt.colorbar()\n        plt.pause(0.1)\n\n        mask[:,:,i,kt]=np.fft.fftshift(mask[:,:,i,kt])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model=model.to(device)\nparams = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)\n# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)\nlr_scheduler = None\n\nnum_epochs = 6\n\n# Initialize metrics (except mask loss)\nbest_map = 0\nmetrics = {'mean_ap' : [],\n           'loss_cls' : [],\n           'loss_reg' : [],\n           'loss_obj' : [],\n           'loss_rpn' : [],\n           'loss_tot' : []}\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)\n\nprint('Done with optim init...')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def mask_input(images,thr_list,pad2):\n    new_images=[]\n    n_thr=len(thr_list)\n    im_masked=np.zeros((3,pad2,pad2))\n    for image in images:\n        # get fft of full image in 3 color channels and mask\n        for j in range(3):\n            for kt in range(n_thr):\n                im_masked[j,:,:]+=np.real(np.fft.ifft2(np.fft.fft2(image[j,:,:])*mask[:,:,j,kt]))\n        # normalize to 0-1\n        im_masked=im_masked-np.min(im_masked)\n        im_masked=im_masked/np.max(im_masked)\n        \n        new_images.append(torch.from_numpy(im_masked).float())\n    return new_images        ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#print(images[0].dtype,new_images[0].dtype)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Train\nloss_hist=Averager()\nloss_cls=Averager()\nloss_reg=Averager()\nloss_obj=Averager()\nloss_rpn=Averager()\n\nitr = 1\n\nfor epoch in range(num_epochs):\n    loss_hist.reset()\n    loss_cls.reset()\n    loss_reg.reset()\n    loss_obj.reset()\n    loss_rpn.reset()\n    image_precisions=[]\n    \n    for images, targets, image_ids in train_data_loader:\n        model.train()\n        # images and targets are both tuples of length batch_size each\n    \n        # mask the input images\n        images=mask_input(images,thr_list,pad2)\n    \n        images = list(image.to(device) for image in images)\n        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n        # images and targets are now both lists of length batch_size each, but sent to device\n        \n        print(images[0].dtype)\n\n        loss_dict = model(images, targets)\n\n        losses = sum(loss for loss in loss_dict.values())\n        loss_value = losses.item()\n\n        optimizer.zero_grad()\n        losses.backward()\n        optimizer.step()\n\n        # Store metrics\n        loss_hist.send(loss_value)\n        losses_batch = np.array([v.cpu().detach().item() for v in loss_dict.values()]) ##\n        loss_cls.send(losses_batch[0]) ##\n        loss_reg.send(losses_batch[1])  ##\n        loss_obj.send(losses_batch[2]) ##\n        loss_rpn.send(losses_batch[3]) ##\n        \n        # Evaluate again to get the APs\n        model.eval()\n        with torch.no_grad():\n            pred = model(images)\n        for idx, image in enumerate(images):\n            # Remove detections below threshold\n            pred[idx]['boxes'] = pred[idx]['boxes'][pred[idx]['scores'] >= 0.5]\n            pred[idx]['scores'] = pred[idx]['scores'][pred[idx]['scores'] >= 0.5]\n\n            preds = pred[idx]['boxes'].cpu().detach().numpy()\n            scores = pred[idx]['scores'].cpu().detach().numpy()\n\n            gt_boxes = targets[idx]['boxes'].cpu().numpy()\n\n            preds_sorted_idx = np.argsort(scores)[::-1]\n            preds_sorted = preds[preds_sorted_idx]\n\n            AP = calculate_image_precision(preds_sorted,\n                                            gt_boxes,\n                                            thresholds=iou_thresholds,\n                                            form='pascal_voc')\n            image_precisions.append(AP)\n        del pred, image, preds, scores, gt_boxes, preds_sorted_idx, preds_sorted, AP ## for CUDA memory\n\n        if itr % 40 == 0:\n            print(f\"Iteration #{itr} loss: {loss_value}\")\n\n        itr += 1\n\n    metrics['mean_ap'].append(np.mean(image_precisions)) ##\n    metrics['loss_cls'].append(loss_cls.value) ##\n    metrics['loss_reg'].append(loss_reg.value) ##\n    metrics['loss_obj'].append(loss_obj.value) ##\n    metrics['loss_rpn'].append(loss_rpn.value) ##\n    metrics['loss_tot'].append(loss_hist.value) ##\n\n#         # Save best model to file\n#         if mean_ap > best_map:\n#             print('mAP improved from {:.4f} to {:.4f}.'.format(best_map, mean_ap))\n#             best_map = mean_ap\n#             torch.save({\n#                 'epoch': epoch,\n#                 'model_state_dict': model.state_dict(),\n#                 }, 'output/best_weights.pth.tar')\n    # update the learning rate\n    if lr_scheduler is not None:\n        lr_scheduler.step()\n\n    print(f\"Epoch #{epoch} loss: {loss_hist.value}\")\n    print(f\"Epoch #{epoch} metrics: {metrics}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(image_precisions)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(metrics)\nprint(epoch)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Get val images\nimages, targets, image_ids = next(iter(valid_data_loader))\n\n# Send input im and target to device\nimages = list(img.to(device) for img in images)\ntargets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n\n# Get prediction\nmodel.eval()\ncpu_device = torch.device(\"cpu\")\n\noutputs = model(images)\noutputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"im_no=3\n\n# Get ground truth for this image\nboxes = targets[im_no]['boxes'].cpu().numpy().astype(np.int32)\n# masks = targets[im_no]['masks'].cpu().numpy().astype(np.int32)\nsample = images[im_no].permute(1,2,0).cpu().numpy()\n\n# Get prediction values for this image\nboxes_out = outputs[im_no]['boxes'].detach().numpy().astype(np.int32)\nscores_out = outputs[im_no]['scores'].detach().numpy()\n# masks_out = outputs[im_no]['masks'].detach().numpy()\nprint(scores_out)\n# print(masks.sum())\n# print(masks_out.sum())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\ni, j = 0, 0\nfor box in boxes:\n    i+=1\n    cv2.rectangle(sample,\n                  (box[0], box[1]),\n                  (box[2], box[3]),\n                  (220, 0, 0), 3)\nfor box in boxes_out:\n    j+=1\n    cv2.rectangle(sample,\n                  (box[0], box[1]),\n                  (box[2], box[3]),\n                  (0, 50, 250), 3)\n\n#ax.set_axis_off()\nax.imshow(sample)\n\nprint(i,j)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save({'epoch': epoch,'model_state_dict': model.state_dict(),'metrics': metrics},'fft_faster_rcnn_weights_w_metric_6_epochs.pth.tar')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"raw","source":"Notes:\n- Need to run for longer epochs\n- Should consider LR scheduler","execution_count":null}],"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}