PyTorch Dataset 클래스 상속과 매직 메서드 이해하기

PyTorch에서 데이터를 모델에 공급하기 위해서는 Dataset과 DataLoader를 사용합니다.

많은 초보자들이 Dataset을 상속할 때 반드시 작성해야 하는 __init__(), __len__(), __getitem__() 함수의 역할을 이해하지 못하는 경우가 많습니다.

이번 글에서는 간단한 학생 점수 데이터셋을 직접 만들어 보면서 Dataset의 동작 원리를 알아보겠습니다.

Dataset이란?

Dataset은 데이터를 저장하고 관리하는 클래스입니다.

예를 들어 다음과 같은 학생 점수 데이터가 있다고 가정해 보겠습니다.

우리는 공부시간을 입력(X)으로 사용하고, 시험점수를 정답(Y)으로 사용하는 데이터셋을 만들 수 있습니다.

Dataset 상속하기

from torch.utils.data import Dataset
import torch

class StudentDataset(Dataset):
    def __init__(self):
        self.study_time = [2, 3, 4, 5, 6]
        self.score = [60, 70, 80, 90, 95]

    def __len__(self):
        return len(self.study_time)

    def __getitem__(self, index):
        x = torch.tensor(self.study_time[index],
                         dtype=torch.float32)

        y = torch.tensor(self.score[index],
                         dtype=torch.float32)

        return x, y

매직 메서드(Magic Method)란?

파이썬에는 특별한 기능을 수행하는 함수들이 있습니다.

대표적으로 다음과 같은 함수들이 있습니다.

__init__()
__len__()
__getitem__()
__str__()

이처럼 양쪽에 __가 붙은 함수를 매직 메서드(Magic Method) 또는 던더 메서드(Dunder Method)라고 부릅니다.

PyTorch의 Dataset 역시 이러한 매직 메서드를 이용하여 동작합니다.

init() 이해하기

def __init__(self):
    self.study_time = [2, 3, 4, 5, 6]
    self.score = [60, 70, 80, 90, 95]

객체를 생성할 때 자동으로 실행됩니다.

dataset = StudentDataset()

실제로는 내부적으로

dataset.__init__()

이 호출됩니다.

역할

데이터를 준비하는 역할을 합니다.

현재는 리스트를 생성했지만 실제 프로젝트에서는

• CSV 파일 읽기
• JSON 파일 읽기
• 이미지 경로 저장
• 데이터베이스 연결

등의 작업을 수행할 수 있습니다.

len() 이해하기

def __len__(self):
    return len(self.study_time)

언제 호출될까?

len(dataset)

를 실행하면

실제로는

dataset.__len__()

이 호출됩니다.

확인해보기

dataset = StudentDataset()

print(len(dataset))

출력

5

총 5개의 데이터가 있다는 의미입니다.

getitem() 이해하기

Dataset에서 가장 중요한 함수입니다.

def __getitem__(self, index):

    x = torch.tensor(
        self.study_time[index],
        dtype=torch.float32
    )

    y = torch.tensor(
        self.score[index],
        dtype=torch.float32
    )

    return x, y

언제 호출될까?

dataset[0]

를 실행하면

실제로는

dataset.__getitem__(0)

이 호출됩니다.

확인해보기

dataset = StudentDataset()

x, y = dataset[0]

print(x)
print(y)

출력

tensor(2.)
tensor(60.)

첫 번째 학생의 데이터가 반환된 것입니다.

DataLoader는 어떻게 사용할까?

from torch.utils.data import DataLoader

dataset = StudentDataset()

loader = DataLoader(
    dataset,
    batch_size=2,
    shuffle=True
)

DataLoader 내부 동작

DataLoader는 직접 데이터를 저장하지 않습니다.

대신 Dataset의 __getitem__()을 반복 호출합니다.

예를 들어 batch_size가 2라면

dataset[0]
dataset[1]

를 호출하여

tensor([2., 3.])
tensor([60., 70.])

형태의 배치를 만듭니다.

배치 확인하기

for x, y in loader:
    print(x)
    print(y)
    break

출력 예시

tensor([4., 2.])
tensor([80., 60.])

shuffle=True 이므로 순서는 매번 달라질 수 있습니다.

DataLoader가 실제로 호출하는 함수

많은 사람들이 DataLoader를 사용하면서 내부 동작을 모르고 넘어갑니다.

사실 DataLoader는 다음 두 개의 함수를 계속 사용합니다.

len(dataset)

↓

dataset.__len__()

그리고

dataset[i]

↓

dataset.__getitem__(i)

즉 DataLoader가 동작하려면 반드시

__len__()
__getitem__()

이 구현되어 있어야 합니다.

정리

PyTorch에서 Dataset을 만들 때 반드시 이해해야 할 세 가지 매직 메서드는 다음과 같습니다.

한 문장으로 정리하면,

Dataset는 데이터를 어떻게 꺼낼지 정의하고, DataLoader는 Dataset를 이용해 배치를 만들어 모델에 공급합니다.

이 원리를 이해하면 이미지 분류, 자연어 처리(NLP), 시계열 데이터 등 어떤 딥러닝 프로젝트에서도 자신만의 Dataset 클래스를 쉽게 구현할 수 있습니다.