Preparing a custom geo dataset for training

[lcoandrade]

Hi there!

I'm working on a Kaggle notebook with a custom dataset based on NAIP images and Chesapeake labels.

I want to perform a segmentation task.

As I'm using an intersection dataset (images from NAIP and labels from Chesapeake):

dataset = naip_images & chesapeake_labels # this means I'm creating an IntersectionDataset

I've created a custom geo datamodule subclassing GeoDatamodule to address the error:

TypeError: IntersectionDataset.__init__() got an unexpected keyword argument 'split'

This is my CustomGeoDatamodule:

class CustomGeoDataModule(GeoDataModule):
    def setup(self, stage: str) -> None:
        """Set up datasets.

        Args:
            stage: Either 'fit', 'validate', 'test', or 'predict'.
        """
        self.dataset = self.dataset_class(**self.kwargs)
        
        generator = torch.Generator().manual_seed(0)
        (
            self.train_dataset,
            self.val_dataset,
            self.test_dataset,
        ) = random_bbox_assignment(dataset, [0.6, 0.2, 0.2], generator)
        
        if stage in ["fit"]:
            self.train_batch_sampler = RandomBatchGeoSampler(
                self.train_dataset, self.patch_size, self.batch_size, self.length
            )
        if stage in ["validate"]:
            self.val_sampler = GridGeoSampler(
                self.val_dataset, self.patch_size, self.patch_size
            )
        if stage in ["test"]:
            self.test_sampler = GridGeoSampler(
                self.test_dataset, self.patch_size, self.patch_size
            )

And this is my datamodule:

datamodule = CustomGeoDataModule(
    dataset_class = type(dataset), # GeoDataModule kwargs
    batch_size = 10, # GeoDataModule kwargs
    patch_size = IMG_SIZE, # GeoDataModule kwargs
    length = SAMPLE_SIZE, # GeoDataModule kwargs
    num_workers = WORKERS, # GeoDataModule kwargs
    dataset1 = naip_images, # IntersectionDataset kwargs
    dataset2 = chesapeake_labels, # IntersectionDataset kwargs
    collate_fn = stack_samples, # IntersectionDataset kwargs
)

But when I try to train with this:

CHECKPOINT_DIR = '../input/savedckpt/'
checkpoint_file = os.path.join(CHECKPOINT_DIR, 'torchgeo_trained.ckpt') 


if os.path.isfile(checkpoint_file):
    print('Resuming training from previous checkpoint...')
    trainer.fit(
        model=task, 
        datamodule=datamodule,
        ckpt_path=checkpoint_file
    )
else:
    print('Starting training from scratch...')
    trainer.fit(
        model=task,
        datamodule = datamodule,
    )

I get this error:

torchgeo.datamodules.utils.MisconfigurationException: CustomGeoDataModule.setup does not define a 'val_dataset'

What am I missing here?

Thanks in advance!

[adamjstewart]

You also need to override __init__ to define the dataset class. However, in this particular case, we already have a builtin datamodule for the task you are trying to solve: https://torchgeo.readthedocs.io/en/stable/api/datamodules.html#naip. You can either use this directly or use it as a reference when writing yours.

[lcoandrade]

This is my dataset:

dataset = naip_images & chesapeake_labels # this means I'm creating an IntersectionDataset

And here I pass the dataset class to my GeoDatamodule:

datamodule = CustomGeoDataModule(
    dataset_class = type(dataset), # GeoDataModule kwargs
    batch_size = 10, # GeoDataModule kwargs
    patch_size = IMG_SIZE, # GeoDataModule kwargs
    length = SAMPLE_SIZE, # GeoDataModule kwargs
    num_workers = WORKERS, # GeoDataModule kwargs
    dataset1 = naip_images, # IntersectionDataset kwargs
    dataset2 = chesapeake_labels, # IntersectionDataset kwargs
    collate_fn = stack_samples, # IntersectionDataset kwargs
)

Making a change in my CustomGeoDataModule to add this:

if stage in ["fit", "validate"]:

It became like this:

class CustomGeoDataModule(GeoDataModule):
    def setup(self, stage: str) -> None:
        """Set up datasets.

        Args:
            stage: Either 'fit', 'validate', 'test', or 'predict'.
        """
        self.dataset = self.dataset_class(**self.kwargs)
        
        generator = torch.Generator().manual_seed(0)
        (
            self.train_dataset,
            self.val_dataset,
            self.test_dataset,
        ) = random_bbox_assignment(dataset, [0.6, 0.2, 0.2], generator)
        
        if stage in ["fit"]:
            self.train_batch_sampler = RandomBatchGeoSampler(
                self.train_dataset, self.patch_size, self.batch_size, self.length
            )
        if stage in ["fit", "validate"]:
            self.val_sampler = GridGeoSampler(
                self.val_dataset, self.patch_size, self.patch_size
            )
        if stage in ["test"]:
            self.test_sampler = GridGeoSampler(
                self.test_dataset, self.patch_size, self.patch_size
            )

The error changed to:

/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:208: UserWarning: num_workers>0, persistent_workers=False, and strategy=ddp_spawn may result in data loading bottlenecks. Consider setting persistent_workers=True (this is a limitation of Python .spawn() and PyTorch)
  rank_zero_warn(
---------------------------------------------------------------------------
ProcessRaisedException                    Traceback (most recent call last)
Cell In[78], line 14
     12 else:
     13     print('Starting training from scratch...')
---> 14     trainer.fit(
     15         model=task,
     16         datamodule = datamodule,
     17     )

File /opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py:608, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    606 model = self._maybe_unwrap_optimized(model)
    607 self.strategy._lightning_module = model
--> 608 call._call_and_handle_interrupt(
    609     self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
    610 )

File /opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py:36, in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs)
     34 try:
     35     if trainer.strategy.launcher is not None:
---> 36         return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
     37     else:
     38         return trainer_fn(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/lightning/pytorch/strategies/launchers/multiprocessing.py:113, in _MultiProcessingLauncher.launch(self, function, trainer, *args, **kwargs)
    110 else:
    111     process_args = [trainer, function, args, kwargs, return_queue]
--> 113 mp.start_processes(
    114     self._wrapping_function,
    115     args=process_args,
    116     nprocs=self._strategy.num_processes,
    117     start_method=self._start_method,
    118 )
    119 worker_output = return_queue.get()
    120 if trainer is None:

File /opt/conda/lib/python3.10/site-packages/torch/multiprocessing/spawn.py:197, in start_processes(fn, args, nprocs, join, daemon, start_method)
    194     return context
    196 # Loop on join until it returns True or raises an exception.
--> 197 while not context.join():
    198     pass

File /opt/conda/lib/python3.10/site-packages/torch/multiprocessing/spawn.py:160, in ProcessContext.join(self, timeout)
    158 msg = "\n\n-- Process %d terminated with the following error:\n" % error_index
    159 msg += original_trace
--> 160 raise ProcessRaisedException(msg, error_index, failed_process.pid)

ProcessRaisedException: 

-- Process 1 terminated with the following error:
Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/strategies/launchers/multiprocessing.py", line 139, in _wrapping_function
    results = function(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 650, in _fit_impl
    self._run(model, ckpt_path=self.ckpt_path)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 1112, in _run
    results = self._run_stage()
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 1191, in _run_stage
    self._run_train()
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 1204, in _run_train
    self._run_sanity_check()
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 1276, in _run_sanity_check
    val_loop.run()
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/loops/dataloader/evaluation_loop.py", line 152, in advance
    dl_outputs = self.epoch_loop.run(self._data_fetcher, dl_max_batches, kwargs)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/loops/epoch/evaluation_epoch_loop.py", line 137, in advance
    output = self._evaluation_step(**kwargs)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/loops/epoch/evaluation_epoch_loop.py", line 234, in _evaluation_step
    output = self.trainer._call_strategy_hook(hook_name, *kwargs.values())
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 1494, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/strategies/ddp_spawn.py", line 288, in validation_step
    return self.model(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1156, in forward
    output = self._run_ddp_forward(*inputs, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1103, in _run_ddp_forward
    inputs, kwargs = _to_kwargs(
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/utils.py", line 100, in _to_kwargs
    _recursive_to(inputs, device_id, use_side_stream_for_tensor_copies)
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/utils.py", line 92, in _recursive_to
    res = to_map(inputs)
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/utils.py", line 83, in to_map
    return list(zip(*map(to_map, obj)))
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/utils.py", line 87, in to_map
    return [type(obj)(i) for i in zip(*map(to_map, obj.items()))]
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/utils.py", line 87, in <listcomp>
    return [type(obj)(i) for i in zip(*map(to_map, obj.items()))]
TypeError: first argument must be callable or None

I was trying to run on GPU T4 x2.

[lcoandrade]

Actually there were 2 problems in my code.

The first one was using:

if stage in ["validate"]:

instead of:

if stage in ["fit", "validate"]:

I don't know why I can't only set validate as a stage in the setup override.

The second error was trying to use the code the way it is to perform a fit with 2 GPUs. It works when changing to uso one GPU.

So, summarizing. To work with a custom geodataset we need to:

1 - Subclass RasterDataset to something like this:

class NAIPImages(RasterDataset):
    filename_glob = "m_*.tif"
    is_image = True
    separate_files = False
    
class ChesapeakeLabels(RasterDataset):
    filename_glob = "m_*.tif"
    is_image = False
    separate_files = False

2 - Instantiate the datasets:

naip_root = os.path.join(INPUT_DIR, 'naip_images')
naip_images = NAIPImages(naip_root)
print(naip_images)

chesapeake_root = os.path.join(INPUT_DIR, "chesapeake_labels")
chesapeake_labels = ChesapeakeLabels(chesapeake_root)
print(chesapeake_labels)

3 - If needed, subclass GeoDatamodule to create a custom geo datamodule:

class CustomGeoDataModule(GeoDataModule):
    def setup(self, stage: str) -> None:
        """Set up datasets.

        Args:
            stage: Either 'fit', 'validate', 'test', or 'predict'.
        """
        self.dataset = self.dataset_class(**self.kwargs)
        
        generator = torch.Generator().manual_seed(0)
        (
            self.train_dataset,
            self.val_dataset,
            self.test_dataset,
        ) = random_bbox_assignment(dataset, [0.6, 0.2, 0.2], generator)
        
        if stage in ["fit"]:
            self.train_batch_sampler = RandomBatchGeoSampler(
                self.train_dataset, self.patch_size, self.batch_size, self.length
            )
        if stage in ["fit", "validate"]:
            self.val_sampler = GridGeoSampler(
                self.val_dataset, self.patch_size, self.patch_size
            )
        if stage in ["test"]:
            self.test_sampler = GridGeoSampler(
                self.test_dataset, self.patch_size, self.patch_size
            )

4 - Instantiate the custom datamodule:

datamodule = CustomGeoDataModule(
    dataset_class = type(dataset), # GeoDataModule kwargs
    batch_size = 10, # GeoDataModule kwargs
    patch_size = IMG_SIZE, # GeoDataModule kwargs
    length = SAMPLE_SIZE, # GeoDataModule kwargs
    num_workers = WORKERS, # GeoDataModule kwargs
    dataset1 = naip_images, # IntersectionDataset kwargs
    dataset2 = chesapeake_labels, # IntersectionDataset kwargs
    collate_fn = stack_samples, # IntersectionDataset kwargs
)

5 - Pass the datamodule to your trainer.

Thanks!

[adamjstewart]

I don't know why I can't only set validate as a stage in the setup override.

The fit step in Lightning runs both train and validation steps. So you need both train and val datasets/samplers.

[lcoandrade]

Thanks!

[tgoelles]

Did you get it running on 2 GPU? I get the same error TypeError: first argument must be callable or None

[adamjstewart]

If you can provide me with code/data to reproduce this, I can try to create a minimal reproducible example. I think it would be best to start by opening an issue with Lightning to see if this is either a bug in Lightning, or a bug in TorchGeo. I've used NonGeoDatasets with multiple GPUs before, but I don't think I've extensively tested GeoDatasets with multiple GPUs, so it's possible that rtree or rasterio doesn't like distributed training. Anyway, we'll know more if we ask the Lightning folks.

[tgoelles]

@adamjstewart

torchgeo_GPU_parallel_issue_new.zip

Here is data and a notebook with a minimal example. There I use the "ddp_notebook" strategy and get the same error as with "ddp" strategy.

edit: new file with the right notebook

[tgoelles]

I updated the zip file. Please use this one

[adamjstewart]

I just remembered I'm on my macOS laptop (no CUDA, no distributed) so I can't actually debug this. Maybe just open an issue with the Lightning folks: https://github.com/Lightning-AI/pytorch-lightning/issues/new/choose. Some of them know of TorchGeo, but they would definitely be very happy if you could find a way to reproduce this using only Lightning.