📂머신러닝

파이토치에서 Numpy 배열로 커스텀 데이터 셋 만들고 사용하는 방법

설명

>>> import numpy as np
>>> import torch
>>> from torch.utils.data import TensorDataset, DataLoader

$32\times 32$의 크기로 된 ‘흑백’사진 100장을 쌓아올린 numpy 배열 $X$와 이에 대한 레이블 $Y$가 준비되어있다고 가정하자. 이를 다음가 같은 코드로 불러왔다고 하자.

>>> X = np.load("X.npy")
>>> X.shape
(100, 32, 32)
>>> Y = np.load("Y.npy")
>>> Y.shape
(100)

학습에 사용할 데이터 셋을 만들기 위해 다음과 같이 numpy배열을 tensor로 바꿔준다.

>>> X = torch.from_numpy(X)
>>> Y = torch.from_numpy(Y)

TensorDataset으로 학습할 데이터를 묶는다. 물론 3 종류 이상의 데이터를 묶어 데이터 셋을 만드는 것도 가능하다.

>>> train_data = TensorDataset(X, Y)

DataLoader로 데이터 순서를 무작위로 섞어 원하는 배치 사이즈로 만들 수 있다. 가령 배치 사이즈를 25로 설정하면 총 100개이므로 $25 \times 32 \times 32$짜리 데이터가 4개 생기는 것이다.

>> train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)

이를 다음과 같이 학습 함수를 정의하여 사용할 수 있다.

def train(model):
    model.train()
            
    for batch_idx, (X, Y) in enumerate(Data_set):
        
        optimizer.zero_grad()

        output = model(X)
        
        loss = criterion(output, Y)
        loss.backward()
        optimizer.step()

    return loss

전체 코드

적당한 model이 정의되어있다고 가정하자.

#필요한 라이브러리 import
import numpy as np
import torch
from torch.utils.data import TensorDataset, DataLoader

#numpy 배열 불러오기
X = np.load("X.npy")
Y = np.load("Y.npy")

#tensor로 데이터 타입 변환
X = torch.from_numpy(X)
Y = torch.from_numpy(Y)

#데이터셋으로 변환
train_data = TensorDataset(X, Y)

#셔플, 배치사이즈 적용
train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)

#학습 함수 정의
def train():
    model.train()
            
    for batch_idx, (X, Y) in enumerate(Data_set):
        
        optimizer.zero_grad()

        output = model(X)
        
        loss = criterion(output, Y)
        loss.backward()
        optimizer.step()

    return loss

#학습
for epoch in range(1, EPOCHS):
    loss = train()
    
    print(loss)

환경

• OS: Windows10
• Version: Python 3.9.2, torch 1.8.1+cu111