如何在pytorch中从多个数据集加载数据

Question

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

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

我正在尝试训练一个模型，在那里我告诉它，"；这是一张包含猫的图片，告诉我它的颜色；无论它被带到哪里(室内、室外、车内、月球上(

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

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

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

Answer 1

假设问题是：

1. 将2个以上的数据集与可能重叠的对象类别相结合(可通过标签区分(
2. 每个对象具有4"；子类别"；每种颜色(可通过标签区分(
3. 每个批次应仅包含一个对象类别

第一步是确保两个数据集的对象标签的一致性，如果还不一致的话。例如，如果dog类在第一个数据集中是标签0，但在第二个数据集中却是标签2，那么我们需要确保这两个dog类别正确合并。我们可以做到这一点"翻译"；使用一个简单的数据集包装器：

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])

最后，我们需要将批处理采样限制在同一个类中，这在创建数据加载器时使用自定义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