如何从pytorch中的乘法数据集加载数据

Question

我有两个图像数据集 - 室内和室外，它们没有相同数量的示例。

每个数据集都有图像包含一定数量的classes（最少1个最多4个），这些classes可以出现在两个数据集中，每个class有4个类别 - 红色、蓝色、绿色、白色。例子：室内 - 猫、狗、马户外 - 狗、人

我正在尝试训练一个模型，我告诉它“这是一张包含猫的图像，告诉我它的颜色”，而不管它是在哪里拍摄的（室内、室外、车内、月亮）

为此，我需要展示我的模型示例，以便每个批次只有一个类别（猫、狗、马或人），但我想从包含这些对象的所有数据集（在本例中为两个）中采样并混合它们。我该怎么做？

必须考虑到每个数据集中的示例数量不同，并且某些类别出现在一个数据集中，而其他类别可能出现在多个数据集中。每个批次只能包含一个类别。

如果有任何帮助，我将不胜感激，我已经尝试解决这个问题好几天了。

Answer 1

假设问题是：

将 2 个以上的数据集与可能重叠的对象类别（可通过标签区分）组合起来
每个对象的每种颜色都有 4 个“子类别”（可通过标签区分）
每个批次应该只包含一个对象类别

第一步是确保来自两个数据集的对象标签的一致性，如果不一致的话。例如，如果狗 class 在第一个数据集中的标签为 0，但在第二个数据集中的标签为 2，那么我们需要确保正确合并两个狗类别。我们可以用一个简单的数据集包装器来完成这个“翻译”：

class TranslatedDataset(Dataset):
  """
  Args:
    dataset: The original dataset.
    translate_label: A lambda (function) that maps the original
      dataset label to the label it should have in the combined data set
  """
  def __init__(self, dataset, translate_label):
    super().__init__()
    self._dataset = dataset
    self._translate_label = translate_label

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

  def __getitem__(self, idx):
    inputs, target = self._dataset[idx]
    return inputs, self._translate_label(target)

下一步是将翻译后的数据集组合在一起，这可以通过 ConcatDataset:

轻松完成

first_original_dataset = ...
second_original_dataset = ...

first_translated = TranslateDataset(
  first_original_dataset, 
  lambda y: 0 if y is 2 else 2 if y is 0 else y, # or similar
)
second_translated = TranslateDataset(
  second_original_dataset, 
  lambda y: y, # or similar
)

combined = ConcatDataset([first_translated, second_translated])

最后，我们需要将批次采样限制为相同的 class，这可以在创建数据加载器时使用自定义 Sampler。

class SingleClassSampler(torch.utils.data.Sampler):
  def __init__(self, dataset, batch_size):
    super().__init__()
    # We need to create sequential groups
    # with batch_size elements from the same class
    indices_for_target = {} # dict to store a list of indices for each target
    
    for i, (_, target) in enumerate(dataset):
      # converting to string since Tensors hash by reference, not value
      str_targ = str(target)
      if str_targ not in indices_for_target:
        indices_for_target[str_targ] = []
      indices_for_target[str_targ] += [i]

    # make sure we have a whole number of batches for each class
    trimmed = { 
      k: v[:-(len(v) % batch_size)] 
      for k, v in indices_for_target.items()
    }

    # concatenate the lists of indices for each class
    self._indices = sum(list(trimmed.values()))
  
  def __len__(self):
    return len(self._indices)

  def __iter__(self):
    yield from self._indices

然后使用采样器：

loader = DataLoader(
  combined, 
  sampler=SingleClassSampler(combined, 64), 
  batch_size=64, 
  shuffle=True
)

我没有运行这个代码，所以它可能不完全正确，但希望它能让你走上正轨。

torch.utils.data Docs

如何从pytorch中的乘法数据集加载数据

How to load data from multiply datasets in pytorch

dataset

multiple-databases

pytorch

dataloader