search

neuronets / nobrainer

A framework for developing neural network models for 3D image processing.
Other
158 stars 45 forks source link

suggested refactoring to avoid OOM errors #313

Open hvgazula opened 7 months ago

hvgazula commented 7 months ago

https://github.com/neuronets/nobrainer/blob/cb855feaadd4ac354e1e2d1c760a649df3f61ab4/nobrainer/dataset.py#L249-L250

Suggestion:

def scalar_labels(self):
        temp_ds = self.dataset.map(
            lambda _, y: tf.experimental.numpy.isscalar(y),
            deterministic=False,
            num_parallel_calls=AUTOTUNE,
        )
        return tf.math.reduce_all(list(temp_ds.as_numpy_iterator())).numpy()

Notes:

  1. The previous snippet collects all label volumes into a list (this is a memory hog and hence the reason for OOM) and then applies _labels_all_scalar.
  2. Refactored the snippet to map each label volume into a isscalar function which returns a bool flag. Subsequently, the collected bool flags are reduced to one final bool flag.
  3. The other (naive) approach is to run nobrainer.tfrecord._is_int_or_float() on each element of the dataset (in a for loop) and then reduce all the bool flags (similar to step 2).
  4. I am unsure why the GPU utilization is non-zero during this operation.
  5. I still maintain that the repeat should be delayed until after this operation. Otherwise, the entire repeated dataset will be used for this operation and is undesirable.

Caveat:

  1. Used a tf.experimental function which may (or may not) be deprecated in the future. @satra what are your thoughts on using experimental features in the nobrainer API?
hvgazula commented 7 months ago

this suggestion works on a small dataset and doesn't fail. however, it is still time-consuming. The non-zero GPU util may have to do with this (🤷‍♂️ ) as there will be an overhead in going back and forth between the CPU and GPU.