PyTorch 프레임워크를 사용해 Tensorflow때와 같이 MNIST 예제를 통해 이미지 데이터를 처리할 것이다.

라이브러리 불러오기

import torch
from torchvision import datasets, transforms

Batch Size 정하기

batch_size = 32
test_batch_size = 32

Data Loader 부르기

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_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= test_batch_size,
    shuffle=True)

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

파이토치는 DataLoader를 불러 model에 넣음

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

데이터 시각화

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

images[0].shape # 이미지 shape

자원수 줄이기

torch_image = torch.squeeze(images[0]) # 앞에 차원수 줄이기
torch_image.shape

label = label[0].numpy()

label.shape # 확인 결과 ()

label 변수를 확인해보면 1의 값이 있다

이미지 띄우기

plt.imshow(torch_image)
plt.title(label)
plt.show()