ENSEEIHT/projet-long
1
0
Fork 0
Code Issues Pull requests Packages Projects Releases Wiki Activity

fix: limit box detections per images to 10

Browse source 
Former-commit-id: 3b49b7b2c9d5093ab989e8757c6b0429eb9746bd [formerly abf1897a28a7591c6fdf9cf51b9c832ca48cc10b]
Former-commit-id: c1db1fe1e9475adbdf10f067db9206d69fe5ed2f
This commit is contained in:
Laurent Fainsin 2022-09-12 09:26:42 +02:00
parent f50b758102
commit c6c08ac98a
1 changed files with 17 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
src/mrcnn/module.py
View file

@ -15,7 +15,10 @@ from torchvision.models.detection.mask_rcnn import (
def get_model_instance_segmentation(num_classes): def get_model_instance_segmentation(num_classes):
# load an instance segmentation model pre-trained on COCO # 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 # get number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features in_features = model.roi_heads.box_predictor.cls_score.in_features
@ -97,20 +100,22 @@ class MRCNNModule(pl.LightningModule):
optimizer = torch.optim.Adam( optimizer = torch.optim.Adam(
self.parameters(), self.parameters(),
lr=wandb.config.LEARNING_RATE, lr=wandb.config.LEARNING_RATE,
momentum=wandb.config.MOMENTUM, # momentum=wandb.config.MOMENTUM,
weight_decay=wandb.config.WEIGHT_DECAY, # weight_decay=wandb.config.WEIGHT_DECAY,
) )
scheduler = LinearWarmupCosineAnnealingLR( # scheduler = LinearWarmupCosineAnnealingLR(
optimizer, # optimizer,
warmup_epochs=10, # warmup_epochs=1,
max_epochs=40, # max_epochs=30,
) # )
return { return {
"optimizer": optimizer, "optimizer": optimizer,
"lr_scheduler": { # "lr_scheduler": {
"scheduler": scheduler, # "scheduler": scheduler,
"monitor": "map", # "interval": "step",
}, # "frequency": 10,
# "monitor": "bbox/map",
# },
} }