update

Author: clee-ai

conf/dataset/default.yaml

@@ -1,8 +1,5 @@
 # @package dataset
-# cfg:
-  # torch data-loader specific arguments
 cfg: 
-  num_classes: 
   feature_dimension: 
   batch_size: ${training.batch_size}
   num_workers: ${training.num_workers}

conf/dataset/segmentation/default.yaml

@@ -1,3 +1,6 @@
 # @package dataset
 defaults:
-  - /dataset/default
+  - /dataset/default
+
+cfg:
+  num_classes: 

torch_points3d/models/base_model.py

@@ -1,4 +1,8 @@
+from typing import Dict, Union
+
+import torch
 import torch.nn as nn
+from torch_geometric.data import Data
 
 from torch_points3d.core.instantiator import Instantiator
 
@@ -9,8 +13,11 @@ def __init__(self, instantiator: Instantiator):
 
         self.instantiator = instantiator
 
-    def set_input(self, data):
+    def set_input(self, data: Data) -> None:
         raise (NotImplementedError("set_input needs to be defined!"))
 
-    def forward(self):
+    def forward(self) -> Union[torch.Tensor, None]:
         raise (NotImplementedError("forward needs to be defined!"))
+
+    def get_losses(self) -> Union[torch.Tensor, None]:
+        raise (NotImplementedError("get_losses needs to be defined!"))

torch_points3d/models/segmentation/base_model.py

@@ -1,7 +1,10 @@
 from omegaconf import DictConfig
+from typing import Dict, Union
 
+import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch_geometric.data import Data
 
 from torch_points3d.core.instantiator import Instantiator
 from torch_points3d.models.base_model import PointCloudBaseModel
@@ -17,19 +20,25 @@ def __init__(self, instantiator: Instantiator, num_classes: int, backbone: DictC
 
         self.head = nn.Sequential(nn.Linear(self.backbone.output_nc, num_classes))
 
-    def set_input(self, data):
+    def set_input(self, data: Data) -> None:
         self.batch_idx = data.batch.squeeze()
         self.input = data
         if data.y is not None:
             self.labels = data.y
         else:
             self.labels = None
 
-    def forward(self):
+    def forward(self) -> Union[torch.Tensor, None]:
         features = self.backbone(self.input).x
         logits = self.head(features)
         self.output = F.log_softmax(logits, dim=-1)
 
+        return self.get_losses()
+
+    def get_losses(self) -> Union[torch.Tensor, None]:
         # only compute loss if loss is defined and the dset has labels
-        if self.labels is not None and self.criterion is not None:
-            return self.criterion(self.output, self.labels)
+        if self.labels is None or self.criterion is None:
+            return
+
+        self.loss = self.criterion(self.output, self.labels)
+        return self.loss