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,
