Claude 를 참고하여 나만의 custom dataset 만들기

Custom dataset 을 정의 하는 이유

Pytorch 로 deep learning 을 실시할 때 일반적으로 데이터를 그냥 하나를 한번에 가져오는게 아니라 Custom dataset class를 통해 가져온다.

 

여기서 Custom data class를 사용하는 이유는 다음과 같다.

1)    Data 를 한 번에 load 하는 것이 컴퓨터 메모리에 부담이 되기 때문이다. 따라서 큰 Data 를 한번에 다 메모리에 올리는 것이 아니라 이를 ‘mini batch’라는 형태로 토막을 낸다. 그 한 토막을 ‘mini batch’라고 한다. 이 작업을 할 때 분석가가 전체 Data 큰 덩어리 1개를 받아와서 전체 Data size (data 를 구성하는 sample 의 수) 는 얼마인지 파악하고, Mini batch size 는 얼마로 할것인지 정하고 (보통 32정도를 사용한다. 즉, 1개의 mini batch 안에 32개의 sample 이 들어가게 되는 것), mini batch가 그렇다면 몇 개가 필요한지 등을 계산해서 for loop 을 돌려서 index 를 iteration 하면서 읽어와야 한다. 이를 체계화, 자동화 한 것이 custom dataset class 이며 pytorch 의 ‘Dataset’ 을 상속받아 작성할 수 있다.

 

2)    그렇다면 메모리 문제만 없다면 큰 Data 1개를 한번에 memory 에 올려도 될까? (즉, batch gradient descent 를 실시해도 될까?) 이는 그다지 추천할 방법은 아니다. 1개만 train 하는 것이므로 training 에 bias 가 일어나게 되어 overfitting 이 일어날 가능성이 높아진다. Mini batch 를 이용하면 mini batch 들 사이에 어느정도 variation 이 생길 수 밖에 없다. 모든 mini batch 에 들어가는 샘플들이 똑같똑 않을 테니까 말이다. 그렇게 되면 좀 더 generalized 된 training 을 할 수 있다. 또한 mini 사이즈로 잘게 쪼개다보니 weight 가 빠르게 업데이트 될 수 있다는 장점이 있다.

 

내가 해야 할 분석에 맞는 Custom dataset code 를 작성하고 싶은데 Claude 는 어떻게 code 를 작성해줄 지 궁금해서 아래와 같은 시도를 해 보았다. 

 

Custom dataset code 생성을 위한 Prompt 작성

기존 신약개발 부트캠프때 연습했던 자료를 정리하여 ‘ultimate_flow.ipynb’ 라는 이름으로 저장하고 현재 신약개발 경진대회에서 사용될 train dataset 의 일부를 sampling 하고 값을 변형하여 csv file 저장했다. 이 2개의 file을 올리면서 ‘ultimate_flow.ipynb’ 에 나온 code 스타일로 training 용 csv file data 를 custom dataset 으로 읽어오는 code 를 만들라고 했다.

 

Prompt 내용은 다음과 같다:

Prompt 실행 결과 

Prompt 실행 했더니 다음과 같은 코드가 생성 되었다.

import torch
from torch.utils.data import Dataset
import pandas as pd
from rdkit import Chem
from rdkit.Chem import AllChem

class QSARDataset(Dataset):
    def __init__(self, csv_file, smiles_col='Smiles', target_col='pIC50'):
        self.data = pd.read_csv(csv_file)
        self.smiles_col = smiles_col
        self.target_col = target_col

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

    def __getitem__(self, idx):
        smiles = self.data.iloc[idx][self.smiles_col]
        target = self.data.iloc[idx][self.target_col]

        # Convert SMILES to molecular fingerprint
        mol = Chem.MolFromSmiles(smiles)
        fingerprint = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=2048)
        
        # Convert fingerprint to PyTorch tensor
        fingerprint_tensor = torch.tensor(list(fingerprint.ToBitString()), dtype=torch.float32)
        
        # Convert target to PyTorch tensor
        target_tensor = torch.tensor(target, dtype=torch.float32)

        sample = {
            'fingerprint': fingerprint_tensor,
            'target': target_tensor,
        }

        return sample
        
# Create dataset
train_dataset = QSARDataset('train_sample.csv')

# Create DataLoader
from torch.utils.data import DataLoader

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# Example of iterating through the data
for batch in train_loader:
    X_batch = batch['fingerprint']
    y_batch = batch['target']
    # Your model and training code here

 

QSARDataset 이라는 이름의 custom dataset class 을 보면 def __init__, __len__, __getitem__ 부분으로 나뉘어져 있는데 이는 custom dataset의 기본 구조로, 변하지 않는 부분이다. 단, 그 구조의 하위에 있는 내용은 바뀔 수 있다. 보통은 Custom dataset 을 정의 할 때 __init__ 부분에는 dataset, feature, label 을 지정하여 따로 변수에 저장한다. 여기서는 feature 에 해당하는 부분이 smiles_col, label (target) 에 해당하는 부분이 ‘pIC50’ 으로 지정되어 있다. __getitem__ 부분을 보면 여기서 SMILES 를 finger print 로 바꾸는 작업을 하고 feature 와 target 을 각각 pytorch tensor 로 바꿔주고 이를 묶어 sample 로 return 하고 있다.