Clean up code for project submission.

Author: zghera

configs/s3dis/__init__.py

@@ -18,7 +18,7 @@
 configs.train = Config()
 configs.train.restart_training = False
 configs.train.num_epochs = 50
-configs.train.batch_size = 32 # Original: 32
+configs.train.batch_size = 32
 configs.train.loss_fn = Config(tf.keras.losses.CategoricalCrossentropy)
 configs.train.loss_fn.axis = 1
 configs.train.loss_fn.reduction = tf.keras.losses.Reduction.SUM

configs/s3dis/pvcnn/__init__.py

@@ -1,7 +1,6 @@
 """Experiment configurations for S3DIS PVCNN model."""
-# TODO: If TF's CosineDecay much different from Pytorch's equivalent?
-import tensorflow as tf
-from tensorflow.keras.optimizers.schedules import CosineDecayRestarts
+# TODO: Is TF's CosineDecay much different from Pytorch's equivalent?
+from tensorflow.keras.optimizers.schedules import CosineDecay
 from tensorflow.keras.regularizers import L2
 
 from modeling import PVCNN, AttentionSchedule
@@ -23,11 +22,7 @@
 configs.dataset.num_points = 4096
 
 # train: scheduler
-# configs.train.optimizer.learning_rate = 1e-3
 configs.train.optimizer.learning_rate = Config(AttentionSchedule)
-configs.train.optimizer.learning_rate.d_model = 100 # Orig: 128
-configs.train.optimizer.learning_rate.warmup_steps = 1000 # Orig: 4000
+configs.train.optimizer.learning_rate.d_model = 100
+configs.train.optimizer.learning_rate.warmup_steps = 1000
 configs.train.optimizer.learning_rate.eps = 1e-8
-# configs.train.optimizer.learning_rate = Config(CosineDecayRestarts)
-# configs.train.optimizer.learning_rate.initial_learning_rate = 1e-3
-# configs.train.optimizer.learning_rate.first_decay_steps = 3200 # Empirical

dataloaders/s3dis.py

@@ -55,7 +55,7 @@ def create_s3dis_dataset(
     lambda label, data_num_points, data,
            use_normalized_coords=use_normalized_coords,
            desired_num_points=num_points,
-           num_classes=num_classes: _random_sample_data(
+           num_classes=num_classes: _pre_process(
       data,
       label,
       data_num_points,
@@ -128,7 +128,7 @@ def _get_file_info(
 
 
 @tf.function
-def _random_sample_data(
+def _pre_process(
   data: tf.Tensor,
   label: tf.Tensor,
   data_num_points: tf.Tensor,
@@ -170,15 +170,13 @@ def sample_without_replacement():
   label = tf.gather(label, indices=indices)
 
   # Filter out any points with class labels > num_classes or < 0
+  # fmt: off
   indices = tf.squeeze(
-    tf.where(
-      tf.math.logical_and(
+    tf.where(tf.math.logical_and(
         tf.less_equal(label, num_classes - 1),
         tf.greater_equal(label, 0),
-      )
-    ),
-    1,
-  )
+    )), 1)
+  # fmt: on
   data = tf.gather(data, indices=indices)
   label = tf.cast(tf.gather(label, indices=indices), dtype=tf.int32)
 

metrics/s3dis.py

@@ -1,8 +1,9 @@
-"""Evaluation metrics for PVCNN S3DIS dataset."""
+"""Custom evaluation metrics for PVCNN S3DIS dataset."""
 import tensorflow as tf
 import numpy as np
 from keras import backend
 
+
 class IouAccuracy(tf.keras.metrics.Metric):
   """Mean IoU accuracy metric."""
 
@@ -58,5 +59,4 @@ def result(self) -> None:
       self._total_correct,
       (self._total_seen + self._total_positive - self._total_correct),
     )
-    # tf.print("IOU accuracy for each class:", iou, summarize=-1)
     return tf.reduce_sum(iou) / self._num_classes

modeling/layers/__init__.py

@@ -1,7 +1,7 @@
 """PVCNN Custom Layer Implementations."""
 from modeling.layers.mlp import ConvBn, DenseBn
 from modeling.layers.pvconv import PVConv
-from modeling.layers.sub_models import (
+from modeling.layers.components import (
   PointFeaturesBranch,
   CloudFeaturesBranch,
   ClassificationHead,

modeling/layers/sub_models.py renamed to modeling/layers/components.py

@@ -147,21 +147,13 @@ def build(self, input_shape) -> None:
     self._num_points = int(input_shape[-1])
     super().build(input_shape)
 
-  # TODO: Delete debugging prints later
   def call(self, inputs, training: bool) -> tf.Tensor:
-    # print("\nCloudFeaturesBranch inputs =", inputs)
     # Get maximum channel value for each channel over all of the points
     x = tf.math.reduce_max(inputs, axis=-1)
-    # print("\nCloudFeaturesBranch reduced inputs shape =", x.shape)
     for layer in self._layers:
       x = layer(x, training=training)
-      # tf.print("CloudFeaturesBranch layer x out nans =", tf.size(tf.where(tf.math.is_nan(x))))
-      # print("\nCloudFeaturesBranch intermed layer shape =", x.shape)
     x = replace_nans_with_norm(x)
     # Duplicate output tensor for N size num_points dimension
-    # print("\nNon-repeated out tensor = ", x)
-    # print(\nNon-repeated out tensor nan idxs = ", tf.where(tf.math.is_nan(x)))
-    # return tf.stack([x] * 4096, axis=-1)
     return tf.repeat(tf.expand_dims(x, axis=-1), self._num_points, axis=-1)
 
 
@@ -192,6 +184,5 @@ def call(self, inputs, training: bool) -> tf.Tensor:
     x = inputs
     for layer in self._layers:
       x = layer(x, training=training)
-      # tf.print("ClassificationHead layer x out nans =", tf.size(tf.where(tf.math.is_nan(x))))
     x = replace_nans_with_norm(x)
     return self._softmax(x)

modeling/layers/nan_replace.py

@@ -1,6 +1,14 @@
+"""A temporary module to fix the NaN issues during training."""
 import tensorflow as tf
 
+
 def replace_nans_with_norm(has_nans: tf.Tensor) -> tf.Tensor:
-  repl_nans_with_zeros = tf.where(tf.math.is_nan(has_nans), tf.zeros_like(has_nans), has_nans)
+  """Replace NaNs in the input tensor with the norm of the tensor once NaNs
+  have been replaced with zeros."""
+  repl_nans_with_zeros = tf.where(
+    tf.math.is_nan(has_nans), tf.zeros_like(has_nans), has_nans
+  )
   norm = tf.norm(repl_nans_with_zeros)
-  return tf.where(tf.math.is_nan(has_nans), tf.fill(has_nans.shape, norm), has_nans)
+  return tf.where(
+    tf.math.is_nan(has_nans), tf.fill(has_nans.shape, norm), has_nans
+  )

modeling/layers/pvconv.py

@@ -11,7 +11,7 @@
 
 
 class PVConv(tf.keras.layers.Layer):
-  """The infamous PVConv block."""
+  """The PVConv block."""
 
   def __init__(
     self,
@@ -56,71 +56,44 @@ def build(self, input_shape) -> None:
         tf.keras.layers.BatchNormalization(axis=1, epsilon=1e-4)
       )
       self._relu_layers.append(tf.keras.layers.LeakyReLU(alpha=0.1))
-    # Add extra batch norm layer to add before 1st conv3d layer
-    # self._bn_layers.append(
-    #   tf.keras.layers.BatchNormalization(axis=1, epsilon=1e-4)
-    # )
 
     self._point_features = ConvBn(
       out_channels=self._out_channels,
       kernel_regularizer=self._kernel_regularizer,
     )
-
     self._squeeze = tf.keras.layers.Reshape((self._out_channels, -1))
-
     super().build(input_shape)
 
-  # TODO: Delete debugging prints later
   def call(self, inputs, training: bool) -> Tuple[tf.Tensor, tf.Tensor]:
     # IC = input channels | OC = output channels (self._out_channels)
     # features = [B, IC, N]  |  coords = [B, 3, N]
     features, coords = inputs
-    # tf.print("\nfeatures nans =", tf.size(tf.where(tf.math.is_nan(features))))
-    # tf.print("coords nans =", tf.size(tf.where(tf.math.is_nan(coords))))
-    # features = tf.debugging.assert_all_finite(features, "features is nan")
-    # coords = tf.debugging.assert_all_finite(coords, "coords is nan")
 
     voxel_features, voxel_coords = self._voxelization((features, coords))
-    # tf.print("voxel features shape=", voxel_features.shape)
     # |--> voxel_features = [B, IC, R, R, R]  |  voxel_coords = [B, 3, N]
-    # tf.print("voxelization features nans =", tf.size(tf.where(tf.math.is_nan(voxel_features))))
     voxel_features = replace_nans_with_norm(voxel_features)
-    # voxel_features = self._voxelization_nan_filter(voxel_features)
-    # tf.print("voxelization features CLIPPED nans =", tf.size(tf.where(tf.math.is_nan(voxel_features))))
-    # tf.print("voxelization coords nans =", tf.size(tf.where(tf.math.is_nan(voxel_coords))))
-    # voxel_features = tf.debugging.assert_all_finite(voxel_features, "voxelization feat is nan")
-    # voxel_coords = tf.debugging.assert_all_finite(voxel_coords, "voxelization coords is nan")
 
-    # voxel_features = self._bn_layers[-1](voxel_features, training=training)
-    for conv, bn, relu in zip(self._conv_layers, self._bn_layers, self._relu_layers):
+    for conv, bn, relu in zip(
+      self._conv_layers, self._bn_layers, self._relu_layers
+    ):
       voxel_features = conv(voxel_features)
-      voxel_features = replace_nans_with_norm(voxel_features)
-      # tf.print("conv3d out features nans =", tf.size(tf.where(tf.math.is_nan(voxel_features))))
-      # voxel_features = tf.debugging.assert_all_finite(voxel_features, "conv3d out feat is nan")
       voxel_features = bn(voxel_features, training=training)
-      # voxel_features = replace_nans_with_norm(voxel_features)
       voxel_features = relu(voxel_features)
-      # voxel_features = replace_nans_with_norm(voxel_features)
-    # voxel_features = replace_nans_with_norm(voxel_features)
+    voxel_features = replace_nans_with_norm(voxel_features)
 
     # |--> voxel_features = [B, OC, R, R, R]
     voxel_features = self._squeeze(voxel_features)
     # |--> voxel_features = [B, OC, R**3]
     voxel_features, _, _ = trilinear_devoxelize(
       voxel_features, voxel_coords, self._resolution, training
     )
-    # tf.print("devox out features nans =", tf.size(tf.where(tf.math.is_nan(voxel_features))))
-    # voxel_features = tf.debugging.assert_all_finite(voxel_features, "devox out feat is nan")
     # |--> voxel_features = [B, OC, N]
     voxel_features = replace_nans_with_norm(voxel_features)
 
     point_features = self._point_features(features, training=training)
-    # tf.print("point feautes nans =", tf.size(tf.where(tf.math.is_nan(point_features))))
     # |--> point_features = [B, OC, N]
     point_features = replace_nans_with_norm(point_features)
 
     fused_features = voxel_features + point_features
-    # tf.print("fused out features nans =", tf.size(tf.where(tf.math.is_nan(fused_features))))
-    # fused_features = tf.debugging.assert_all_finite(fused_features, "fused feat is nan")
     # |--> fused_features = [B, OC, N]
     return fused_features, coords

modeling/layers/voxelization.py

@@ -8,9 +8,7 @@
 class Voxelization(tf.keras.layers.Layer):
   """Voxelization layer."""
 
-  def __init__(
-    self, resolution: int, normalize: bool, eps: float, **kwargs
-  ):
+  def __init__(self, resolution: int, normalize: bool, eps: float, **kwargs):
     super().__init__(**kwargs)
     self._resolution = resolution
     self._normalize = normalize

modeling/s3dis/lr_schedule.py

@@ -1,10 +1,14 @@
-"""TODO if we end up using"""
+"""Define custom learning rate schedulers."""
+
 import tensorflow as tf
 
+
 class AttentionSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
-  """TODO if we end up using
-  https://www.tensorflow.org/text/tutorials/transformer#optimizer
+  """Learning rate scheduler based off of the scheduler used in 
+     Transformers (Vaswani, Ashish, et al.). Adapted from 
+     https://www.tensorflow.org/text/tutorials/transformer#optimizer.
   """
+
   def __init__(self, d_model=128, warmup_steps=4000, eps=1e-8):
     super().__init__()
     self.d_model = d_model
@@ -15,22 +19,5 @@ def __init__(self, d_model=128, warmup_steps=4000, eps=1e-8):
   def __call__(self, step):
     arg1 = tf.math.rsqrt(step)
     arg2 = step * (self.warmup_steps ** -1.5)
-    # tf.print("----------------")
-    # tf.print("step = ", step)
-    # tf.print("arg1 = ", arg1)
-    # tf.print("arg2 = ", arg2)
-
-    lr = tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2) + self.eps
-    # tf.print("lr = ", lr)
-    # tf.print("----------------")
-    return lr
-
-
-# -------------------------------------------------------------------
-# import matplotlib.pyplot as plt
 
-# temp_learning_rate_schedule = AttentionSchedule()
-# plt.plot(temp_learning_rate_schedule(tf.range(80000, dtype=tf.float32)))
-# plt.ylabel("Learning Rate")
-# plt.xlabel("Train Step")
-# plt.savefig("custom-schedule.png")
+    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2) + self.eps