파이토치에서 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