[Python] Torch tensor dataloader shape issue

I have a simple application of torch.DataLoader that gets a nice performance boost. It's created by the tensor_loader in the following example.

from torch.utils.data import DataLoader, TensorDataset, BatchSampler, RandomSampler
import numpy as np
import pandas as pd

def tensor_loader(dataset: TensorDataset, batch_size: int):
return DataLoader(
dataset=dataset,
sampler=BatchSampler(
sampler=RandomSampler(dataset), # randomly sample indexes, same as shuffle=True
batch_size=batch_size,
drop_last=True
)
)

dataset = TensorDataset(torch.tensor(np.random.random(1_000_000).reshape(-1, 10), dtype=torch.float32))

start = pd.Timestamp.now()
for i in tensor_loader(dataset, 4096):
i[0]
end = pd.Timestamp.now()
print(end - start)
assert i[0].shape == torch.Size([1, 4096, 10])

start = pd.Timestamp.now()
simple_loader = DataLoader(dataset=dataset, batch_size=4096)
for i in simple_loader:
pass
end = pd.Timestamp.now()
print(end - start)
assert next(iter(simple_loader))[0].shape == torch.Size([4096, 10])


However, the difference in the shapes is a little annoying: the data_loader adds an exterior dimension that I don't want. It means the the two loaders can't be substituted for one another, which would mean a lot of niggling changes to existing code to substitute the tensor_loader for the existing one.

Obviously, I can subclass the DataLoader to drop the outer dimension, but this feels more complicated than it should be. Is there are way to create something like the above data_loader's return value that will produce the shape of the simpler_loader when iterating?

Continue reading...