[PyTorch] 데이터 로드 및 전처리 기본

학습 데이터와 테스트 데이터를 로드하는 방법과, 이를 PyTorch에서 다룰 수 있도록 tensor 변환 등 전처리하는 방법에 대해 기술한다. 

 

import torch
import torchvision
import torchvision.transforms as tr
from torch.utils.data import DataLoader, Dataset
import numpy as np
import matplotlib.pyplot as plt

# 1. GPU 장치 설정
device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)  # GPU 사용 여부 설정

라이브러리 import 및 GPU 설정

 

# 2. 데이터 전처리 변환
transf = tr.Compose(
    [tr.Resize(16), tr.ToTensor()]
)  # 이미지 크기를 16x16으로 조정 후 텐서로 변환

tr.compose는 입력된 전처리를 순서대로 적용한다. tr.Resize는 이미지의 크기를 변경하며, 여기서는 16x16으로 변환한다. 이때 만약 이미지의 너비와 높이가 다를 경우 이를 각각 지정해주어야 한다. 즉, tr.Resize((16, 16))으로 입력해주어야 한다. 

 

ToTensor()을 사용하여 이미지를 PyTorch가 처리할 수 있는 tensor로 변환한다. 그 결과 이미지의 픽셀값이 [0, 255]에서 [0, 1]로 정규화된다. 

 

# 3. CIFAR-10 데이터셋 로드
trainset = torchvision.datasets.CIFAR10(
    root="./data", train=True, download=True, transform=transf
)
testset = torchvision.datasets.CIFAR10(
    root="./data", train=False, download=True, transform=transf
)
print(trainset[0][0].size())  # 첫 번째 샘플의 이미지 크기 출력

torchvision.datasets에서 제공하는 CIFAR-10 데이터를 불러온다. CIFAR-10은 10개의 클래스를 가진 데이터 셋으로 학습 데이터 50,000개, 테스트 데이터 10,000개를 포함한다. root는 저장할 폴더이며, train=True이면 학습 데이터를, train=False이면 테스트 데이터를 불러온다. 데이터 전처리로는 앞서 정의한 trasf를 지정한다. 

 

일반적으로 데이터 셋은 이미지와 라벨이 동시에 들어있는 tuple 형태이다. 따라서 trainset[0]은 학습 데이터의 첫 번째 이미지로 이미지 한장과 라벨 숫자 하나가 저장되어 있다. 즉, trainset[0][0]은 이미지이며 trainset[0][1]은 라벨이다. 위 코드에서 출력 결과 이미지의 사이즈는 3 x 16 x 16이다.

torch.Size([3, 16, 16])

여기서 3은 채널 수, 16 x 16은 이미지의 너비와 높이이다. 일반적으로 컬러(RGB) 이미지는 (너비)x(높이)x(채널 수)로 크기가 표현되지만 PyTorch에서는 (채널 수)x(너비)x(높이)로 표현된다. 

 

# 4. DataLoader 정의
trainloader = DataLoader(trainset, batch_size=50, shuffle=True)  # 학습 데이터 배치 정의
testloader = DataLoader(
    testset, batch_size=50, shuffle=False
)  # 테스트 데이터 배치 정의
print(len(trainloader))  # trainloader에 포함된 배치 수 출력

DataLoader는 데이터를 미니 batch의 형태로 만들어준다. 따라서 배치 사이즈 및 셔플 여부 등을 결정할 수 있다. batch_size=50, shuffle=True는 무작위로 데이터를 섞어 한번에 50개의 데이터를 묶은 배치로 제공한다. 

 

학습 데이터는 과적합을 방지하기 위해 무작위로 섞고, 테스트 데이터는 순서대로 불러온다. CIFAR-10의 학습 데이터는 50,000개이고 배치 사이즈가 50이므로 trainloar의 배치 수는 1000이 된다. 

1000

 

# 5. DataLoader에서 데이터 가져오기
images, labels = next(iter(trainloader))  # 첫 번째 배치의 데이터 가져오기
images, labels = images.to(device), labels.to(device)  # GPU로 이동
print(images.size())  # 배치 크기 출력

iter(trainloader)은 DataLoader 객체인 trainloader를 반복자로 변환한다. next()는 반복자에서 다음 배치를 반환하므로 여기서 image와 labels는 첫 번째 배치의 데이터이다. 이 데이터를 GPU로 이동시킨다.

 

배치 사이즈는 (배치 크기)x(채널 수)x(너비)x(높이)로 다음 출력 결과는 배치 하나에 50개의 이미지가 들어가 있음을 의미한다. 

torch.Size([50, 3, 16, 16])

 

# 6. 이미지 시각화
oneshot = images[1].cpu().permute(1, 2, 0).numpy()  # GPU 텐서를 CPU로 이동 후 시각화
plt.figure(figsize=(2, 2))
plt.imshow(oneshot)
plt.axis("off")
plt.show()

image[1]은 선택한 배치에서 두 번째 이미지이다. Matplotlib은 GPU 텐서를 처리할 수 없으므로 .cpu()로 CPU로 이동시킨다. permute(1, 2, 0)은 기존 차원의 인덱스인 0, 1, 2를 1, 2, 0으로 바꿔주는 함수이다. image[1]의 크기는 (3, 16, 16)이고 이는 (채널 수)x(너비)x(높이)이다. 이는 일반적인 RGB 이미지 (너비)x(높이)x(채널 수)와 다르므로 Matplotlib과 호환되도록 채널을 마지막 차원으로 보낸다. 

 

왼쪽 결과는 위 코드대로 실행한 결과이다. CIFAR-10의 원본 이미지 크기는 32x32이다. 따라서 이를 아래와 같이 조정하지 않고 그대로 전처리 하면, 보다 선명한 이미지를 얻는 것을 확인할 수 있다(오른쪽)

transf = tr.Compose([tr.ToTensor()])  # 이미지 크기 조정 없이 텐서로 변환

 

전체 코드

import torch
import torchvision
import torchvision.transforms as tr
from torch.utils.data import DataLoader, Dataset
import numpy as np
import matplotlib.pyplot as plt

# 1. GPU 장치 설정
device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)  # GPU 사용 여부 설정

# 2. 데이터 전처리 변환
transf = tr.Compose(
    [tr.Resize(16), tr.ToTensor()]
)  # 이미지 크기를 16x16으로 조정 후 텐서로 변환

# 3. CIFAR-10 데이터셋 로드
trainset = torchvision.datasets.CIFAR10(
    root="./data", train=True, download=True, transform=transf
)
testset = torchvision.datasets.CIFAR10(
    root="./data", train=False, download=True, transform=transf
)
print(trainset[0][0].size())  # 첫 번째 샘플의 이미지 크기 출력

# 4. DataLoader 정의
trainloader = DataLoader(trainset, batch_size=50, shuffle=True)  # 학습 데이터 배치 정의
testloader = DataLoader(
    testset, batch_size=50, shuffle=False
)  # 테스트 데이터 배치 정의
print(len(trainloader))  # trainloader에 포함된 배치 수 출력

# 5. DataLoader에서 데이터 가져오기
images, labels = next(iter(trainloader))  # 첫 번째 배치의 데이터 가져오기
images, labels = images.to(device), labels.to(device)  # GPU로 이동
print(images.size())  # 배치 크기 출력

# 6. 이미지 시각화
oneshot = images[1].cpu().permute(1, 2, 0).numpy()  # GPU 텐서를 CPU로 이동 후 시각화
plt.figure(figsize=(2, 2))
plt.imshow(oneshot)
plt.axis("off")
plt.show()