Repository: microsoft/torchgeo

Issue:
Samplers don't work in parallel
When I try to use RandomGeoSampler with `num_workers` >= 1, I get the following error:
```
Traceback (most recent call last):
  File "/anaconda/envs/azureml_py38/lib/python3.9/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/anaconda/envs/azureml_py38/lib/python3.9/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/anaconda/envs/azureml_py38/lib/python3.9/site-packages/torch/utils/data/_utils/worker.py", line 260, in _worker_loop
    r = index_queue.get(timeout=MP_STATUS_CHECK_INTERVAL)
  File "/anaconda/envs/azureml_py38/lib/python3.9/multiprocessing/queues.py", line 122, in get
    return _ForkingPickler.loads(res)
TypeError: __new__() missing 5 required positional arguments: 'maxx', 'miny', 'maxy', 'mint', and 'maxt'
```


Discussion:
May be related to https://github.com/pytorch/pytorch/issues/63483

Resolving patch:
diff --git a/torchgeo/datasets/utils.py b/torchgeo/datasets/utils.py
--- a/torchgeo/datasets/utils.py
+++ b/torchgeo/datasets/utils.py
@@ -214,6 +214,14 @@ def __init__(
         self.mint = mint
         self.maxt = maxt
 
+    def __getnewargs__(self) -> Tuple[float, float, float, float, float, float]:
+        """Values passed to the ``__new__()`` method upon unpickling.
+
+        Returns:
+            tuple of bounds
+        """
+        return self.minx, self.maxx, self.miny, self.maxy, self.mint, self.maxt
+
     def __repr__(self) -> str:
         """Return the formal string representation of the object.
 


Test patch:
diff --git a/tests/datasets/test_geo.py b/tests/datasets/test_geo.py
--- a/tests/datasets/test_geo.py
+++ b/tests/datasets/test_geo.py
@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Any, Dict
+from typing import Dict
 
 import pytest
 from rasterio.crs import CRS
@@ -27,12 +27,12 @@ def __init__(
         self.crs = crs
         self.res = res
 
-    def __getitem__(self, query: BoundingBox) -> Dict[str, Any]:
+    def __getitem__(self, query: BoundingBox) -> Dict[str, BoundingBox]:
         return {"index": query}
 
 
 class CustomVisionDataset(VisionDataset):
-    def __getitem__(self, index: int) -> Dict[str, Any]:
+    def __getitem__(self, index: int) -> Dict[str, int]:
         return {"index": index}
 
     def __len__(self) -> int:
diff --git a/tests/datasets/test_utils.py b/tests/datasets/test_utils.py
--- a/tests/datasets/test_utils.py
+++ b/tests/datasets/test_utils.py
@@ -1,6 +1,7 @@
 import builtins
 import glob
 import os
+import pickle
 import shutil
 import sys
 from pathlib import Path
@@ -177,6 +178,12 @@ def test_intersects(
         bbox2 = BoundingBox(*test_input)
         assert bbox1.intersects(bbox2) == bbox2.intersects(bbox1) == expected
 
+    def test_picklable(self) -> None:
+        bbox = BoundingBox(0, 1, 2, 3, 4, 5)
+        x = pickle.dumps(bbox)
+        y = pickle.loads(x)
+        assert bbox == y
+
     def test_invalid_x(self) -> None:
         with pytest.raises(
             ValueError, match="Bounding box is invalid: 'minx=1' > 'maxx=0'"
diff --git a/tests/samplers/test_batch.py b/tests/samplers/test_batch.py
--- a/tests/samplers/test_batch.py
+++ b/tests/samplers/test_batch.py
@@ -1,11 +1,13 @@
 import math
-from typing import Iterator, List
+from typing import Dict, Iterator, List
 
 import pytest
 from _pytest.fixtures import SubRequest
+from rasterio.crs import CRS
 from rtree.index import Index, Property
+from torch.utils.data import DataLoader
 
-from torchgeo.datasets import BoundingBox
+from torchgeo.datasets import BoundingBox, GeoDataset
 from torchgeo.samplers import BatchGeoSampler, RandomBatchGeoSampler
 
 
@@ -21,6 +23,22 @@ def __len__(self) -> int:
         return 2
 
 
+class CustomGeoDataset(GeoDataset):
+    def __init__(
+        self,
+        bounds: BoundingBox = BoundingBox(0, 1, 2, 3, 4, 5),
+        crs: CRS = CRS.from_epsg(3005),
+        res: float = 1,
+    ) -> None:
+        super().__init__()
+        self.index.insert(0, bounds)
+        self.crs = crs
+        self.res = res
+
+    def __getitem__(self, query: BoundingBox) -> Dict[str, BoundingBox]:
+        return {"index": query}
+
+
 class TestBatchGeoSampler:
     @pytest.fixture(scope="function")
     def sampler(self) -> CustomBatchGeoSampler:
@@ -33,6 +51,15 @@ def test_iter(self, sampler: CustomBatchGeoSampler) -> None:
     def test_len(self, sampler: CustomBatchGeoSampler) -> None:
         assert len(sampler) == 2
 
+    @pytest.mark.parametrize("num_workers", [0, 1, 2])
+    def test_dataloader(self, sampler: CustomBatchGeoSampler, num_workers: int) -> None:
+        ds = CustomGeoDataset()
+        dl = DataLoader(
+            ds, batch_sampler=sampler, num_workers=num_workers  # type: ignore[arg-type]
+        )
+        for _ in dl:
+            continue
+
     def test_abstract(self) -> None:
         with pytest.raises(TypeError, match="Can't instantiate abstract class"):
             BatchGeoSampler(None)  # type: ignore[abstract]
@@ -62,3 +89,12 @@ def test_iter(self, sampler: RandomBatchGeoSampler) -> None:
 
     def test_len(self, sampler: RandomBatchGeoSampler) -> None:
         assert len(sampler) == sampler.length // sampler.batch_size
+
+    @pytest.mark.parametrize("num_workers", [0, 1, 2])
+    def test_dataloader(self, sampler: RandomBatchGeoSampler, num_workers: int) -> None:
+        ds = CustomGeoDataset()
+        dl = DataLoader(
+            ds, batch_sampler=sampler, num_workers=num_workers  # type: ignore[arg-type]
+        )
+        for _ in dl:
+            continue
diff --git a/tests/samplers/test_single.py b/tests/samplers/test_single.py
--- a/tests/samplers/test_single.py
+++ b/tests/samplers/test_single.py
@@ -1,11 +1,13 @@
 import math
-from typing import Iterator
+from typing import Dict, Iterator
 
 import pytest
 from _pytest.fixtures import SubRequest
+from rasterio.crs import CRS
 from rtree.index import Index, Property
+from torch.utils.data import DataLoader
 
-from torchgeo.datasets import BoundingBox
+from torchgeo.datasets import BoundingBox, GeoDataset
 from torchgeo.samplers import GeoSampler, GridGeoSampler, RandomGeoSampler
 
 
@@ -21,6 +23,22 @@ def __len__(self) -> int:
         return 2
 
 
+class CustomGeoDataset(GeoDataset):
+    def __init__(
+        self,
+        bounds: BoundingBox = BoundingBox(0, 1, 2, 3, 4, 5),
+        crs: CRS = CRS.from_epsg(3005),
+        res: float = 1,
+    ) -> None:
+        super().__init__()
+        self.index.insert(0, bounds)
+        self.crs = crs
+        self.res = res
+
+    def __getitem__(self, query: BoundingBox) -> Dict[str, BoundingBox]:
+        return {"index": query}
+
+
 class TestGeoSampler:
     @pytest.fixture(scope="function")
     def sampler(self) -> CustomGeoSampler:
@@ -32,6 +50,15 @@ def test_iter(self, sampler: CustomGeoSampler) -> None:
     def test_len(self, sampler: CustomGeoSampler) -> None:
         assert len(sampler) == 2
 
+    @pytest.mark.parametrize("num_workers", [0, 1, 2])
+    def test_dataloader(self, sampler: CustomGeoSampler, num_workers: int) -> None:
+        ds = CustomGeoDataset()
+        dl = DataLoader(
+            ds, sampler=sampler, num_workers=num_workers  # type: ignore[arg-type]
+        )
+        for _ in dl:
+            continue
+
     def test_abstract(self) -> None:
         with pytest.raises(TypeError, match="Can't instantiate abstract class"):
             GeoSampler(None)  # type: ignore[abstract]
@@ -61,6 +88,15 @@ def test_iter(self, sampler: RandomGeoSampler) -> None:
     def test_len(self, sampler: RandomGeoSampler) -> None:
         assert len(sampler) == sampler.length
 
+    @pytest.mark.parametrize("num_workers", [0, 1, 2])
+    def test_dataloader(self, sampler: RandomGeoSampler, num_workers: int) -> None:
+        ds = CustomGeoDataset()
+        dl = DataLoader(
+            ds, sampler=sampler, num_workers=num_workers  # type: ignore[arg-type]
+        )
+        for _ in dl:
+            continue
+
 
 class TestGridGeoSampler:
     @pytest.fixture(
@@ -92,3 +128,12 @@ def test_iter(self, sampler: GridGeoSampler) -> None:
             assert math.isclose(
                 query.maxt - query.mint, sampler.roi.maxt - sampler.roi.mint
             )
+
+    @pytest.mark.parametrize("num_workers", [0, 1, 2])
+    def test_dataloader(self, sampler: GridGeoSampler, num_workers: int) -> None:
+        ds = CustomGeoDataset()
+        dl = DataLoader(
+            ds, sampler=sampler, num_workers=num_workers  # type: ignore[arg-type]
+        )
+        for _ in dl:
+            continue