how to load two dataset images simultaneously for train two streams(Pytorch)

Question

i need load identical two dataset suppose one dataset has RGB images and another dataset contain same image with different processed(grey images) with same order same size,

datasetA=[1.jpg,2.jpg,..........n.jpg] // RGB
datasetA=[g1.jpg,g2.jpg,..........gn.jpg] //grey

so I need to feed the same order images to two independent networks using DataLoader with random_split, so how to use

rgb = datasets.ImageFolder(rgb images)
grey = datasets.ImageFolder(gray images)

train_data1, test_data = random_split(rgb, [train_data, test_data])
train_data2, test_data = random_split(grey, [train_data, test_data])

train_loader1 = DataLoader(train_data1, batch_size=batch_size, shuffle=True)
train_loader2 = DataLoader(train_data2, batch_size=batch_size, shuffle=True)

so need to load same order images touple like (1.jpg,g1.jpg) for train both network independantly

and how to use

trainiter1 = iter(train_loader1)
features, labels = next(trainiter)

please explain process

Answer 1

I think he easiest way to go about this is to construct a custom Dataset that handles both:

class JointImageDataset(torch.utils.data.Dataset):
  def __init__(self, args_rgb_dict, args_grey_dict):
    # construct the two individual datasets
    self.rgb_dataset = ImageFolder(**args_rgb_dict)
    self.grey_dataset = ImageFolder(**args_grey_dict)

  def __len__(self):
    return min(len(self.rgb_dataset), len(selg.grey_dataset))

  def __getitem__(self, index):
    rgb_x, rgb_y = self.rgb_dataset[index]
    grey_x, grey_y = self.grey_dataset[index]
    return rgb_x, grey_x, rgb_y, grey_y

Now you can construct a single DataLoader from the JoindImageDataset and iterate over the joint batches:

joint_data = JoindImageDataset(...)
train_loader = DataLoader(joint_data, batch_size=...)
for rgb_batch, grey_batch, rgb_ys, grey_ys in train_loader:
  # do your stuff here...