add criterion

Author: clee-ai

conf/model/segmentation/default.yaml

@@ -6,6 +6,8 @@ model:
   _recursive_: false
   _target_: torch_points3d.models.segmentation.base_model.SegmentationBaseModel
   num_classes: ${dataset.cfg.num_classes}
+  criterion:
+    _target_: torch.nn.NLLLoss
 
   backbone:
     input_nc: ${dataset.cfg.feature_dimension}

torch_points3d/models/segmentation/base_model.py

@@ -8,11 +8,12 @@
 
 
 class SegmentationBaseModel(PointCloudBaseModel):
-    def __init__(self, instantiator: Instantiator, num_classes: int, backbone: DictConfig):
+    def __init__(self, instantiator: Instantiator, num_classes: int, backbone: DictConfig, criterion: DictConfig):
         super().__init__(instantiator)
 
         print(backbone)
         self.backbone = self.instantiator.backbone(backbone)
+        self.criterion = self.instantiator.instantiate(criterion)
 
         self.head = nn.Sequential(nn.Linear(self.backbone.output_nc, num_classes))