Get single random example from PyTorch DataLoader

If your DataLoader is something like this:

test_loader = DataLoader(image_datasets['val'], batch_size=batch_size, shuffle=True)

it is giving you a batch of size batch_size, and you can pick out a single random example by directly indexing the batch:

for test_images, test_labels in test_loader:  
    sample_image = test_images[0]    # Reshape them according to your needs.
    sample_label = test_labels[0]

Alternative solutions

  1. You can use RandomSampler to obtain random samples.

  2. Use a batch_size of 1 in your DataLoader.

  3. Directly take samples from your DataSet like so:

     mnist_test = datasets.MNIST('../MNIST/', train=False, transform=transform)
    

    Now use this dataset to take samples:

     for image, label in mnist_test:
          # do something with image and other attributes
    
  4. (Probably the best) See here:

     inputs, classes = next(iter(dataloader))   
    

If you want to choose specific images from your Trainloader/Testloader, you should check out the Subset function from master:

Here's an example how to use it:

testset = ImageFolderWithPaths(root="path/to/your/Image_Data/Test/", transform=transform)
subset_indices = [0] # select your indices here as a list
subset = torch.utils.data.Subset(testset, subset_indices)
testloader_subset = torch.utils.data.DataLoader(subset, batch_size=1, num_workers=0, shuffle=False)

This way you can use exactly one image and label. However, you can of course use more than just one index in your subset_indices.

If you want to use a specific image from your DataFolder, you can use dataset.sample and build a dictionary to get the index of the image you want to use.

Tags:

Python

Pytorch