Is "batched" PreChippedGeoSampler supported?

[DimitrisMantas]

I'm attempting to use PreChippedGeoSampler with the LandCoverAI dataset to test some custom training code with a dataset I know "works", but I'm running into some weird issues. In particular, whenever I use the sampler with a batch size which is greater than one, I sometimes get weird sample collation errors.

I am using the chipped version of LandCoverAI (i.e., whatever is in the output folder after downloading and extracting it), with the only difference between the official implementation and mine being that I have assigned a CRS to both the images and masks so that I can use them with my own Dataset implementation. I have confirmed that this works by plotting a bunch of random images alongside their corresponding masks.

My DataModule implementation looks like this:

def setup(self, stage: str) -> None:
    …
    if stage in [“fit”]:
        self.train_sampler = PreChippedGeoSampler(self.train_dataset, shuffle=True)
    …

and I am initializing it like so: datamodule = MyDataModule(batch_size=8, ...).

When I try to train a model, the dev run completes successfully, and I can even get ca. 30 steps through the first epoch until I get the following sample collation error:

RuntimeError: Caught RuntimeError in DataLoader worker process 6.
Original Traceback (most recent call last):
  File "C:\Documents\RoofSense\venv\Lib\site-packages\torch\utils\data\_utils\worker.py", line 308, in _worker_loop
    data = fetcher.fetch(index)
           ^^^^^^^^^^^^^^^^^^^^
  File "C:\Documents\RoofSense\venv\Lib\site-packages\torch\utils\data\_utils\fetch.py", line 54, in fetch
    return self.collate_fn(data)
           ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Documents\RoofSense\venv\Lib\site-packages\torchgeo\datasets\utils.py", line 560, in stack_samples
    collated[key] = torch.stack(value)
                    ^^^^^^^^^^^^^^^^^^
RuntimeError: stack expects each tensor to be equal size, but got [3, 1024, 1024] at entry 0 and [3, 512, 512] at entry 1

I checked the PreChippedGeoSampler tests and wrote the following to investigate the issue:

datamodule.setup("fit")

and

from torchgeo.datasets import stack_samples
from torch.utils.data import DataLoader


dataloader = DataLoader(datamodule.train_dataset,batch_size=8,sampler=datamodule.train_sampler,collate_fn=stack_samples).
for _ in dataloader:
    continue

This throws the same error as above, but if I rerun it, it appears to work fine. Additionally, I have noticed that the error returns if I set up the data module again.

I'm not 100% certain this is a bug on TorchGeo's end, because pretty much all my code is custom, but I'd be grateful if somebody could replicate it.

[DimitrisMantas]

By the way, here's a suggestion that just came to mind: since pre-chipped datasets come in this "standardized" form that we can exploit to sample them simply using a file index, why are we relying on a spatial index at all?

I get that they are GeoDatasets and come with a CRS, and so the logical thing is put them into a spatial index, but this is kind of irrelevant in the case of chip sampling IMO, and reading 10k+ chips into memory to construct this index takes time, and so does the actual indexing. This can significantly increase the training setup and epoch durations, potentially resulting in reduced batch sizes, etc.

I think a "hybrid" dataset that behaves like a Geo- or NonGeo-Dataset on demand (e.g., via a use_index parameter in its initialized and type-dispatched getters) would be a nice addition.

Relevant discussions/issues: #479, #1353

[adamjstewart]

Is "batched" PreChippedGeoSampler supported?

Yes, iff all chips are the same size in the same CRS. If not, you'll need to random/center crop or resize all images to be the same size in order to collate them.

why are we relying on a spatial index at all?

You don't have to. That's why there is a NonGeoDataset version of LandCover.ai. The benefit of a GeoDataset is solely that you can combine it with any other GeoDataset. So for example, you could create a dataset that includes LandCover.ai imagery and masks and combine it with Sentinel-2 imagery or a DEM to see how that impacts performance. I don't personally see any other advantage of using GeoDataset instead of NonGeoDataset. If you're only interested in benchmarking, NonGeoDataset is better.

A hybrid approach that combined Geo and NonGeo would be interesting, although confusing (you have to explain that sometimes this dataset requires a GeoSampler and sometimes it doesn't) and a lot of code in a single class. But it would negate the need for duplicated classes for datasets like LandCover.ai which would be nice. I haven't thought too much in this direction yet.

#409 is also a relevant discussion, and would negate most of the benefits of #1353.

[DimitrisMantas]

Reading your reply, I just remembered that the LandCoverAI imagery is of varying resolution, and I haven’t accounted for that. I’ll check it out again later and mark your input as the answer assuming everything works fine (I guess it should – that was a pretty big oversight on my part).

I can’t really comment on the reprojection stuff in #409, because I’m not familiar with warped files and VRTs that TorchGeo apparently uses.

My use case is that I have a dataset that I have pre-chipped into patches for easier annotation, but I’m still experimenting with the patch size for training. In my mind, patch sizes less than the chip size should be handled by BatchGeoSampler, with PreChippedGeoSampler (or an equivalent) being responsible for every other case. However, since all these chips and their corresponding masks are georeferenced, my dataset implementation extends RasterDataset, and therefore relies on a spatial index.

I’m currently looking to implement this uses_index parameter so that I can instantiate the appropriate “version” (i.e., geo- or "non-geo") of my dataset in the corresponding data module and get rid of PreChippedGeoSampler all together.

[adamjstewart]

I’m not familiar with warped files and VRTs that TorchGeo apparently uses.

Our paper does a decent job of explaining the rationale behind it. Our blog is a shorter, easier to read summary of this.

TL;DR: if you want to combine various data sources, all data needs to be reprojected to a common coordinate system in order to reason about overlap or sample patches. The user can either do this themselves as a pre-processing step, or TorchGeo can do it automatically.

P.S. I'm not trying to suggest that your workflow is incorrect or non-standard. For now, cropping or resizing imagery should work well for experimentation, but you may want to consider a NonGeoDataset for benchmarking and reproducibility. TorchGeo supports both Geo and NonGeo, and I would like to make both easier to use and more flexible in the future, but the technical difficulties of this give me pause. We're currently working on redesigning much of the internals to handle time-series support which is proving to be difficult enough.

[DimitrisMantas]

Of course, I see the benefits of NonGeoDataset. I guess this is why I’m trying to avoid relying on GeoDataset functionality, at least where I can. I may rewrite my dataset implementation to be completely non-geo in the future if I end up not relying on spatial indexing for training or inference.

Really appreciate your work and looking forward to new updates!