TypeError: list indices must be integers or slices, not DefaultDataKeys when training Object Detection Model

🐛 Bug

I’ve spent days making the data augmentation work for Object Detection but errors keep poping up. I don’t know if I’m reinventing the wheels or you are missing a lot in term data preparation/augmentation documentation for object detection. I’m about to give up…

Following #409 (always not resolved) I’ve created a custom data augmentation transformation using albumentations. However it fails with a weird message when starting training (when we fix the error I can make a PR for integrating albumentations with pytorch lightning flash):

File "train.py", line 93, in train
    trainer.finetune(model, datamodule=datamodule)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/flash/core/trainer.py", line 148, in finetune
    return super().fit(model, train_dataloader, val_dataloaders, datamodule)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 458, in fit
    self._run(model)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 756, in _run
    self.dispatch()
  File "/home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 797, in dispatch
    self.accelerator.start_training(self)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 96, in start_training
    self.training_type_plugin.start_training(trainer)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 144, in start_training
    self._results = trainer.run_stage()
  File "/home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 807, in run_stage
    return self.run_train()
  File "/home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 842, in run_train
    self.run_sanity_check(self.lightning_module)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1107, in run_sanity_check
    self.run_evaluation()
  File "/home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 962, in run_evaluation
    output = self.evaluation_loop.evaluation_step(batch, batch_idx, dataloader_idx)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 174, in evaluation_step
    output = self.trainer.accelerator.validation_step(args)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 226, in validation_step
    return self.training_type_plugin.validation_step(*args)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp.py", line 322, in validation_step
    return self.model(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/nn/parallel/distributed.py", line 705, in forward
    output = self.module(*inputs[0], **kwargs[0])
  File "/usr/lib/python3/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/overrides/base.py", line 57, in forward
    output = self.module.validation_step(*inputs, **kwargs)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/flash/image/detection/model.py", line 179, in validation_step
    images, targets = batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET]
TypeError: list indices must be integers or slices, not DefaultDataKeys

Before that it was failing with RuntimeError: each element in list of batch should be of equal size but this torch vision tip of custom collate : lambda x:x “fixes” it https://github.com/pytorch/vision/issues/2624

What is going on?

To Reproduce

import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2
from PIL import Image
import cv2

import flash
from flash.core.data.utils import download_data
from flash.image import ObjectDetectionData, ObjectDetector
from pytorch_lightning import seed_everything
import numpy

import logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

seed_everything(42)

   image_size = 1024

    train_transform = A.Compose(
        [
            A.Resize(height=image_size, width=image_size, p=1),
            A.OneOf([
                A.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=20,
                                     val_shift_limit=20, p=0.5),
                A.RandomBrightnessContrast(brightness_limit=0.2,
                                           contrast_limit=0.2, p=0.5),
            ], p=0.9),
            A.ToGray(p=0.01),
            A.VerticalFlip(p=0.5),
            A.HorizontalFlip(p=0.5),
            A.ShiftScaleRotate(p=0.5),
            A.Cutout(num_holes=10, max_h_size=32, max_w_size=32, fill_value=0, p=0.5),
            ToTensorV2(p=1)
        ],
        p=1.0,
        bbox_params=A.BboxParams(
            format='pascal_voc',
            min_area=0,
            min_visibility=0,
            label_fields=['labels']
        )
    )

    valid_transform = A.Compose(
        [
            A.Resize(height=image_size, width=image_size, p=1),
            ToTensorV2(p=1)
        ],
        p=1.0,
        bbox_params=A.BboxParams(
            format='pascal_voc',
            min_area=0,
            min_visibility=0,
            label_fields=['labels']
        )
    )

    test_transform = A.Compose(
        [
            A.Resize(height=image_size, width=image_size, p=1),
            ToTensorV2(p=1)
        ],
        p=1.0,
        bbox_params=A.BboxParams(
            format='pascal_voc',
            min_area=0,
            min_visibility=0,
            label_fields=['labels']
        )
    )

    datamodule = ObjectDetectionData.from_coco(
        train_folder="data_coco/train",
        train_ann_file="data_coco/train/_annotations.coco.json",
        train_transform={
            'pre_tensor_transform': lambda sample: transform_using_albu(sample, train_transform),
            'collate' : lambda x: x
        },
        val_transform={
         'pre_tensor_transform': lambda sample: transform_using_albu(sample, valid_transform),
          'collate': lambda x: x
        },
        test_transform={
         'pre_tensor_transform': lambda sample: transform_using_albu(sample, test_transform),
         'collate': lambda x: x
       },
        val_split=0.2,
        batch_size=8,
        num_workers=4,
    )

    model = ObjectDetector(model="retinanet", backbone="resnet101", num_classes=datamodule.num_classes, fpn=True)

    # 4. Create the trainer
    trainer = flash.Trainer(max_epochs=1, gpus=2, accelerator='ddp', limit_train_batches=1, limit_val_batches=1, checkpoint_callback=True)

    # 5. Finetune the model
    trainer.finetune(model, datamodule=datamodule)


def transform_using_albu(sample, train_transform):
        labels = sample['target']['labels']
        image = to_cv(sample['input'])
        transformed = train_transform(image=image, bboxes=sample['target']['boxes'], labels=sample['target']['labels'])
        trans_bboxes = [list(boxes) for boxes in transformed["bboxes"]]
        area = [calculate_area(boxes) for boxes in trans_bboxes]
        return {
            'input': transformed["image"],
            'target': {
              'boxes': trans_bboxes,
              'labels': labels,
              'image_id': sample['target']['image_id'],
              'area': area,
              'iscrowd': [0 for _ in trans_bboxes]}
            }

Environment

• PyTorch Version: 1.8
• OS (e.g., Linux): MacOS
• How you installed PyTorch: pip
• Python version: 3.7
• CUDA/cuDNN version: 11
• GPU models and configuration: 2 A 6000

Additional context

It is necessary to provide a clear and working example of augmenting and resizing images for object detection using torchvision transformers or albumentations.

[ethanwharris]

Ah yes, this one is caused by the most recent pytorch lightning version. They are working on a patch, but if you could do pip install 'pytorch-lightning==1.3.6' that should fix the error 😃

[ethanwharris]

Hi @hzitoun Thanks for the issue! Sorry that your experience with Flash hasn’t been very smooth. Using a custom collate is the right idea, but rather than an identity function you need the following:

def collate(samples: Sequence[Dict[str, Any]]) -> Dict[str, Sequence[Any]]:
    return {key: [sample[key] for sample in samples] for key in samples[0]}

Your collate gives an output like this:

[{sample 1}, {sample 2}, {sample 3}, ...]

The new collate does this:

{key1: [sample1[key1], sample2[key1], ...], key2: [sample1[key2], sample2[key2], ...], ...}

Hope that helps 😃