numpy.core._exceptions._ArrayMemoryError with SemanticSegmentationSlidingWindowGeoDataset

[MathiasBaumgartinger]

Hello, I am using a configuration which closely resembles the one provided at: https://docs.rastervision.io/en/0.20/usage/tutorials/pred_and_eval_ss.html. My configuration has worked on smaller test scenarios. When trying to predict a fairly large ortho image (i.e 90 GB) I am receiving an numpy.core._exceptions._ArrayMemoryError. Any idea what might cause this?

The exact error output is: numpy.core._exceptions._ArrayMemoryError: Unable to allocate 412. GiB for an array with shape (3, 271959, 271076) and data type float16. Maybe I am getting something wrong here but isn't the SemanticSegmentationSlidingWindowGeoDataset class's purpose to split a large dataset into smaller subdatasets using rasterio windows?

Thank you :)

[AdeelH]

Please share the full stack trace.

[MathiasBaumgartinger]

Sure:

Traceback (most recent call last):
  File "/home/user/rastervision/code/rv_predict.py", line 61, in <module>
    pred_labels = SemanticSegmentationLabels.from_predictions(
  File "/home/user/anaconda3/envs/rastervision/lib/python3.9/site-packages/rastervision/core/data/label/semantic_segmentation_labels.py", line 219, in from_predictions
    labels = cls.make_empty(extent, num_classes, smooth=smooth)
  File "/home/user/anaconda3/envs/rastervision/lib/python3.9/site-packages/rastervision/core/data/label/semantic_segmentation_labels.py", line 182, in make_empty
    return SemanticSegmentationSmoothLabels.make_empty(
  File "/home/user/anaconda3/envs/rastervision/lib/python3.9/site-packages/rastervision/core/data/label/semantic_segmentation_labels.py", line 519, in make_empty
    return cls(extent=extent, num_classes=num_classes)
  File "/home/user/anaconda3/envs/rastervision/lib/python3.9/site-packages/rastervision/core/data/label/semantic_segmentation_labels.py", line 443, in __init__
    self.pixel_scores = np.zeros(
numpy.core._exceptions._ArrayMemoryError: Unable to allocate 412. GiB for an array with shape (3, 271959, 271076) and data type float16

[AdeelH]

WRT to that notebook, this is the part where you're running into an error. It is not related to SemanticSegmentationSlidingWindowGeoDataset.

pred_labels = SemanticSegmentationLabels.from_predictions(
    ds.windows,
    predictions,
    smooth=True,
    extent=ds.scene.extent,
    num_classes=len(class_config))

The above code tries to load ALL the predicted chips into memory. This is obviously not feasible when your image is too large to fit into memory.

You can try one of the following:

• Split your large TIFF into smaller ones and do prediction on them one at a time.
• Write custom code for writing the predictions to a GeoTIFF that doesn't use as much memory. predictions in the above snippet is a generator that yields chips (the window for each chip is in the corresponding item in ds.windows); you could potentially loop over it and write out chips immediately without storing them. Take a look at the source for SemanticSegmentationLabelStore.write_smooth_raster_output().