def test_frcnn_nonfpn(self): for bbone in non_fpn_supported_models: backbone = faster_rcnn.create_fastercnn_backbone(backbone=bbone, pretrained=None, fpn=False) self.assertTrue(isinstance(backbone, nn.Module)) frcnn_model = faster_rcnn.create_vision_fastercnn(num_classes=3, backbone=backbone) self.assertTrue(isinstance(frcnn_model, nn.Module))
def test_sanity_fit_cuda(self): for bbone in fpn_supported_models: backbone = faster_rcnn.create_fastercnn_backbone(backbone=bbone, pretrained=None) self.assertTrue(isinstance(backbone, nn.Module)) frcnn_model = faster_rcnn.create_vision_fastercnn(num_classes=3, backbone=backbone) self.assertTrue(isinstance(frcnn_model, nn.Module)) result = faster_rcnn.sanity_fit(frcnn_model, train_loader, val_loader, "cuda", num_batches=10, fp16=True) self.assertTrue(result)
def test_val_step_fpn(self): for bbone in fpn_supported_models: backbone = faster_rcnn.create_fastercnn_backbone(backbone=bbone, pretrained=None) self.assertTrue(isinstance(backbone, nn.Module)) frcnn_model = faster_rcnn.create_vision_fastercnn(num_classes=3, backbone=backbone) self.assertTrue(isinstance(frcnn_model, nn.Module)) val_metrics = faster_rcnn.val_step(frcnn_model, train_loader, "cpu", num_batches=10) self.assertIsInstance(val_metrics, Dict) exp_keys = ("iou", "giou") for exp_k in exp_keys: self.assertTrue(exp_k in val_metrics.keys())
def test_train_step_fpn(self): for bbone in fpn_supported_models: backbone = faster_rcnn.create_fastercnn_backbone(backbone=bbone, pretrained=None) self.assertTrue(isinstance(backbone, nn.Module)) frcnn_model = faster_rcnn.create_vision_fastercnn(num_classes=3, backbone=backbone) self.assertTrue(isinstance(frcnn_model, nn.Module)) opt = torch.optim.SGD(frcnn_model.parameters(), lr=1e-3) train_metrics = faster_rcnn.train_step(frcnn_model, train_loader, "cpu", opt, num_batches=10) self.assertIsInstance(train_metrics, Dict) exp_keys = ("loss_classifier", "loss_box_reg", "loss_objectness", "loss_rpn_box_reg") for exp_k in exp_keys: self.assertTrue(exp_k in train_metrics.keys())
def test_fit_cuda(self): for bbone in fpn_supported_models: backbone = faster_rcnn.create_fastercnn_backbone(backbone=bbone, pretrained=None) self.assertTrue(isinstance(backbone, nn.Module)) frcnn_model = faster_rcnn.create_vision_fastercnn(num_classes=3, backbone=backbone) self.assertTrue(isinstance(frcnn_model, nn.Module)) opt = torch.optim.SGD(frcnn_model.parameters(), lr=1e-3) history = faster_rcnn.fit(frcnn_model, 1, train_loader, val_loader, "cuda", opt, num_batches=4, fp16=True) self.assertIsInstance(history, Dict) exp_keys = ("train", "val") for exp_k in exp_keys: self.assertTrue(exp_k in history.keys())
def test_val_sanity_fit(self): for bbone in fpn_supported_models: backbone = faster_rcnn.create_fastercnn_backbone(backbone=bbone, pretrained=None) self.assertTrue(isinstance(backbone, nn.Module)) frcnn_model = faster_rcnn.create_fastercnn(num_classes=3, backbone=backbone) frcnn_model = frcnn_model.cpu() self.assertTrue(isinstance(frcnn_model, nn.Module)) result = faster_rcnn.val_sanity_fit(frcnn_model, val_loader, "cpu", num_batches=10) self.assertTrue(result)
def __init__( self, learning_rate: float = 0.0001, num_classes: int = 91, backbone: str = None, fpn: bool = True, pretrained_backbone: str = None, trainable_backbone_layers: int = 3, **kwargs, ): """ Args: learning_rate: the learning rate num_classes: number of detection classes (including background) pretrained: if true, returns a model pre-trained on COCO train2017 pretrained_backbone (str): if "imagenet", returns a model with backbone pre-trained on Imagenet trainable_backbone_layers: number of trainable resnet layers starting from final block """ super().__init__() self.learning_rate = learning_rate self.num_classes = num_classes self.backbone = backbone if backbone is None: self.model = fasterrcnn_resnet50_fpn( pretrained=True, trainable_backbone_layers=trainable_backbone_layers, ) in_features = self.model.roi_heads.box_predictor.cls_score.in_features self.model.roi_heads.box_predictor = FastRCNNPredictor( in_features, self.num_classes) else: backbone_model = create_fastercnn_backbone( self.backbone, fpn, pretrained_backbone, trainable_backbone_layers, **kwargs, ) self.model = FasterRCNN(backbone_model, num_classes=num_classes, **kwargs)
batch_size=1, shuffle=False, num_workers=4, collate_fn=collate_fn) # Now let's instantiate the model and the optimizer # And now let's train the model for 10 epochs, evaluating at the end of every epoch. # Quickvision to Train Detection Models ! # Training Faster RCNN backbone = faster_rcnn.create_fastercnn_backbone( "resnet101", fpn=True, pretrained="imagenet", trainable_backbone_layers=3) # our dataset has two classes only - background and person num_classes = 2 model = faster_rcnn.create_vision_fastercnn(num_classes=num_classes, backbone=backbone) # - Quickvision supports Mixed Precision training as well ! # Let's use Mixed Precision training scaler = amp.GradScaler() params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params,