suggested refactoring to avoid OOM errors

[hvgazula]

nobrainer/nobrainer/dataset.py

Lines 249 to 250 in cb855fe

	def scalar_labels(self):
	return _labels_all_scalar([y for _, y in self.dataset.as_numpy_iterator()])

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]

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.