[Python, Pytorch] TensorDataset & DataLoader

안녕하세요, 오늘은 Pytorch 내 TensorDataset과 DataLoader에 대해 알아보도록 하겠습니다.

 

코드를 구현하기 위한 과정은 다음과 같습니다.

• 임의의 numpy.narray 타입의 데이터 셋 생성(train_x, train_y)
• numpy.narray 타입의 데이터 셋을 torch.Tensor 타입으로 변환
• TensorDataset을 이용하여 데이터 셋 구성
• DataLoader를 이용하여 Batch로 구성

 

위의 과정에 대해 전체 코드는 다음과 같습니다.

 

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

train_x = np.random.randint(0, 256, (64, 3, 256, 256)).astype(np.uint8)
train_y = np.random.randint(0, 2, (64,)).astype(np.uint8)
print(type(train_x))

train_x_tensor = torch.tensor(train_x)
train_y_tensor = torch.tensor(train_y)
print(type(train_x_tensor))

train_dataset = TensorDataset(train_x_tensor, train_y_tensor)
print(type(train_dataset))

train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True)
print(type(train_dataloader))

for i, train_batch in enumerate(train_dataloader):
  train_batch_x, train_batch_y = train_batch
  print("Batch - {}".format(i))
  print(train_batch_x.size())
  print(train_batch_y.size(), "\n")

출력된 결과는 다음과 같습니다.

<class 'numpy.ndarray'>
<class 'torch.Tensor'>
<class 'torch.utils.data.dataset.TensorDataset'>
<class 'torch.utils.data.dataloader.DataLoader'>
Batch - 0
torch.Size([16, 3, 256, 256])
torch.Size([16]) 

Batch - 1
torch.Size([16, 3, 256, 256])
torch.Size([16]) 

Batch - 2
torch.Size([16, 3, 256, 256])
torch.Size([16]) 

Batch - 3
torch.Size([16, 3, 256, 256])
torch.Size([16])

---

 

위의 전체 코드에 대해 해당 코드들을 세부적으로 알아보겠습니다.

 

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

train_x = np.random.randint(0, 256, (64, 3, 256, 256)).astype(np.uint8)
train_y = np.random.randint(0, 2, (64,)).astype(np.uint8)
print(type(train_x))

train_x_tensor = torch.tensor(train_x)
train_y_tensor = torch.tensor(train_y)
print(type(train_x_tensor))

• np.random.randint() : numpy.ndarray 타입의 train_x, train_y 임의 생성
• torch.tensor() : numpy.ndarray 타입의 데이터 셋을 torch.Tensor 타입으로 변환

 

train_dataset = TensorDataset(train_x_tensor, train_y_tensor)
print(type(train_dataset))

train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True)
print(type(train_dataloader))

for i, train_batch in enumerate(train_dataloader):
  train_batch_x, train_batch_y = train_batch
  print("Batch - {}".format(i))
  print(train_batch_x.size())
  print(train_batch_y.size(), "\n")

• TensorDataset : torch.Tensor 타입의 데이터들을 사용해서 TensorDataet 구성
• DataLoader : TensorDataset을 batch_size=16을 통해 각 Batch 구성
• enumerate : 각 Batch 별 index 정보와 데이터 (train_batch_x, train_batch_y)를 불러오기
  **각 Batch 별 size()를 통해 torch.Size를 확인하였습니다.

 

지금까지, Pytorch 내 TensorDataset과 DataLoader를 사용하는 방법에 대해서 알아보았습니다.