From c6c08ac98a4725ce9a421d631ba66c0e3ed39ab7 Mon Sep 17 00:00:00 2001
From: Laurent Fainsin <laurentfainsin@protonmail.com>
Date: Mon, 12 Sep 2022 09:26:42 +0200
Subject: [PATCH] fix: limit box detections per images to 10

Former-commit-id: 3b49b7b2c9d5093ab989e8757c6b0429eb9746bd [formerly abf1897a28a7591c6fdf9cf51b9c832ca48cc10b]
Former-commit-id: c1db1fe1e9475adbdf10f067db9206d69fe5ed2f
---
 src/mrcnn/module.py | 29 +++++++++++++++++------------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/src/mrcnn/module.py b/src/mrcnn/module.py
index 21a7720..fcee13b 100644
--- a/src/mrcnn/module.py
+++ b/src/mrcnn/module.py
@@ -15,7 +15,10 @@ from torchvision.models.detection.mask_rcnn import (
 
 def get_model_instance_segmentation(num_classes):
     # load an instance segmentation model pre-trained on COCO
-    model = torchvision.models.detection.maskrcnn_resnet50_fpn(weights=MaskRCNN_ResNet50_FPN_Weights.DEFAULT)
+    model = torchvision.models.detection.maskrcnn_resnet50_fpn(
+        weights=MaskRCNN_ResNet50_FPN_Weights.DEFAULT,
+        box_detections_per_img=10,  # cap numbers of detections, else memory explosion
+    )
 
     # get number of input features for the classifier
     in_features = model.roi_heads.box_predictor.cls_score.in_features
@@ -97,20 +100,22 @@ class MRCNNModule(pl.LightningModule):
         optimizer = torch.optim.Adam(
             self.parameters(),
             lr=wandb.config.LEARNING_RATE,
-            momentum=wandb.config.MOMENTUM,
-            weight_decay=wandb.config.WEIGHT_DECAY,
+            # momentum=wandb.config.MOMENTUM,
+            # weight_decay=wandb.config.WEIGHT_DECAY,
         )
 
-        scheduler = LinearWarmupCosineAnnealingLR(
-            optimizer,
-            warmup_epochs=10,
-            max_epochs=40,
-        )
+        # scheduler = LinearWarmupCosineAnnealingLR(
+        #     optimizer,
+        #     warmup_epochs=1,
+        #     max_epochs=30,
+        # )
 
         return {
             "optimizer": optimizer,
-            "lr_scheduler": {
-                "scheduler": scheduler,
-                "monitor": "map",
-            },
+            # "lr_scheduler": {
+            #     "scheduler": scheduler,
+            #     "interval": "step",
+            #     "frequency": 10,
+            #     "monitor": "bbox/map",
+            # },
         }