MarStarck
MarStarck

Reputation: 443

How to create Datasets Like MNIST in Pytorch?

I have looked Pytorch source code of MNIST dataset but it seems to read numpy array directly from binaries. How can I just create train_data and train_labels like it? I have already prepared images and txt with labels.

I have learned how to read image and label and write get_item and len, what really confused me is how to make train_data and train_labels, which is torch.Tensor. I tried to arrange them into python lists and convert to torch.Tensor but failed:

for index in range(0,len(self.files)):
  fn, label = self.files[index]
  img = self.loader(fn)
  if self.transform is not None:
    img = self.transform(img)
  train_data.append(img)
self.train_data = torch.tensor(train_data)

ValueError: only one element tensors can be converted to Python scalars

Upvotes: 4

Views: 2613

Answers (1)

Jatentaki
Jatentaki

Reputation: 13103

There are two ways to go. First, the manual. Torchvision.datasets states the following:

datasets are subclasses of torch.utils.data.Dataset i.e, they have __getitem__ and __len__ methods implemented. Hence, they can all be passed to a torch.utils.data.DataLoader which can load multiple samples parallelly using torch.multiprocessing workers.

So you can just implement your own class which scans for all the images and labels, keeps a list of their paths (so that you don't have to keep them in RAM) and has the __getitem__ method which given index i reads the i-th file, its label and returns them. This minimal interface is enough to work with the parallel dataloader in torch.utils.data.

Secondly, if your data directory can be rearranged into either structure, you can use DatasetFolder and ImageFolder pre-built loaders. This will save you some coding and automatically provide support for data augumentation routines from torchvision.transforms.

Upvotes: 2

Related Questions