[PyTorch] 데이터 불러오기

PyTorch Data Preprocess

import torch
from torchvision import datasets, transforms

Import Error

ImportError: cannot import name 'PILLOW_VERSION' from 'PIL'

pillow 버전이 7.0.0 이상 일경우 Import 에러 나는 경우가 있습니다.
아래 처럼 pillow 버전을 내려주면 해결이 됩니다.

$ pip install pillow==6.2.2

Data Loader 부르기

Pytorch는 DataLoader를 불러 model에 넣습니다.

batch_size = 32

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('dataset/', train=True, download=True,
                transform=transforms.Compose([
                  transforms.ToTensor(),
                  transforms.Normalize(mean=(0.5,), std=(0.5,))
                ])),
                batch_size=batch_size,
                shuffle=True)

test_batch_size = 32

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('dataset', train=False,
                transform=transforms.Compose([
                  transforms.ToTensor(),
                  transforms.Normalize((0.5,), (0.5))
                ])),
                batch_size=test_batch_size,
                shuffle=True)

첫번째 iteration에서 나오는 데이터 확인

images, labels = next(iter(train_loader))
image.shape, label.shape

=> torch.Size([32, 1, 28, 28]),  torch.Size([32])

데이터 시각화

import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

# squeeze() 함수는 차원의 원소가 1인 차원을 없애줍니다.
torch_image = torch.squeeze(images[0])
torch_image.shape

=> torch.Size([28, 28])

image = torch_image.numpy()
image.shape

=> (28, 28)

plt.title(label)
plt.imshow(image, 'gray')
plt.show()