파이토치, colab에서 MNIST파일 dataloader로 불러오기

2021. 3. 17. 06:48

pytorch에는 이미지 관련된 학습을 편리하게 도와주는 torchvision이라는 패키지가 있다. 이 패키지를 이용하면 테스트 목적의 데이터셋도 쉽게 불렁로 수 있다. dataset과 dataloader함수를 이용하면 된다. 테스트 목적으로 가장 많이 사용되는 MNIST의 경우, 아래와 같이 입력하면 dataset을 가져오고 배치를 위한 입력준비까지 쉽게 할 수 있다.

 

import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms

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

train_loader = torch.utils.data.DataLoader(datasets.MNIST('./data', train=True, download=True, transform=transform), batch_size = 64, shuffle=True)

 

하지만, 이를 코랩(colab)에서 사용했더니 'HTTP Error~'와 같이 에러가 발생한다. 에러 메시지를 보니, 관련 IP가 차단된 것으로 보인다. 예전에는 됐던 것 같은데, 최근에는 download가 되지 못하도록 막힌 듯 하다.

 

 

pytorch MNIST

 

 

그래서 방법을 찾던 중에, colab을 실행하면 아래와 같이 sample_data가 있다는 사실을 확인할 수 있었다. 그래서 이 samle_data에 있는 mnist_train_small.csv를 DataLoader에 넣기로 하였다.

 

먼저 필요한 모듈을 import한다.

 

import pandas as pd
import numpy as np
from torch.utils.data import DataLoader

import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import Dataset

 

 

간단하게 DataLoader에 넣을 수 있다고 생각했는데, csv파일을 데이터와 target값으로 분리해서 넣어야 한다.

 

df_train = pd.read_csv('./sample_data/mnist_train_small.csv', header=None)
df_test = pd.read_csv('./sample_data/mnist_test.csv', header=None)

train_labels = df_train.iloc[:, 0]
train_images = df_train.iloc[:, 1:]
test_labels = df_test.iloc[:, 0]
test_images = df_test.iloc[:, 1:]

 

 

다음 transform을 정의해 준다. ReSize는 없어도 된다. 필자는 오토인코더에 학습시키려고 하는데, Resize를 하지 않으면 입력과 출력의 shape가 달라서 에러가 발생했다. 그래서 ReSize를 추가하였다.

 

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

 

 

다음으로 Dataset클래스를 상속한 MNISTDataset 클래스를 생성해준다.

 

class MNISTDataset(Dataset):
    def __init__(self, images, labels=None, transforms=None):
        self.X = images
        self.y = labels
        self.transforms = transforms
         
    def __len__(self):
        return (len(self.X))
    
    def __getitem__(self, i):

        data = self.X.iloc[i, :]
        data = np.asarray(data).astype(np.uint8).reshape(28, 28, 1)
        
        if self.transforms:
            data = self.transforms(data)
            
        if self.y is not None:
            return (data, self.y[i])
        else:
            return data


train_data = MNISTDataset(train_images, train_labels, transform)
test_data = MNISTDataset(test_images, test_labels, transform)

# dataloaders
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = DataLoader(test_data, batch_size=64, shuffle=True)

 

make_grid함수를 이용해서 MNIST결과를 출력해보았다. make_grid에 nrow는 하나의 행에 몇 개의 그림을 보여줄지 입력하는 파라미터이다.

 

images, labels = next(iter(train_loader))
grid_img = torchvision.utils.make_grid(images, nrow=5)

 

 

imshow함수를 이용해, 이미지를 확인해보았다.

 

import matplotlib.pyplot as plt

plt.imshow(grid_img.permute(1, 2, 0))

 

 

MNIST 이미지

 

 

오늘은 이렇게 파이토치, colab에서 MNIST CSV파일을 불러오는 방법에 대해서 알아보았다. 조금 번거롭기는 하지만, 이미 폴더에 있는 파일을 사용하는 것은 편리한 듯 하다.

댓글()