如何同时加载两个数据集图像以训练两个流（Pytorch）

Question

我需要加载相同的两个数据集，假设一个数据集有 RGB 图像，另一个数据集包含相同的图像，不同的处理（灰色图像）具有相同的顺序相同的大小，

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

所以我需要使用带有 random_split 的 DataLoader 将相同的订单图像提供给两个独立的网络，那么如何使用

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)

因此需要加载相同顺序的图像，例如 (1.jpg,g1.jpg) 以独立训练两个网络

以及如何使用

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

请说明过程

Answer 1

我认为解决此问题的最简单方法是构建一个同时处理两者的custom Dataset：

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

现在您可以从JoindImageDataset 构造一个单个 DataLoader 并迭代联合批次：

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...