[python] 파이썬 PyTorch에서 데이터셋을 생성하는 방법은?

22 Nov 2023

python

PyTorch는 딥러닝 모델을 구축할 때 많이 사용되는 파이썬 기반의 오픈 소스 라이브러리입니다. 데이터셋을 생성하기 위해 PyTorch의 torch.utils.data.Dataset 클래스를 사용할 수 있습니다. 데이터셋을 생성하는 방법은 다음과 같습니다:

torch.utils.data.Dataset 클래스를 상속하여 새로운 클래스를 만듭니다.

import torch
from torch.utils.data import Dataset

class CustomDataset(Dataset):
    def __init__(self, data):
        self.data = data

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

    def __getitem__(self, idx):
        # idx에 해당하는 데이터를 불러오고 반환합니다.
        return self.data[idx]

__init__ 메서드에서 데이터를 초기화하고, __len__ 메서드에서 데이터셋의 총 개수를 반환하며, __getitem__ 메서드에서 인덱스에 해당하는 데이터를 반환합니다.
데이터를 생성하여 CustomDataset 클래스에 전달합니다.
```
data = [1, 2, 3, 4, 5]
dataset = CustomDataset(data)
```
생성한 데이터셋을 PyTorch의 DataLoader에 넣어 배치 크기와 데이터 로딩 방식 등을 설정할 수 있습니다.
```
dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=True)
```

생성한 데이터셋을 이용하여 학습 과정에서 배치 단위로 데이터를 로딩하고 모델에 전달할 수 있습니다.

for batch in dataloader:
    inputs, labels = batch
    # 모델에 데이터를 전달하고 학습 또는 추론을 수행합니다.

PyTorch를 사용하여 데이터셋을 생성하고 모델을 훈련시키는 것은 딥러닝 모델 개발 과정에서 매우 중요합니다. 데이터셋을 효율적으로 관리하고 이용하는 것은 모델의 성능 향상에 큰 영향을 미칠 수 있습니다.

참고 자료: