move scheduler/optim to be contained within the model, not in global scope

separate "backbones" from "lightning models"

fix sp3d sparsetensor device

add basic data transforms (need to structure better than original repo)

add training methods

Author: clee-ai

conf/config.yaml

@@ -1,7 +1,5 @@
 defaults: # loads default configs
   - dataset: ???
-  - optimizer: sgd
-  - scheduler: default
   - model: ???
   - training: default
   - trainer: default

conf/model/default.yaml

@@ -1,6 +1,8 @@
+# @package model
+defaults:
+  - /optimizer: sgd
+  - /scheduler: 
 # By default we turn off recursive instantiation, allowing the user to instantiate themselves at the appropriate times.
 _recursive_: false
 
-_target_: torch_points3d.models.base_model.PointCloudBaseModel
-optimizer: ${optimizer}
-scheduler: ${scheduler}
+_target_: torch_points3d.models.base_model.PointCloudBaseModel

conf/model/segmentation/default.yaml

@@ -0,0 +1,8 @@
+# @package model
+defaults:
+  - /model/default
+
+_target_: torch_points3d.models.segmentation.base_model.SegmentationBaseModel
+
+backbone:
+  architecture: unet

conf/model/segmentation/sparseconv3d.yaml

@@ -0,0 +1,9 @@
+# @package model
+defaults:
+  - /model/segmentation/ResUNet32
+
+backbone:
+  down_conv:
+    N: [ 0, 2, 3, 4, 6 ]
+  up_conv:
+    N: [ 1, 1, 1, 1, 1 ]

conf/model/segmentation/sparseconv3d/Res16UNet34.yaml

@@ -0,0 +1,40 @@
+# @package model
+defaults:
+  - /model/segmentation/default
+
+backbone:
+  _target_: torch_points3d.applications.sparseconv3d.SparseConv3d
+  backend: torchsparse
+
+  config:
+    define_constants:
+      in_feat: 32
+      block: ResBlock # Can be any of the blocks in modules/MinkowskiEngine/api_modules.py
+    down_conv:
+      module_name: ResNetDown
+      block: block
+      N: [ 0, 1, 2, 2, 3 ]
+      down_conv_nn:
+        [
+          [ FEAT, in_feat ],
+          [ in_feat, in_feat ],
+          [ in_feat, 2*in_feat ],
+          [ 2*in_feat, 4*in_feat ],
+          [ 4*in_feat, 8*in_feat ],
+        ]
+      kernel_size: 3
+      stride: [ 1, 2, 2, 2, 2 ]
+    up_conv:
+      block: block
+      module_name: ResNetUp
+      N: [ 1, 1, 1, 1, 0 ]
+      up_conv_nn:
+        [
+          [ 8*in_feat, 4*in_feat ],
+          [ 4*in_feat + 4*in_feat, 4*in_feat ],
+          [ 4*in_feat + 2*in_feat, 3*in_feat ],
+          [ 3*in_feat + in_feat, 3*in_feat ],
+          [ 3*in_feat + in_feat, 3*in_feat ],
+        ]
+      kernel_size: 3
+      stride: [ 2, 2, 2, 2, 1 ]

conf/model/segmentation/sparseconv3d/ResUNet32.yaml

@@ -6,7 +6,7 @@ gradient_clip_val: 0.0
 process_position: 0
 num_nodes: 1
 num_processes: 1
-gpus: null
+gpus: 1
 auto_select_gpus: False
 tpu_cores: null
 log_gpu_memory: null

conf/scheduler/default.yaml

@@ -135,7 +135,7 @@ def _set_input(self, data):
         data:
             a dictionary that contains the data itself and its metadata information.
         """
-        self.input = sp3d.nn.SparseTensor(data.x, data.coords, data.batch)
+        self.input = sp3d.nn.SparseTensor(data.x, data.coords, data.batch, self.device)
         if data.pos is not None:
             self.xyz = data.pos
         else:
@@ -163,7 +163,7 @@ def forward(self, data, *args, **kwargs):
         for i in range(len(self.down_modules)):
             data = self.down_modules[i](data)
 
-        out = Batch(x=data.F, batch=data.C[:, 0].long().to(data.F.device))
+        out = Batch(x=data.F, batch=data.C[:, 0].long())
         if not isinstance(self.inner_modules[0], Identity):
             out = self.inner_modules[0](out)
 

conf/trainer/default.yaml

@@ -0,0 +1,147 @@
+import torch
+import re
+import logging
+import torch.nn.functional as F
+from torch_scatter import scatter_mean, scatter_add
+from torch_geometric.nn.pool.consecutive import consecutive_cluster
+from torch_geometric.nn import voxel_grid
+from torch_cluster import grid_cluster
+
+log = logging.getLogger(__name__)
+
+class AddOnes(object):
+    """
+    Add ones tensor to data
+    """
+
+    def __call__(self, data):
+        num_nodes = data.pos.shape[0]
+        data.ones = torch.ones((num_nodes, 1)).float()
+        return data
+
+    def __repr__(self):
+        return "{}()".format(self.__class__.__name__)
+
+
+# Label will be the majority label in each voxel
+_INTEGER_LABEL_KEYS = ["y", "instance_labels"]
+
+def shuffle_data(data):
+    num_points = data.pos.shape[0]
+    shuffle_idx = torch.randperm(num_points)
+    for key in set(data.keys):
+        item = data[key]
+        if torch.is_tensor(item) and num_points == item.shape[0]:
+            data[key] = item[shuffle_idx]
+    return data
+
+
+def group_data(data, cluster=None, unique_pos_indices=None, mode="last", skip_keys=[]):
+    """ Group data based on indices in cluster.
+    The option ``mode`` controls how data gets agregated within each cluster.
+
+    Parameters
+    ----------
+    data : Data
+        [description]
+    cluster : torch.Tensor
+        Tensor of the same size as the number of points in data. Each element is the cluster index of that point.
+    unique_pos_indices : torch.tensor
+        Tensor containing one index per cluster, this index will be used to select features and labels
+    mode : str
+        Option to select how the features and labels for each voxel is computed. Can be ``last`` or ``mean``.
+        ``last`` selects the last point falling in a voxel as the representent, ``mean`` takes the average.
+    skip_keys: list
+        Keys of attributes to skip in the grouping
+    """
+
+    assert mode in ["mean", "last"]
+    if mode == "mean" and cluster is None:
+        raise ValueError("In mean mode the cluster argument needs to be specified")
+    if mode == "last" and unique_pos_indices is None:
+        raise ValueError("In last mode the unique_pos_indices argument needs to be specified")
+
+    num_nodes = data.num_nodes
+    for key, item in data:
+        if bool(re.search("edge", key)):
+            raise ValueError("Edges not supported. Wrong data type.")
+        if key in skip_keys:
+            continue
+
+        if torch.is_tensor(item) and item.size(0) == num_nodes:
+            if mode == "last" or key == "batch": #or key == SaveOriginalPosId.KEY:
+                data[key] = item[unique_pos_indices]
+            elif mode == "mean":
+                is_item_bool = item.dtype == torch.bool
+                if is_item_bool:
+                    item = item.int()
+                if key in _INTEGER_LABEL_KEYS:
+                    item_min = item.min()
+                    item = F.one_hot(item - item_min)
+                    item = scatter_add(item, cluster, dim=0)
+                    data[key] = item.argmax(dim=-1) + item_min
+                else:
+                    data[key] = scatter_mean(item, cluster, dim=0)
+                if is_item_bool:
+                    data[key] = data[key].bool()
+    return data
+
+# todo: replace these with minkowski/torchsparse impl?
+class GridSampling3D:
+    """ Clusters points into voxels with size :attr:`size`.
+    Parameters
+    ----------
+    size: float
+        Size of a voxel (in each dimension).
+    quantize_coords: bool
+        If True, it will convert the points into their associated sparse coordinates within the grid and store
+        the value into a new `coords` attribute
+    mode: string:
+        The mode can be either `last` or `mean`.
+        If mode is `mean`, all the points and their features within a cell will be averaged
+        If mode is `last`, one random points per cell will be selected with its associated features
+    """
+
+    def __init__(self, size, quantize_coords=False, mode="mean", verbose=False):
+        self._grid_size = size
+        self._quantize_coords = quantize_coords
+        self._mode = mode
+        if verbose:
+            log.warning(
+                "If you need to keep track of the position of your points, use SaveOriginalPosId transform before using GridSampling3D"
+            )
+
+            if self._mode == "last":
+                log.warning(
+                    "The tensors within data will be shuffled each time this transform is applied. Be careful that if an attribute doesn't have the size of num_points, it won't be shuffled"
+                )
+
+    def _process(self, data):
+        if self._mode == "last":
+            data = shuffle_data(data)
+
+        coords = torch.round((data.pos) / self._grid_size)
+        if "batch" not in data:
+            cluster = grid_cluster(coords, torch.tensor([1, 1, 1]))
+        else:
+            cluster = voxel_grid(coords, data.batch, 1)
+        cluster, unique_pos_indices = consecutive_cluster(cluster)
+
+        data = group_data(data, cluster, unique_pos_indices, mode=self._mode)
+        if self._quantize_coords:
+            data.coords = coords[unique_pos_indices].int()
+
+        data.grid_size = torch.tensor([self._grid_size])
+        return data
+
+    def __call__(self, data):
+        if isinstance(data, list):
+            data = [self._process(d) for d in data]
+        else:
+            data = self._process(data)
+        return data
+
+    def __repr__(self):
+        return "{}(grid_size={}, quantize_coords={}, mode={})".format(
+            self.__class__.__name__, self._grid_size, self._quantize_coords, self._mode
+        )