Pytorch数据加载器-选择类STL10数据集
在PyTorchPytorch数据加载器-选择类STL10数据集,pytorch,torch,torchvision,Pytorch,Torch,Torchvision,在PyTorchtorchvision中,是否可以仅在STL10数据集中class=0的位置拉取?我能够在循环中检查它们,但需要接收批量的0类图像 # STL10 dataset train_dataset = torchvision.datasets.STL10(root='./data/', transform=transforms.Compose([
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))
如果您只需要一个类中的样本,那么可以从
数据集
实例中获取具有相同类的样本的索引,方法如下
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
然后,您可以使用substrandomsampler
仅从一个类的索引列表中提取样本
torch.utils.data.sampler.SubsetRandomSampler(indices)