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_detr(num_classes=3,
num_queries=5,
backbone=backbone)
detr_model = detr_model.cpu()
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_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_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_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_detr(num_classes=91,
num_queries=5,
backbone=bbone)
self.assertTrue(isinstance(bbone, nn.Module))
def test_sanity_fit_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_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.sanity_fit(detr_model,
train_loader,
val_loader,
criterion,
"cuda",
num_batches=4)
self.assertTrue(ret)