[torchvision] datasets.MNIST 데이터 내려받기

 

- 목차

 

키워드.

• - torchvision
• - MNIST

 

들어가며.

 

이번 글에서는 "torchvision.datasets" 의 MNIST 모듈에 대해서 알아보려고 합니다.

"torchvision.datasets" 의 "MNIST" 모듈은 이미지 데이터들을 손쉽게 제공하는 기능이 존재합니다.

CNN 모델의 테스팅을 위해서 MNIST 이미지 데이터를 내려받을 경우를 위해서 관련 기능들을 정리합니다.

 

먼저 간단한 코드 예시는 아래와 같습니다.

torchvistion 의 datasets 모듈 내부에 MNIST Class 가 존재하구요.

이는 각종 Vision 과 관련된 Train/Test 데이터셋을 다루기 위해서 사용됩니다.

from torchvision import datasets
from torchvision.transforms import ToTensor

trainset = datasets.MNIST(
    root="./mnist",
    train=True,
    transform=ToTensor(),
    download=True
)

 

Arguments 알아보기.

기본적으로 아래의 모듈들은 import 합니다.

from torchvision import datasets
from torchvision.transforms import ToTensor

 

torchvision 모듈을 import 한 이후에 아래와 같이 데이터들을 내려받을 수 있게 됩니다.

trainset 은 MNIST 클래스의 객체로 표현됩니다.

실질적인 Image Tensor 는 trainset 객체의 data 변수에 할당되어 있구요.

shape 는 60000, 28, 28 로 표현됩니다.

즉, 6만개의 이미지 케이스와 28x28 의 사이즈의 이미지를 의미합니다.

trainset = datasets.MNIST(
    root="./mnist",
    train=True,
    transform=ToTensor(),
    download=True
)

type(trainset)
--> <class 'torchvision.datasets.mnist.MNIST'>

type(trainset.data)
--> <class 'torch.Tensor'>

trainset.data.shape
--> torch.Size([60000, 28, 28])

 

root.

root 인자는 이미지를 저장할 폴더의 위치를 지정할 수 있습니다.

저의 경우에는 "./mnist" 와 같이 실행 스크립트와 동일한 위치에 데이터를 저장할 폴더를 지정합니다.

저의 실행파일의 이름은 test_load_mnist.py 이구요.

test_load_mnist.py 와 동일한 폴더에 mnist 폴더가 생성됩니다.

├── mnist
│   └── MNIST
│       └── raw
│           ├── t10k-images-idx3-ubyte
│           ├── t10k-images-idx3-ubyte.gz
│           ├── t10k-labels-idx1-ubyte
│           ├── t10k-labels-idx1-ubyte.gz
│           ├── train-images-idx3-ubyte
│           ├── train-images-idx3-ubyte.gz
│           ├── train-labels-idx1-ubyte
│           └── train-labels-idx1-ubyte.gz
├── test_load_mnist.py

 

train.

train 인자는 Trainset 의 MNIST 이미지 데이터를 내려받을지, Testset 의 MNIST 이미지 데이터를 내려받을지 결정합니다.

그래서 아래와 같이 train=True/False 설정을 통해서 결정이 가능합니다.

trainset = datasets.MNIST(
    root="./mnist2",
    train=True,
    transform=ToTensor(),
    download=True
)

testset = datasets.MNIST(
    root="./mnist2",
    train=False,
    transform=ToTensor(),
    download=True
)

 

trainset 과 testset 의 Shape 는 아래와 같습니다.

trainset.data.shape : torch.Size([60000, 28, 28]) 
testset.data.shape  : torch.Size([10000, 28, 28])

 

transform.

transform 인자는 내려받은 MNIST 이미지 데이터를 일차적으로 변형할 수 있는 기능을 제공합니다.

일반적으로 ToTensor Transformation 을 적용하여 MNIST 이미지 데이터를 Tensor 로 변형하게 됩니다.

 

내려받은 이미지 확인해보기.

아래 코드를 통해서 내려받은 이미지를 matplotlib 라이브러리를 통해 시각화할 수 있습니다.

 

from torchvision import datasets
from torchvision.transforms import ToTensor, ToPILImage
import matplotlib.pyplot as plt

trainset = datasets.MNIST(
    root="./mnist",
    train=True,
    transform=ToTensor(),
    download=True
)


to_image = ToPILImage()

fig = plt.figure()
col_count = 5
row_count = 5
for index, image_tensor in enumerate(trainset.data[0:25]):
    sub_plot = fig.add_subplot(row_count, col_count, index + 1)
    img = to_image(image_tensor)
    sub_plot.imshow(img)
plt.show()