Custom Image Folder

Hello!

I want to create a custom ImageFolder were instead of the default (sample, target), it will also return an annotated version of each data sample. Therefore (sample, annotated_sample, target).

Below

class AugmentatedImageFolder(ImageFolder):
    """Augmentated Images dataset."""

    def __getitem__(self, index: int):
        """
        Args:
            index (int): Index

        Returns:
            tuple: (sample, augmented_sample, target) where target is class_index of the target class.
        """
        path, target = self.samples[index]
        sample = self.loader(path)

        if self.transform is not None:
            sample_augmented = self.transform(sample)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return sample, sample_augmented, target

Is the above approach correct?

I run my experiments in an HPC, and I get the below memory error when I do this customization:
slurmstepd: error: Detected 2 oom_kill events in StepId=1819449.batch. Some of the step tasks have been OOM Killed.

Is there something in the custom class that could be affecting the memory?

Thank you! :slight_smile:

It’s unclear what is causing the OOM issues but depending on the sample size and batch size the data loading could be responsible for it.
Your code looks generally alright, but you should receive an error if the if conditions are not met since the returned variables would be undefined.

Hello @ptrblck , thank you for your response and your points.

I had some extra lines in the class that i skipped in my post, where I was composing multiple transforms using transorms.Compose and the applied them to the data. It seems that this was causing the OOM issue, so I used another approach instead, and it worked fine.

Good to hear it’s working now! Just out of interest, which transformations were you using?