28 lines
830 B
Python
28 lines
830 B
Python
import numpy as np
|
|
from .batch import Batch
|
|
|
|
|
|
class DataLoader:
|
|
|
|
def __init__(self, dataset, batch_size, sampler=None):
|
|
self.dataset = dataset
|
|
self.batch_size = batch_size
|
|
self.sampler = sampler
|
|
|
|
def __iter__(self):
|
|
if self.sampler is not None:
|
|
indices = list(iter(self.sampler))
|
|
|
|
else:
|
|
indices = list(range(len(self.dataset)))
|
|
|
|
for i in range(0, len(indices), self.batch_size):
|
|
batch_indices = indices[i:i + self.batch_size]
|
|
batch_data = self.dataset[batch_indices]
|
|
yield Batch(batch_data, len(batch_data))
|
|
|
|
def __len__(self):
|
|
if self.sampler is not None:
|
|
return len(self.sampler) // self.batch_size
|
|
|
|
return len(self.dataset) // self.batch_size
|