Loading data in PyTorch

pytorch_imgloader.py

!pip install Pillow==4.0.0
!pip install PIL
!pip install image

import matplotlib.pyplot as plt
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, transforms

dataset = datasets.ImageFolder(
    '<path/to/data>',
    transforms.Compose([
        transforms.ToTensor(),
    ]))

data_loader = DataLoader(dataset=dataset,
                         batch_size=2,
                         shuffle=False,
                         num_workers=2)

for i, data in enumerate(data_loader):
  img, lbl = data
  print(img.numpy()[0].T.shape)
  plt.imshow(np.moveaxis(img.numpy()[0], 0, -1))
  break