MONAI icon indicating copy to clipboard operation
MONAI copied to clipboard

DataLoader's collate_fn does not work with Compose transforms

Open nslay opened this issue 2 years ago • 1 comments

Describe the bug A clear and concise description of what the bug is.

DataLoader provides a collate_fn argument to pass a function that transforms a list of tensors (or dictionaries of tensors) into a list of tensors (or dictionaries of tensors) that have the same size to appease DataLoader. You can pass, for example, a lambda expression of pad_list_data_collate or a standalone PadListDataCollate transform. If you try to pass a Compose transform as collate_fn and DataLoader is operating on dictionaries, you will get an exception from Compose. The Compose transform can even be of a single PadListDataCollate transform.

To Reproduce Steps to reproduce the behavior:

  1. Go to '...'
  2. Install '....'
  3. Run commands '....'

The below works:

collate_fn = lambda batch : pad_list_data_collate(batch, mode="constant", constant_values=(-1,))

The below also works:

collate_fn = PadListDataCollate(mode="constant", constant_values=(-1,))

The below does not work

        collate_fn = Compose(
            [
                PadListDataCollate(mode="constant", constant_values=(-1,)),
            ]
        )

Expected behavior A clear and concise description of what you expected to happen.

I expected collate_fn option to operate with Compose wrapping a single PadListDataCollate transform just as it can with a standalone PadListDataCollate transform. I further expect it to work in non-trivial cases with more transforms chained together in the Compose.

Screenshots If applicable, add screenshots to help explain your problem.

You can see a backtrace here: https://github.com/Project-MONAI/MONAI/issues/6264#issuecomment-1494471024

Environment

Ensuring you use the relevant python executable, please paste the output of: Python 3.9.15

python -c 'import monai; monai.config.print_debug_info()'

================================ Printing MONAI config... ================================ MONAI version: 1.1.0 Numpy version: 1.24.2 Pytorch version: 1.12.1+cu113 MONAI flags: HAS_EXT = False, USE_COMPILED = False, USE_META_DICT = False MONAI rev id: a2ec3752f54bfc3b40e7952234fbeb5452ed63e3 MONAI file: /gpfs/gsfs12/users/layns/torch_monai_v100/lib/python3.9/site-packages/monai/init.py

Optional dependencies: Pytorch Ignite version: NOT INSTALLED or UNKNOWN VERSION. Nibabel version: 5.0.0 scikit-image version: NOT INSTALLED or UNKNOWN VERSION. Pillow version: 9.4.0 Tensorboard version: NOT INSTALLED or UNKNOWN VERSION. gdown version: NOT INSTALLED or UNKNOWN VERSION. TorchVision version: 0.13.1+cu113 tqdm version: 4.65.0 lmdb version: NOT INSTALLED or UNKNOWN VERSION. psutil version: NOT INSTALLED or UNKNOWN VERSION. pandas version: NOT INSTALLED or UNKNOWN VERSION. einops version: 0.6.0 transformers version: NOT INSTALLED or UNKNOWN VERSION. mlflow version: NOT INSTALLED or UNKNOWN VERSION. pynrrd version: NOT INSTALLED or UNKNOWN VERSION.

For details about installing the optional dependencies, please visit: https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies

================================ Printing system config... ================================ psutil required for print_system_info

================================ Printing GPU config... ================================ Num GPUs: 1 Has CUDA: True CUDA version: 11.3 cuDNN enabled: True cuDNN version: 8302 Current device: 0 Library compiled for CUDA architectures: ['sm_37', 'sm_50', 'sm_60', 'sm_70', 'sm_75', 'sm_80', 'sm_86'] GPU 0 Name: Tesla V100-SXM2-32GB GPU 0 Is integrated: False GPU 0 Is multi GPU board: False GPU 0 Multi processor count: 80 GPU 0 Total memory (GB): 31.7 GPU 0 CUDA capability (maj.min): 7.0

Additional context collate_fn should ideally be able to operate with Compose transforms because, for example, input tensors to some models need to be divisible by some power of 2. So it would be natural to chain PadListDataCollate with DivisiblePadd.

nslay avatar Apr 03 '23 15:04 nslay

To add to additional context, you could work around this and use DivisiblePadd on the Dataset side. Then the collate data spatial size should remain divisible by some value (e.g. 16). But I also need to do a custom transform on label after the collate (it changes the label tensor dimensions in weird ways!). Anyway, DivisiblePadd is not the best context example.

My current workaround is to just do this after getting the batch from DataLoader.

nslay avatar Apr 03 '23 15:04 nslay