Is it possible to pull only where class = 0 in the STL10 dataset in PyTorch torchvision
? I am able to check them in a loop, but need to receive batches of class 0 images
# 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:...
edit based on iacolippo's answer - this is now working:
# 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))
If you only want samples from one class, you can get the indices of samples with the same class from the Dataset
instance with something like
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
then you can use SubsetRandomSampler
to draw samples only from the list of indices of one class
torch.utils.data.sampler.SubsetRandomSampler(indices)