[Lucid] Dataset과 DataLoader 설계

📦 Dataset과 DataLoader 설계

Lucid에서 데이터 파이프라인은 사용자가 가장 자주 만나는 영역이다. PyTorch의 Dataset/DataLoader를 거의 습관처럼 쓰던 입장에서, Lucid도 익숙한 사용감을 제공하는 것이 중요했다. 동시에 MLX와 NumPy를 혼용하는 환경에서 가볍고 예측 가능한 동작을 유지해야 했다. 이 문서는 lucid/data/_base.py, _util.py를 구현하며 어떤 선택을 했고, 어디서 시간을 쏟았는지를 기술한다.

---

🧭 설계 원칙

• 호환성: __getitem__, __len__, DataLoader 인터페이스를 PyTorch와 최대한 유사하게 유지.
• 단순화: 멀티프로세싱/worker, pin_memory 등은 제외. 단일 프로세스에 집중.
• 디바이스 흐름: TensorDataset.to(device)로 모델 .to()와 자연스럽게 연동.
• 명확한 에러: 길이 불일치, 잘못된 인덱스, 디바이스 깨짐 등에 대해 즉시 명확한 예외를 던진다.

처음에는 PyTorch 옵션을 모두 따라갈 생각이었지만, MLX의 lazy 특성과 Mac 환경을 고려해 "필요한 최소"를 선택했다. 나머지는 API 모양을 맞춰 두고 후속 단계에서 확장할 수 있도록 열어두었다.

🧱 Dataset 베이스 – 기본 틀 마련

경로: lucid/data/_base.py

class Dataset(ABC):
    @abstractmethod
    def __getitem__(self, idx): ...

    @abstractmethod
    def __len__(self): ...

    def __add__(self, other): 
        return ConcatDataset([self, other])

    def __iter__(self):
        for i in range(len(self)): yield self[i]

PyTorch처럼 베이스는 아무 것도 안 한다. transform 훅이나 캐싱을 넣을까 고민했지만, 강제 기능이 늘어날수록 사용자 자유도가 줄었다. __add__로 데이터셋 결합을 자연스럽게 지원하고, __iter__를 제공해 단순 순회도 가능하게 했다.

예시

class MyDataset(Dataset):
    def __init__(self, X, y): 
        self.X, self.y = X, y

    def __getitem__(self, idx): 
        return self.X[idx], self.y[idx]

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

베이스를 가볍게 유지한 덕분에 사용자는 transform, caching, on-the-fly augmentation 등을 상속으로 자유롭게 덧붙일 수 있다. 초기에는 로깅 훅을 시도했지만, 사용자 정의 __getitem__과 충돌해 제거했다.

🧩 Subset – 속성 위임과 인덱스 래핑

경로: lucid/data/_base.py

• 아이디어: dataset[indices[idx]]를 래핑.
• __iter__를 indices 기반으로 재정의.
• __getattr__로 원본 속성 위임(PyTorch 패턴).

초기에는 속성 위임 중 일부가 덮여 재귀 호출이 발생하는 문제가 있었다. 위임 로직을 단순히 getattr(self.dataset, name)으로 제한해 해결했다. 음수 인덱스는 원본 길이를 더해 처리해 PyTorch 동작과 맞췄다.

예시

train_ds = Subset(full_ds, list(range(0, 800)))
val_ds   = Subset(full_ds, list(range(800, 1000)))

🧊 TensorDataset – 길이 검증과 디바이스 보존

경로: lucid/data/_base.py

• 최소 1개 텐서, 모두 같은 길이, 1D 이상 강제.
• 입력이 Tensor/array 혼합이어도 _check_is_tensor로 승격해 디바이스/백엔드 정보를 유지.
• to(device)로 내부 텐서 일괄 이동.

길이 검증을 소홀히 했다가 서로 다른 길이 텐서를 묶었을 때 인덱싱이 뒤틀리는 문제를 겪었다. PyTorch처럼 "길이가 다르면 예외"를 명시적으로 넣었다. MLX 텐서와 NumPy 텐서가 섞일 때 장치 정보가 사라지는 문제는 모든 입력을 Tensor로 승격하고, .to(device)에서 새 튜플로 재할당하여 해결했다.

예시

X = lucid.randn(100, 32); y = lucid.randint(0, 10, (100,))
ds = TensorDataset(X, y).to("gpu")

🧱 ConcatDataset – 음수 인덱스와 누적 크기

경로: lucid/data/_base.py

여러 Dataset을 이어붙이고 cumulative_sizes로 위치를 빠르게 찾는다. 음수 인덱스를 처리하다가 off-by-one 버그가 있었는데, PyTorch와 동일하게 길이에 더하는 방식으로 통일했다. 누적 크기를 매 접근마다 계산해 $\mathcal{O}(n)$ 비용이 발생하던 초기 버전은, 생성자에서 한 번만 계산하도록 변경해 해결했다.

예시

full = ConcatDataset([train_ds, extra_ds])
last = full[-1]

🚚 DataLoader – 데이터 불러오기의 핵심

경로: lucid/data/_base.py

• 인덱스 목록을 만들어 배치 단위 슬라이스, shuffle=True면 epoch마다 재셔플.
• default_collate: 튜플/리스트 batch는 전치 후 lucid.stack, Tensor는 바로 stack, 기타는 원본 유지.

가장 많은 수정이 collate였다. 처음 np.stack을 썼더니 MLX 텐서가 섞이면 디바이스가 깨졌다. Lucid의 stack으로 통일해 디바이스/백엔드 일관성을 확보했다. 튜플 전치 과정에서 축 순서를 잘못 잡아 shape mismatch가 났으나, zip(*)로 전치 후 각 요소를 stack하는 방식으로 정리했다. 마지막 배치가 덜 찼을 때는 min(start+batch_size, len(indices))로 안전하게 슬라이스했다.

사용 예

loader = DataLoader(ds, batch_size=32, shuffle=True)
for x, y in loader:
    ...

커스텀 collate

def my_collate(batch):
    xs, ys = zip(*batch)
    return lucid.stack(xs, axis=0), lucid.stack(ys, axis=0)
loader = DataLoader(ds, batch_size=16, collate_fn=my_collate)

🔀 random_split – 소수 길이와 시드 관리

경로: lucid/data/_util.py

비율 분할 시 합이 $1$이 아니면 예외를 던지고, 반올림 오차로 남는 샘플은 앞에서부터 분배한다. 부동소수점 합이 $0.999999\ldots$ 로 떨어지는 사례가 있어 math.isclose를 적용했다. 인덱스 셔플은 random.Random(seed)로 지역 RNG를 만들어 전역 상태를 오염시키지 않는다.

예시

train_ds, val_ds = random_split(ds, [0.8, 0.2], seed=42)

🧭 PyTorch와의 호환성/차이 – 의도적 간소화

• 비슷한 점: 인터페이스, Subset/ConcatDataset, random_split, collate_fn, shuffle.
• 다른 점: 멀티프로세싱/worker 미지원, pin_memory 없음, TensorDataset.to(device) 같은 편의 추가, default_collate가 Lucid Tensor 스택을 사용.

멀티프로세싱을 뺀 이유는 MLX lazy 특성과 Mac의 포크 제약을 감안했을 때 단순하고 예측 가능한 쪽이 낫다고 판단했기 때문이다. pin_memory도 Apple 실리콘 통합 메모리 환경에서 이점이 제한적이라 제외했다. 대신 API 모양을 PyTorch와 맞춰두어, 추후 필요 시 옵션을 추가해도 사용자 경험을 유지할 수 있도록 했다.

🔎 추가 구현 디테일 – 에러 메시지와 로깅

• 길이 불일치, $\mathbb{R}^0$ 텐서 입력, 빈 데이터셋 등은 즉시 명확한 예외 메시지를 던지도록 했다.
• DataLoader에서 인덱스 초과 시 StopIteration으로 종료, 내부 상태를 초기화해 재사용 가능하게 했다.
• __repr__를 간단히 제공해 디버깅 시 데이터셋 크기, 텐서 shape, 디바이스 정보를 바로 확인할 수 있다(TensorDataset, Subset 등).

이런 세부 로깅은 개발 중 혼선을 줄였고, 사용자가 데이터 구성을 눈으로 검증하기 쉽게 해주었다.

---

🧵 실전 예제: 이미지 분류 미니 루프

import lucid
import lucid.nn as nn
from lucid.data import TensorDataset, DataLoader, random_split

X = lucid.randn(1000, 3, 32, 32)
y = lucid.randint(0, 10, (1000,))

dataset = TensorDataset(X, y)
train_ds, val_ds = random_split(dataset, [0.8, 0.2], seed=123)

train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=128, shuffle=False)

model = nn.Sequential(
    nn.Linear(3 * 32 * 32, 128), 
    nn.ReLU(), 
    nn.Linear(128, 10),
).to("gpu")

opt = lucid.optim.SGD(model.parameters(), defaults={"lr": 1e-2})

for epoch in range(5):
    model.train()
    for x, y in train_loader:
        x = x.to("gpu").reshape(x.shape[0], -1)
        y = y.to("gpu")
        opt.zero_grad()

        loss = nn.functional.cross_entropy(model(x), y)
        loss.eval()
        loss.backward()
        opt.step()

    model.eval()
    correct = total = 0
    for x, y in val_loader:
        x = x.to("gpu").reshape(x.shape[0], -1)
        y = y.to("gpu")

        preds = model(x).argmax(axis=1)
        correct += (preds == y).sum().item()
        total += y.size

    print(f"Epoch {epoch} acc={correct/total:.3f}")

PyTorch 사용감과 거의 동일한 흐름을 Lucid에서 재현한다. MLX의 lazy 특성만 loss.eval()로 챙기면 된다.

---

✅ 정리

lucid.data는 "익숙한데 가벼운" 데이터 파이프라인을 목표로, PyTorch 핵심 인터페이스를 최소 구현으로 재현했다. 길이 검증, 음수 인덱스, 디바이스 깨짐 같은 문제를 밟고 수정하면서 MLX/NumPy 혼용 환경에서도 일관되게 동작하는 뼈대를 만들었다. 단일 프로세스와 기본 collate만 지원하지만, API 모양과 규약을 PyTorch와 맞춰 두었기에 이후 기능을 추가하더라도 사용자 경험을 해치지 않는 범위에서 확장할 수 있을 것이다.