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)
