test + metrics

Author: humanpose1

test/test_confusion_matrix.py

@@ -11,10 +11,10 @@
 sys.path.insert(0, ROOT)
 sys.path.append('.')
 
-from torch_points3d.metrics.segmentation.segmentation_tracker import compute_intersection_union_per_class
-from torch_points3d.metrics.segmentation.segmentation_tracker import compute_average_intersection_union
-from torch_points3d.metrics.segmentation.segmentation_tracker import compute_overall_accuracy
-from torch_points3d.metrics.segmentation.segmentation_tracker import compute_mean_class_accuracy
+from torch_points3d.metrics.segmentation.metrics import compute_intersection_union_per_class
+from torch_points3d.metrics.segmentation.metrics import compute_average_intersection_union
+from torch_points3d.metrics.segmentation.metrics import compute_overall_accuracy
+from torch_points3d.metrics.segmentation.metrics import compute_mean_class_accuracy
 
 
 

torch_points3d/metrics/base_tracker.py

@@ -21,9 +21,9 @@ def track(self, output_model, *args, **kwargs) -> Dict[str, Any]:
     def track_loss(self, losses: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
         out_loss = dict()
         for key, loss in losses.items():
-            loss_key = "%s_%s" % (self.stage, key)
-            if loss_key not in self.loss_metrics:
-                self.loss_metrics[loss_key] = AverageMeter()
+            loss_key = f"{self.stage}_{key}"
+            if loss_key not in self.loss_metrics.keys():
+                self.loss_metrics[loss_key] = AverageMeter().to(loss)
             val = self.loss_metrics[loss_key](loss)
             out_loss[loss_key] = val
         return out_loss

torch_points3d/tasks/base_model.py

@@ -98,7 +98,7 @@ def configure_metrics(self, stage: str) -> Optional[Any]:
     def _step(self, batch, batch_idx, stage: str):
         self.model.set_input(batch)
         loss = self.model.forward()
-        losses = self.model.get_losses(stage=stage)
+        losses = self.model.get_losses()
         outputs = self.model.get_outputs()
         metric_dict = self.tracker(outputs, losses)
         self.log_dict(metric_dict, prog_bar=True, on_step=False, on_epoch=True)