Pytorch DataLoader - 选择类 STL10 数据集答案

【问题标题】：Pytorch DataLoader - Choose Class STL10 DatasetPytorch DataLoader - 选择类 STL10 数据集
【发布时间】：2018-12-22 09:49:44
【问题描述】：

是否可以仅在 PyTorch torchvision 的 STL10 数据集中的 class= 0 处拉取？我可以循环检查它们，但需要批量接收 0 类图像

# STL10 dataset
train_dataset = torchvision.datasets.STL10(root='./data/',
                                           transform=transforms.Compose([
                                               transforms.Grayscale(),
                                               transforms.ToTensor()
                                           ]),
                                           split='train',
                                           download=True)


# Data loader
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size, 
                                           shuffle=True)

for i, (images, labels) in enumerate(train_loader):
    if labels[0] == 0:...

根据 iacolippo 的回答进行编辑 - 现在正在运行：

# Set params
batch_size = 25
label_class = 0   # only airplane images

# Return only images of certain class (eg. airplanes = class 0)
def get_same_index(target, label):
    label_indices = []

    for i in range(len(target)):
        if target[i] == label:
            label_indices.append(i)

    return label_indices

# STL10 dataset
train_dataset = torchvision.datasets.STL10(root='./data/',
                                           transform=transforms.Compose([
                                               transforms.Grayscale(),
                                               transforms.ToTensor()
                                           ]),
                                           split='train',
                                           download=True)

# Get indices of label_class
train_indices = get_same_index(train_dataset.labels, label_class)

# Data loader
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           sampler=torch.utils.data.sampler.SubsetRandomSampler(train_indices))

【问题讨论】：

标签： pytorch torch torchvision

【解决方案1】：

如果您只想要一个类的样本，您可以从 Dataset 实例中获取具有相同类的样本的索引，例如

def get_same_index(target, label):
    label_indices = []

    for i in range(len(target)):
        if target[i] == label:
            label_indices.append(i)

    return label_indices

那么您可以使用SubsetRandomSampler 仅从一个类的索引列表中抽取样本

torch.utils.data.sampler.SubsetRandomSampler(indices)

【讨论】：