def test_fit_cuda(self):
for bbone in some_supported_backbones:
backbone = detr.create_detr_backbone(bbone, pretrained="coco")
self.assertTrue(isinstance(backbone, nn.Module))
detr_model = detr.create_vision_detr(num_classes=3,
num_queries=5,
backbone=backbone)
self.assertTrue(isinstance(detr_model, nn.Module))
matcher = detr_loss.HungarianMatcher()
weight_dict = {"loss_ce": 1, "loss_bbox": 1, "loss_giou": 1}
losses = ["labels", "boxes", "cardinality"]
opt = torch.optim.SGD(detr_model.parameters(), lr=1e-3)
criterion = detr_loss.SetCriterion(2,
matcher,
weight_dict,
eos_coef=0.5,
losses=losses)
history = detr.fit(detr_model,
1,
train_loader,
val_loader,
criterion,
device="cuda",
optimizer=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_step(self):
for bbone in some_supported_backbones:
backbone = detr.create_detr_backbone(bbone, pretrained=None)
self.assertTrue(isinstance(backbone, nn.Module))
detr_model = detr.create_vision_detr(num_classes=3,
num_queries=5,
backbone=backbone)
self.assertTrue(isinstance(detr_model, nn.Module))
matcher = detr_loss.HungarianMatcher()
weight_dict = {"loss_ce": 1, "loss_bbox": 1, "loss_giou": 1}
losses = ["labels", "boxes", "cardinality"]
criterion = detr_loss.SetCriterion(2,
matcher,
weight_dict,
eos_coef=0.5,
losses=losses)
met = detr.val_step(detr_model,
train_loader,
criterion,
"cpu",
num_batches=4)
self.assertIsInstance(met, Dict)
exp_keys = ("total_loss", "giou_loss", "bbox_loss", "labels_loss")
for exp_k in exp_keys:
self.assertTrue(exp_k in met.keys())
def test_train_step_cuda(self):
for bbone in some_supported_backbones:
backbone = detr.create_detr_backbone(bbone, pretrained=None)
self.assertTrue(isinstance(backbone, nn.Module))
detr_model = detr.create_vision_detr(num_classes=3,
num_queries=5,
backbone=backbone)
self.assertTrue(isinstance(detr_model, nn.Module))
opt = torch.optim.SGD(detr_model.parameters(), lr=1e-3)
scaler = amp.GradScaler()
matcher = detr_loss.HungarianMatcher()
weight_dict = {"loss_ce": 1, "loss_bbox": 1, "loss_giou": 1}
losses = ["labels", "boxes", "cardinality"]
criterion = detr_loss.SetCriterion(2,
matcher,
weight_dict,
eos_coef=0.5,
losses=losses)
met = detr.train_step(detr_model,
train_loader,
criterion,
"cuda",
opt,
num_batches=4,
scaler=scaler)
self.assertIsInstance(met, Dict)
exp_keys = ("total_loss", "giou_loss", "bbox_loss", "labels_loss")
for exp_k in exp_keys:
self.assertTrue(exp_k in met.keys())
def test_create_vision_detr(self):
for supp_bb in supported_detr_backbones:
bbone = detr.create_detr_backbone(supp_bb, pretrained=None)
self.assertTrue(isinstance(bbone, nn.Module))
model = detr.create_vision_detr(num_classes=91,
num_queries=5,
backbone=bbone)
self.assertTrue(isinstance(bbone, nn.Module))
def __init__(
self,
learning_rate: float = 1e-3,
num_classes: int = 91,
num_queries: int = 5,
pretrained: str = None,
backbone: str = "resnet50",
**kwargs,
):
"""
PyTorch Lightning implementation of `Detr: End-to-End Object Detection with Transformers
During training, the model expects both the input tensors, as well as targets (list of dictionary),
containing:
- boxes (`FloatTensor[N, 4]`): the ground truth boxes in `[x1, y1, x2, y2]` format.
- labels (`Int64Tensor[N]`): the class label for each ground truh box
Args:
learning_rate: the learning rate
num_classes: number of detection classes (including background)
num_queries: number of queries to the transformer module.
pretrained: if "coco", returns a model pre-trained on COCO train2017
backbone: Supported Detection backbones are "resnet50", "resnet101", "resnet50_dc5", "resnet101_dc5".
It returns a dict with the following elements:
- "pred_logits": the classification logits (including no-object) for all queries.
Shape= [batch_size x num_queries x (num_classes + 1)]
- "pred_boxes": The normalized boxes coordinates for all queries, represented as
(center_x, center_y, height, width). These values are normalized in [0, 1],
relative to the size of each individual image (disregarding possible padding).
See PostProcess for information on how to retrieve the unnormalized bounding box.
"""
super().__init__()
self.model = create_detr_backbone(backbone, pretrained)
in_features = self.model.class_embed.in_features
self.model.class_embed = nn.Linear(in_features=in_features,
out_features=num_classes)
self.model.num_queries = num_queries
self.learning_rate = learning_rate
matcher = detr_loss.HungarianMatcher()
weight_dict = {"loss_ce": 1, "loss_bbox": 1, "loss_giou": 1}
losses = ['labels', 'boxes', 'cardinality']
self.criterion = detr_loss.SetCriterion(num_classes - 1,
matcher,
weight_dict,
eos_coef=0.5,
losses=losses)
def test_val_sanity_fit(self):
for bbone in some_supported_backbones:
backbone = detr.create_detr_backbone(bbone, pretrained=None)
self.assertTrue(isinstance(backbone, nn.Module))
detr_model = detr.create_vision_detr(num_classes=3,
num_queries=5,
backbone=backbone)
self.assertTrue(isinstance(detr_model, nn.Module))
matcher = detr_loss.HungarianMatcher()
weight_dict = {"loss_ce": 1, "loss_bbox": 1, "loss_giou": 1}
losses = ["labels", "boxes", "cardinality"]
criterion = detr_loss.SetCriterion(2,
matcher,
weight_dict,
eos_coef=0.5,
losses=losses)
ret = detr.val_sanity_fit(detr_model,
val_loader,
criterion,
"cpu",
num_batches=4)
self.assertTrue(ret)
def test_detr_backbones(self):
for supp_bb in supported_detr_backbones:
bbone = detr.create_detr_backbone(supp_bb, pretrained=None)
self.assertTrue(isinstance(bbone, nn.Module))
