Tensor Loader
TensorLoader is similar to the combination of PyTorch's TensorDataset and DataLoader. It is faster and has better type hints.
Installation
Install from PyPI:
pip install tensorloaderInstall from source:
git clone https://github.com/zhb2000/tensorloader.git
cd tensorloader
pip install .Usage
This package only contains a TensorLoader class.
from tensorloader import TensorLoaderUse a single tensor as data:
X = torch.tensor(...)
dataloader = TensorLoader(X)
for x in dataloader:
...Use a tuple of tensor as data:
X = torch.tensor(...)
Y = torch.tensor(...)
dataloader = TensorLoader((X, Y))
for x, y in dataloader: # unpack the batch tuple as x, y
...Use a namedtuple of tensor as data:
from collections import namedtuple
Batch = namedtuple('Batch', ['x', 'y'])
X = torch.tensor(...)
Y = torch.tensor(...)
# set unpack_args=True when using a namedtuple as data
dataloader = TensorLoader(Batch(X, Y), unpack_args=True)
for batch in dataloader:
assert isinstance(batch, Batch)
assert isinstance(batch.x, torch.Tensor)
assert isinstance(batch.y, torch.Tensor)
x, y = batch
...PS: Namedtuples are similar to common tuples and they allow field access by name which makes code more readable. For more information, see the documentation of namedtuple.
Speed Test
TensorLoader is much faster than TensorDataset + DataLoader, for it uses vectorized operations instead of creating costly Python lists.
import timeit
import torch
from torch.utils.data import TensorDataset, DataLoader
from tensorloader import TensorLoader
def speed_test(epoch_num: int, **kwargs):
sample_num = int(1e6)
X = torch.zeros(sample_num, 10)
Y = torch.zeros(sample_num)
tensorloader = TensorLoader((X, Y), **kwargs)
torchloader = DataLoader(TensorDataset(X, Y), **kwargs)
def loop(loader):
for _ in loader:
pass
t1 = timeit.timeit(lambda: loop(tensorloader), number=epoch_num)
t2 = timeit.timeit(lambda: loop(torchloader), number=epoch_num)
print(f'TensorLoader: {t1:.4g}s, TensorDatset + DataLoader: {t2:.4g}s.')>>> speed_test(epoch_num=10, batch_size=128, shuffle=False)
TensorLoader: 0.363s, TensorDatset + DataLoader: 54.39s.
>>> speed_test(epoch_num=10, batch_size=128, shuffle=True)
TensorLoader: 0.9296s, TensorDatset + DataLoader: 56.54s.
>>> speed_test(epoch_num=10, batch_size=10000, shuffle=False)
TensorLoader: 0.005262s, TensorDatset + DataLoader: 55.57s.
>>> speed_test(epoch_num=10, batch_size=10000, shuffle=True)
TensorLoader: 0.5682s, TensorDatset + DataLoader: 57.71s.
