Hi there,
I am implementing a Siamese Neural Network and therefore wrote a custom Dataset for it. The Dataset is implemented to fit my folder structure.
The folder structure looks like the following:
root/
  0001/
     0001-normal.jpg
     0001-blurry.jpg
     0001-shifted.jpg
     ....
   0002/
     ....

In total I have more than 6000 folders with 8 images of the same person.
Here’s my implementation of the Dataset:
class SiameseDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        super().__init__()
        self.h_table = {}
        self.root_dir = root_dir
        self.folders = os.listdir(root_dir)
        self.folders.sort()
        self.num_folders = len(self.folders) - 1
        # store all files in h_table[label] as list
        for folder in self.folders:
            self.h_table[str(int(folder) - 1)] = os.listdir("{}/{}".format(root_dir, folder))

    def __len__(self):
        return len(self.h_table)

    def __getitem__(self, index):
        same = random.uniform(0,1) > 0.5
        h_len = len(self.h_table[str(index)]) - 1
        h_ = self.h_table[str(index)]
        if same:
            first_idx = random.randint(0, h_len)
            while True:
                second_idx = random.randint(0, h_len)
                if first_idx != second_idx:
                    break

            first_path = "{}/{:04d}/{}".format(self.root_dir, index + 1, h_[first_idx])
            second_path = "{}/{:04d}/{}".format(self.root_dir, index + 1, h_[second_idx])

            first = plt.imread(first_path)
            second = plt.imread(second_path)

            first = transforms.ToTensor()(first)
            second = transforms.ToTensor()(second)

            return (first, index), (second, index)
        else:
            first_idx = random.randint(0, h_len)
            while True:
                sec_class_idx = random.randint(0, self.num_folders)
                if sec_class_idx != index:
                    break
            second_idx = random.randint(0, h_len)
            second = self.h_table[str(sec_class_idx)][second_idx]

            first_path = "{}/{:04d}/{}".format(self.root_dir, index + 1, h_[first_idx])
            second_path = "{}/{:04d}/{}".format(self.root_dir, index + 1, h_[second_idx])

            first = plt.imread(first_path)
            second = plt.imread(second_path)

            first = transforms.ToTensor()(first)
            second = transforms.ToTensor()(second)

            return (first, index), (second, sec_class_idx)

I think the code is very simple. If same is True, the __getitem__ method simply takes two random images of the same folder, applies a ToTensor transformation and returns both images with the index each. If same is False, then it takes a random image of the folder at the given index, chooses another random folder index and takes another random image of that folder.
The code works if I do not use a DataLoader.
root_dir = "data/256x256/1"
dataset = SiameseDataset(root_dir, device)
(x1, y1), (x2, y2) = dataset[0] # <- this works

But if I use a DataLoader object and try to iterate over it in a training loop for my Siamese Network it simply loads forever.
loader = DataLoader(dataset, batch_size=128)
for i, data in enumerate(loader): # <- this takes forever
    # do stuff

I’m working on a large machine provided by my university so performance should not be an issue and I never had issues like this before, when I used the ImageFolder dataset class.
I tried opening all images and saving them in h_table, such that they are already in memory. But that did not help either.
I also tried calling __getitem__ manually. This works without issues.
for i in range(10):
    (x1,y1), (x2,y2) = dataset.__getitem__(i)
    print(y1,y2)

This returns:

0 585
1 5848
2 262
3 3
4 4001
5 5
6 1696
7 2027
8 8
9 9