mirror of
https://github.com/Laurent2916/REVA-QCAV.git
synced 2024-11-09 15:02:03 +00:00
fix: limit box detections per images to 10
Former-commit-id: 3b49b7b2c9d5093ab989e8757c6b0429eb9746bd [formerly abf1897a28a7591c6fdf9cf51b9c832ca48cc10b] Former-commit-id: c1db1fe1e9475adbdf10f067db9206d69fe5ed2f
This commit is contained in:
parent
f50b758102
commit
c6c08ac98a
|
@ -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",
|
||||||
|
# },
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in a new issue