Update modeling_diffusiondet.py

modeling_diffusiondet.py

@@ -238,7 +238,7 @@ class DiffusionDet(PreTrainedModel):
 
         return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise
 
-    def forward(self, pixel_values, labels):
+    def forward(self, pixel_values, labels=None):
         """
         Args:
         """
@@ -256,6 +256,16 @@ class DiffusionDet(PreTrainedModel):
         features = self.fpn(features)  # [144, 72, 36, 18]
         features = [features[f] for f in features.keys()]
 
+        if not self.training:
+            pred_logits, pred_labels, pred_boxes = self.ddim_sample(pixel_values, features, images_whwh)
+            return DiffusionDetOutput(
+                # loss=loss_dict['loss'],
+                # loss_dict=loss_dict,
+                logits=pred_logits,
+                labels=pred_labels,
+                pred_boxes=pred_boxes,
+            )
+
         # if self.training:
         labels = list(map(lambda tensor: tensor.to(self.device), labels))
         targets, x_boxes, noises, ts = self.prepare_targets(labels)
@@ -277,23 +287,6 @@ class DiffusionDet(PreTrainedModel):
                 loss_dict[k] *= weight_dict[k]
         loss_dict['loss'] = sum([loss_dict[k] for k in weight_dict.keys()])
 
-        wandb_logs_values = ["loss_ce", "loss_bbox", "loss_giou"]
-
-        if self.training:
-            wandb.log({f'train/{k}': v.detach().cpu().numpy() for k, v in loss_dict.items() if k in wandb_logs_values})
-        else:
-            wandb.log({f'eval/{k}': v.detach().cpu().numpy() for k, v in loss_dict.items() if k in wandb_logs_values})
-
-        if not self.training:
-            pred_logits, pred_labels, pred_boxes  = self.ddim_sample(pixel_values, features, images_whwh)
-            return DiffusionDetOutput(
-                loss=loss_dict['loss'],
-                loss_dict=loss_dict,
-                logits=pred_logits,
-                labels=pred_labels,
-                pred_boxes=pred_boxes,
-            )
-
         return DiffusionDetOutput(
             loss=loss_dict['loss'],
             loss_dict=loss_dict,