How do you load MNIST images into Pytorch DataLoader?

Question

The pytorch tutorial for data loading and processing is quite specific to one example, could someone help me with what the function should look like for a more generic simple loading of images?

Tutorial: http://pytorch.org/tutorials/beginner/data_loading_tutorial.html

My Data:

I have the MINST dataset as jpg's in the following folder structure. (I know I can just use the dataset class, but this is purely to see how to load simple images into pytorch without csv's or complex features).

The folder name is the label and the images are 28x28 png's in greyscale, no transformations required.

Duane · Accepted Answer · 2019-12-10 05:30:15Z

44

Here's what I did for pytorch 0.4.1 (should still work in 1.3)

def load_dataset():
    data_path = 'data/train/'
    train_dataset = torchvision.datasets.ImageFolder(
        root=data_path,
        transform=torchvision.transforms.ToTensor()
    )
    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=64,
        num_workers=0,
        shuffle=True
    )
    return train_loader

for batch_idx, (data, target) in enumerate(load_dataset()):
    #train network

edited Dec 10, 2019 at 5:30

answered Aug 5, 2018 at 20:31

Duane

5,2307 gold badges36 silver badges36 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

Arturo Over a year ago

How is the class label specified in your load_dataset() function?

Sami Over a year ago

It's generated by ImageFolder depending on the class folder: pytorch.org/docs/stable/torchvision/…

Andrey Over a year ago

For MNIST It's may be necessary to use "transforms.Grayscale()" : test_dataset = torchvision.datasets.ImageFolder( root=data_path, transform=transforms.Compose([transforms.Grayscale(), transforms.ToTensor()]) )

Ari K · Accepted Answer · 2018-04-26 22:00:53Z

13

If you're using mnist, there's already a preset in pytorch via torchvision.
You could do

import torch
import torchvision
import torchvision.transforms as transforms
import pandas as pd

transform = transforms.Compose(
[transforms.ToTensor(),
 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

mnistTrainSet = torchvision.datasets.MNIST(root='./data', train=True,
                                    download=True, transform=transform)
mnistTrainLoader = torch.utils.data.DataLoader(mnistTrainSet, batch_size=16,
                                      shuffle=True, num_workers=2)

If you want to generalize to a directory of images (same imports as above), you could do

class mnistmTrainingDataset(torch.utils.data.Dataset):

    def __init__(self,text_file,root_dir,transform=transformMnistm):
        """
        Args:
            text_file(string): path to text file
            root_dir(string): directory with all train images
        """
        self.name_frame = pd.read_csv(text_file,sep=" ",usecols=range(1))
        self.label_frame = pd.read_csv(text_file,sep=" ",usecols=range(1,2))
        self.root_dir = root_dir
        self.transform = transform

    def __len__(self):
        return len(self.name_frame)

    def __getitem__(self, idx):
        img_name = os.path.join(self.root_dir, self.name_frame.iloc[idx, 0])
        image = Image.open(img_name)
        image = self.transform(image)
        labels = self.label_frame.iloc[idx, 0]
        #labels = labels.reshape(-1, 2)
        sample = {'image': image, 'labels': labels}

        return sample


mnistmTrainSet = mnistmTrainingDataset(text_file ='Downloads/mnist_m/mnist_m_train_labels.txt',
                                   root_dir = 'Downloads/mnist_m/mnist_m_train')

mnistmTrainLoader = torch.utils.data.DataLoader(mnistmTrainSet,batch_size=16,shuffle=True, num_workers=2)

You can then iterate over it like:

for i_batch,sample_batched in enumerate(mnistmTrainLoader,0):
    print("training sample for mnist-m")
    print(i_batch,sample_batched['image'],sample_batched['labels'])

There are a bunch of ways to generalize pytorch for image dataset loading, the method that I know of is subclassing torch.utils.data.dataset

edited Apr 26, 2018 at 22:00

answered Apr 26, 2018 at 21:51

Ari K

4442 silver badges18 bronze badges

2 Comments

moi Over a year ago

Loading the same file and accessing its two columns two times independently is highly inefficient!

moi Over a year ago

Yes. Rather use one single data frame. Load it with index_col=False in read_csv to obtain a numeric index. Then use self.df.at[idx, "filename"] and self.df.at[idx, "label"] in __getitem__.

Collectives™ on Stack Overflow

How do you load MNIST images into Pytorch DataLoader?

2 Answers 2

3 Comments

2 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

3 Comments

2 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related