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_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_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 __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)