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 = retinanet_resnet50_fpn(pretrained=True, **kwargs)
self.model.head = RetinaNetHead(in_channels=self.model.backbone.out_channels,
num_anchors=self.model.head.classification_head.num_anchors,
num_classes=num_classes, **kwargs)
else:
backbone_model = create_retinanet_backbone(self.backbone, fpn, pretrained_backbone,
trainable_backbone_layers, **kwargs)
self.model = RetinaNet(backbone_model, num_classes=num_classes, **kwargs)
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_retinanet(num_classes=3,
backbone=backbone)
self.assertTrue(isinstance(retina_model, nn.Module))
def test_train_sanity_fit_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_retinanet(num_classes=3,
backbone=backbone)
self.assertTrue(isinstance(retina_model, nn.Module))
result = retinanet.train_sanity_fit(retina_model,
train_loader,
"cuda",
num_batches=10)
self.assertTrue(result)
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_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())
train_loader,
device="cuda",
optimizer=optimizer,
scheduler=lr_scheduler,
log_interval=100,
scaler=scaler)
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)
