def test_retina_fpn(self): for bbone in fpn_supported_models: backbone = retinanet.create_retinanet_backbone(backbone=bbone, pretrained=None) self.assertTrue(isinstance(backbone, nn.Module)) retina_model = retinanet.create_vision_retinanet(num_classes=3, backbone=backbone) self.assertTrue(isinstance(retina_model, nn.Module))
def test_infer(self): # Infer over an image image = Image.open("tests/assets/grace_hopper_517x606.jpg") tensor = im2tensor(image) self.assertEqual(tensor.ndim, 4) retina_model = retinanet.create_vision_retinanet() retina_model.eval() out = retina_model(tensor) self.assertIsInstance(out, list) self.assertIsInstance(out[0], Dict) self.assertIsInstance(out[0]["boxes"], torch.Tensor) self.assertIsInstance(out[0]["labels"], torch.Tensor) self.assertIsInstance(out[0]["scores"], torch.Tensor)
def test_val_sanity_fit(self): for bbone in fpn_supported_models: backbone = retinanet.create_retinanet_backbone(backbone=bbone, pretrained=None) self.assertTrue(isinstance(backbone, nn.Module)) retina_model = retinanet.create_vision_retinanet(num_classes=3, backbone=backbone) self.assertTrue(isinstance(retina_model, nn.Module)) result = retinanet.val_sanity_fit(retina_model, val_loader, "cpu", num_batches=10) self.assertTrue(result)
def test_train(self): # Read Image using PIL Here # Do forward over image image = Image.open("tests/assets/grace_hopper_517x606.jpg") img_tensor = im2tensor(image) self.assertEqual(img_tensor.ndim, 4) boxes = torch.tensor([[0, 0, 100, 100], [0, 1, 2, 2], [10, 15, 30, 35], [23, 35, 93, 95]], dtype=torch.float) labels = torch.tensor([1, 2, 3, 4], dtype=torch.int64) targets = [{"boxes": boxes, "labels": labels}] retina_model = retinanet.create_vision_retinanet(num_classes=5) out = retina_model(img_tensor, targets) self.assertIsInstance(out, Dict) self.assertIsInstance(out["classification"], torch.Tensor) self.assertIsInstance(out["bbox_regression"], torch.Tensor)
def test_val_step_fpn(self): for bbone in fpn_supported_models: backbone = retinanet.create_retinanet_backbone(backbone=bbone, pretrained=None) self.assertTrue(isinstance(backbone, nn.Module)) retina_model = retinanet.create_vision_retinanet(num_classes=3, backbone=backbone) self.assertTrue(isinstance(retina_model, nn.Module)) val_metrics = retinanet.val_step(retina_model, train_loader, "cpu", num_batches=4) 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_fit(self): for bbone in fpn_supported_models: backbone = retinanet.create_retinanet_backbone(backbone=bbone, pretrained=None) self.assertTrue(isinstance(backbone, nn.Module)) retina_model = retinanet.create_vision_retinanet(num_classes=3, backbone=backbone) self.assertTrue(isinstance(retina_model, nn.Module)) opt = torch.optim.SGD(retina_model.parameters(), lr=1e-3) history = retinanet.fit(retina_model, 1, train_loader, val_loader, "cpu", opt, num_batches=4) self.assertIsInstance(history, Dict) exp_keys = ("train", "val") for exp_k in exp_keys: self.assertTrue(exp_k in history.keys())
def test_train_step_fpn_cuda(self): for bbone in fpn_supported_models: backbone = retinanet.create_retinanet_backbone(backbone=bbone, pretrained=None) self.assertTrue(isinstance(backbone, nn.Module)) retina_model = retinanet.create_vision_retinanet(num_classes=3, backbone=backbone) self.assertTrue(isinstance(retina_model, nn.Module)) opt = torch.optim.SGD(retina_model.parameters(), lr=1e-3) train_metrics = retinanet.train_step(retina_model, train_loader, "cuda", opt, num_batches=4) self.assertIsInstance(train_metrics, Dict) exp_keys = ( "loss_classifier", "loss_box_reg", ) for exp_k in exp_keys: self.assertTrue(exp_k in train_metrics.keys())
val_metrics = faster_rcnn.val_step(model, test_loader, device="cuda") print("Training Metrics: ") pprint(train_metrics) print("Validation metrics") pprint(val_metrics) # Training Retina Net backbone = retinanet.create_retinanet_backbone("resnet50", fpn=True, pretrained="imagenet", trainable_backbone_layers=3) # our dataset has two classes only - background and person num_classes = 2 model = retinanet.create_vision_retinanet(num_classes=num_classes, backbone=backbone) # 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, lr=0.005, momentum=0.9, weight_decay=0.0005) # and a learning rate scheduler which decreases the learning rate by 10x every 3 epochs lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)