PyTorch 데이터 세트 만들기

pytorch에서 제공되는 데이터 세트이 아닌 다른 데이터 세트를 이용하고 싶을 때 데이터 로더를 작성하는 방법을 설명합니다. 예시로 sklearn에서 제공되는 데이터 세트를 사용합니다.

구현

import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms
from torch import optim
from torch.utils.data import DataLoader, Dataset, TensorDataset

# MNIST 데이터 세트 로드
from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split

mnist_X, mnist_y = fetch_openml ( 'mnist_784', version = 1, data_home = "./ mnist"return_X_y = True)

X_all = mnist_X.astype (np.float32) / 255 # 데이터를 [0,1]로 변환
y_all = mnist_y.astype (np.int32)

# 분할
X_train, X_test, y_train, y_test = train_test_split (X_all, y_all, train_size = 0.8, random_state = 1)

print (X_train.shape) # (56000, 784)

X_train = X_train.reshape (len (X_train), 1, 28, 28) .astype ( 'float32')
X_test = X_test.reshape (len (X_test), 1, 28, 28) .astype ( 'float32')

y_train = y_train.astype ( 'long')
y_test = y_test.astype ( 'long')

train_set = TensorDataset (torch.from_numpy (X_train) torch.from_numpy (y_train))
test_set = TensorDataset (torch.from_numpy (X_test) torch.from_numpy (y_test))

bs = 1000

train_loader = DataLoader (train_set, batch_size = bs, shuffle = True)
test_loader = DataLoader (test_set, batch_size = bs, shuffle = False)

해설

train_set = TensorDataset (torch.from_numpy (X_train) torch.from_numpy (y_train))
test_set = TensorDataset (torch.from_numpy (X_test) torch.from_numpy (y_test))

여기에서 텐서 데이터 세트를 작성하기 때문에, 입력 이미지 데이터 (X) (데이터 수, 채널 수, 높이, 폭)로 설정해야만 합니다.

따라서

X_train = X_train.reshape (len (X_train), 1, 28, 28) .astype ( 'float32')
X_test = X_test.reshape (len (X_test), 1, 28, 28) .astype ( 'float32')

로 변환합니다.

그리고 레이블은

y_train = y_train.astype ( 'long')
y_test = y_test.astype ( 'long')

처럼 long 형으로 설정합니다.