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