Preparing a custom geo dataset for training #1469

lcoandrade · 2023-07-10T14:12:56Z

lcoandrade
Jul 10, 2023

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!

Answered by lcoandrade

Jul 10, 2023

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 - In…

View full answer

adamjstewart · 2023-07-10T14:52:28Z

adamjstewart
Jul 10, 2023
Maintainer

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.

1 reply

lcoandrade Jul 10, 2023
Author

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 · 2023-07-10T15:25:52Z

lcoandrade
Jul 10, 2023
Author

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!

2 replies

adamjstewart Jul 10, 2023
Maintainer

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 Jul 10, 2023
Author

Thanks!

tgoelles · 2024-07-15T10:22:41Z

tgoelles
Jul 15, 2024

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

1 reply

adamjstewart Jul 16, 2024
Maintainer

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 · 2024-07-17T08:40:12Z

tgoelles
Jul 17, 2024

@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

2 replies

tgoelles Jul 17, 2024

I updated the zip file. Please use this one

adamjstewart Jul 18, 2024
Maintainer

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.

This comment has been hidden.

Sign in to view

This comment has been hidden.

Sign in to view

This comment has been hidden.

Sign in to view

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Preparing a custom geo dataset for training #1469

{{title}}

{{editor}}'s edit

{{editor}}'s edit

Replies: 5 comments 7 replies

This comment has been hidden.

This comment has been hidden.

{{title}}

{{title}}

{{title}}

{{editor}}'s edit

{{editor}}'s edit

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{editor}}'s edit

{{editor}}'s edit

{{title}}

{{title}}

Select a reply

Preparing a custom geo dataset for training #1469

lcoandrade Jul 10, 2023

Replies: 5 comments · 7 replies

This comment has been hidden.

This comment has been hidden.

adamjstewart Jul 10, 2023 Maintainer

lcoandrade Jul 10, 2023 Author

lcoandrade Jul 10, 2023 Author

adamjstewart Jul 10, 2023 Maintainer

lcoandrade Jul 10, 2023 Author

tgoelles Jul 15, 2024

adamjstewart Jul 16, 2024 Maintainer

tgoelles Jul 17, 2024

tgoelles Jul 17, 2024

adamjstewart Jul 18, 2024 Maintainer

lcoandrade
Jul 10, 2023

Replies: 5 comments 7 replies

adamjstewart
Jul 10, 2023
Maintainer

lcoandrade Jul 10, 2023
Author

lcoandrade
Jul 10, 2023
Author

adamjstewart Jul 10, 2023
Maintainer

lcoandrade Jul 10, 2023
Author

tgoelles
Jul 15, 2024

adamjstewart Jul 16, 2024
Maintainer

tgoelles
Jul 17, 2024

adamjstewart Jul 18, 2024
Maintainer