FP16

Author: harrydobbs

.gitignore

@@ -138,8 +138,10 @@ wandb/
 *.out
 *.sl
 speed-tree-outputs/
+smart_tree/conf/tree-dataset-test.yaml
 smart_tree/conf/tree-split-test.json
 smart_tree/conf/apple-trellis-split-test.json
 smart_tree/conf/apple-trellis.yaml
 smart_tree/conf/apple-trellis-split.json
-FRNN/
+FRNN/
+debug/

smart_tree/conf/tree-dataset.yaml

@@ -7,7 +7,6 @@ training:
   lr_decay: True
   early_stop_epoch: 20
   early_stop: True
-  use_colour: False
 
   dataset:
     _target_: smart_tree.dataset.dataset.TreeDataset
@@ -18,16 +17,14 @@ training:
     block_size: 4
     buffer_size: 0.4
 
-    # augmentation:
-    #     _target_:  smart_tree.dataset.augmentations.Scale
-    #     min_scale: 0.9
-    #     max_scale: 1.1
-      
+    augmentation:
+    # Scale:
+    #    _target_:  smart_tree.dataset.augmentations.Scale
+    #    min_scale: 0.1
+    #    max_scale: 0.1
+    # Dropout:
     #     _target_:  smart_tree.dataset.augmentations.RandomDropout
     #     max_drop_out: 0.1
-      
-    #     _target_:  smart_tree.dataset.augmentations.RandomColourDropout
-    #     max_drop_out: 0.1
 
   data_loader:
     _target_: torch.utils.data.DataLoader
@@ -67,7 +64,6 @@ model_inference:
   buffer_size: 0.4
   num_workers : 8
   batch_size : 4
-  use_colour: False
 
 skeletonizer:
   K: 16
@@ -92,20 +88,19 @@ pipeline:
     # VoxelDownsample:
     #    _target_:  smart_tree.dataset.augmentations.VoxelDownsample
     #    voxel_size : 0.01
-
-    # FixedRotate:
-    #    _target_:  smart_tree.dataset.augmentations.FixedRotate
-    #    xyz: [0, 0, 90]
+    #FixedRotate:
+    #   _target_:  smart_tree.dataset.augmentations.FixedRotate
+    #   xyz: [0, 0, 90]
     # CentreCloud:
-    #    _target_:  smart_tree.dataset.augmentations.CentreCloud
+    #   _target_:  smart_tree.dataset.augmentations.CentreCloud
 
 
   repair_skeletons : True
   smooth_skeletons : True
   prune_skeletons : True
   min_skeleton_radius : 0.001 #0.005
   min_skeleton_length : 0.02
-  view_model_output : True
+  view_model_output : False
   view_skeletons : True
   save_outputs : False
   branch_classes: [0]

smart_tree/data_types/cloud.py

@@ -193,8 +193,8 @@ def medial_pts(self):
     @staticmethod
     def from_numpy(xyz, rgb, vector, class_l):
         return LabelledCloud(
-            torch.from_numpy(xyz).float(),  # float64 -> these data types are stupid...
+            torch.from_numpy(xyz).float(),  # -> these data types are stupid...
             torch.from_numpy(rgb).float(),  # float64
             torch.from_numpy(vector).float(),  # float32
-            torch.from_numpy(class_l).float(),  # int64
+            torch.from_numpy(class_l).int(),  # int64
         )

smart_tree/dataset/dataset.py

@@ -103,7 +103,7 @@ def __getitem__(self, idx):
         if self.blocking:
             loss_mask = cube_filter(feats[:, :3], block_center, self.block_size)
 
-        return feats.float(), coords.int(), loss_mask
+        return feats, coords.int(), loss_mask
 
     def __len__(self):
         return len(self.tree_paths)

smart_tree/model/model.py

@@ -5,7 +5,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from smart_tree.model.model_blocks import SparseFC, UBlock, SubMConvBlock
+from smart_tree.model.model_blocks import MLP, UBlock, SubMConvBlock
 from smart_tree.util.math.maths import torch_normalized
 
 spconv.constants.SPCONV_ALLOW_TF32 = True
@@ -30,9 +30,7 @@ def __init__(
 
         self.branch_classes = torch.tensor(branch_classes, device=device)
 
-        norm_fn = functools.partial(
-            nn.BatchNorm1d, eps=1e-4, momentum=0.1
-        )  # , momentum=0.99)
+        norm_fn = functools.partial(nn.BatchNorm1d, eps=1e-4)  # , momentum=0.99)
         activation_fn = nn.ReLU
 
         self.radius_loss = nn.L1Loss()
@@ -55,24 +53,9 @@ def __init__(
             algo=algo,
         )
 
-        self.radius_head = SparseFC(
-            radius_fc_planes,
-            norm_fn,
-            activation_fn,
-            algo=algo,
-        )
-        self.direction_head = SparseFC(
-            direction_fc_planes,
-            norm_fn,
-            activation_fn,
-            algo=algo,
-        )
-        self.class_head = SparseFC(
-            class_fc_planes,
-            norm_fn,
-            activation_fn,
-            algo=algo,
-        )
+        self.radius_head = MLP(radius_fc_planes, norm_fn, activation_fn)
+        self.direction_head = MLP(direction_fc_planes, norm_fn, activation_fn)
+        self.class_head = MLP(class_fc_planes, norm_fn, activation_fn)
 
         self.apply(self.set_bn_init)
 
@@ -87,12 +70,12 @@ def forward(self, input):
         x = self.input_conv(input)
         unet_out = self.UNet(x)
 
-        radius = self.radius_head(unet_out)
-        direction = self.direction_head(unet_out)
-        class_l = self.class_head(unet_out)
+        radius = self.radius_head(unet_out).features
+        direction = self.direction_head(unet_out).features
+        class_l = self.class_head(unet_out).features
 
         return torch.cat(
-            [radius.features, F.normalize(direction.features), class_l.features],
+            [radius, direction, class_l],
             dim=1,
         )
 
@@ -105,7 +88,7 @@ def compute_loss(self, outputs, targets, mask=None):
             targets = targets[mask]
 
         radius_pred = outputs[:, [0]]
-        direction_pred = outputs[:, 1:4]
+        direction_pred = F.normalize(outputs[:, 1:4])
         class_pred = outputs[:, 4:]
 
         class_target = targets[:, [3]]
@@ -128,26 +111,23 @@ def compute_loss(self, outputs, targets, mask=None):
 
         return losses
 
-    # @force_fp32(apply_to=("outputs", "targets"))
+    @force_fp32(apply_to=("outputs", "targets"))
     def compute_radius_loss(self, outputs, targets):
         return self.radius_loss(outputs, torch.log(targets))
 
-    # @force_fp32(apply_to=("outputs", "targets"))
+    @force_fp32(apply_to=("outputs", "targets"))
     def compute_direction_loss(self, outputs, targets):
         return torch.mean(1 - self.direction_loss(outputs, targets))
 
-    # @force_fp32(apply_to=("outputs", "targets"))
+    @force_fp32(apply_to=("outputs", "targets"))
     def compute_class_loss(self, outputs, targets):
-        return self.focal_loss(outputs, targets.long())
+        return self.dice_loss(outputs, targets.long())
 
     def dice_loss(self, outputs, targets):
         # https://gist.github.com/jeremyjordan/9ea3032a32909f71dd2ab35fe3bacc08
         smooth = 1
         outputs = F.softmax(outputs, dim=1)
-        targets = F.one_hot(targets)
-
-        outputs = outputs.view(-1)
-        targets = targets.view(-1)
+        targets = F.one_hot(targets).reshape(-1, 1)
 
         intersection = (outputs * targets).sum()
 

smart_tree/model/model_blocks.py

@@ -127,19 +127,19 @@ def __init__(
                     output_channels,
                     kernel_size=1,
                     padding=1,
-                    bias=False,
+                    bias=bias,
                     algo=algo,
                 )
             )
 
         self.sequence = spconv.SparseSequential(
             spconv.SubMConv3d(
-                input_channels, output_channels, kernel_size, bias=False, algo=algo
+                input_channels, output_channels, kernel_size, bias=bias, algo=algo
             ),
             norm_fn(output_channels),
             activation_fn(),
             spconv.SubMConv3d(
-                output_channels, output_channels, kernel_size, bias=False, algo=algo
+                output_channels, output_channels, kernel_size, bias=bias, algo=algo
             ),
             norm_fn(output_channels),
         )
@@ -244,43 +244,23 @@ def forward(self, input):
         return output
 
 
-class SparseFC(nn.Module):
+class MLP(nn.Module):
     def __init__(
         self,
         n_planes,
         norm_fn,
         activation_fn=None,
-        kernel_size=1,
-        algo=spconv.ConvAlgo.Native,
-        bias=False,
     ):
         super().__init__()
 
         self.sequence = spconv.SparseSequential()
+
         for i in range(len(n_planes) - 2):
-            self.sequence.add(
-                spconv.SubMConv3d(
-                    n_planes[i],
-                    n_planes[i + 1],
-                    kernel_size=kernel_size,
-                    bias=False,
-                    algo=algo,
-                    padding=0,
-                )
-            )
+            self.sequence.add(nn.Linear(n_planes[i], n_planes[i + 1]))
             self.sequence.add(norm_fn(n_planes[i + 1]))
             self.sequence.add(activation_fn())
 
-        self.sequence.add(
-            spconv.SubMConv3d(
-                n_planes[-2],
-                n_planes[-1],
-                kernel_size=kernel_size,
-                bias=False,
-                algo=algo,
-                padding=0,
-            )
-        )
+        self.sequence.add(nn.Linear(n_planes[-2], n_planes[-1]))
 
     def forward(self, input):
         return self.sequence(input)

smart_tree/model/model_inference.py

@@ -9,6 +9,8 @@
 from omegaconf import DictConfig, OmegaConf
 from py_structs.torch import map_tensors
 from torch.utils.data import DataLoader
+import torch.nn.functional as F
+
 from tqdm import tqdm
 
 import wandb
@@ -159,7 +161,7 @@ def test():
 
         xyz = xyz[:, :3]  # [mask]
         radius = out[:, [0]]  # [mask]
-        direction = out[:, 1:4]  # [mask]
+        direction = F.normalize(out[:, 1:4])  # [mask]
 
         new_xyz = xyz + torch.exp(radius) * direction