🎓 A1Learner
🔍

PyTorch TensorDataset

TensorDataset is a ready-made Dataset that wraps aligned tensors. It serves one row at a time, with no custom class to write.

Wrap two tensors

One line wraps the pair. The dataset reports the row count of the tensors.

data = TensorDataset(features, targets)
print(len(data))

Fetch one row

One index returns one pair. Row 2 is 4 hours with a score of 71.

x_row, y_row = data[2]
print(x_row)
print(y_row)

Fetch a slice

A slice returns a small batch. Both parts keep the same number of rows.

xb, yb = data[0:2]
print(xb.shape)
print(yb.shape)

Example

import torch
from torch.utils.data import TensorDataset

features = torch.tensor([[2.0], [3.0], [4.0], [5.0], [6.0]])
targets = torch.tensor([55.0, 62.0, 71.0, 79.0, 88.0])
data = TensorDataset(features, targets)

print("Rows:", len(data))
for start in range(0, 5, 2):
    xb, yb = data[start:start + 2]
    print(f"Batch from row {start}: inputs {xb.shape}, targets {yb.shape}")

Key points

  • TensorDataset: a ready-made Dataset that wraps aligned tensors
  • Use TensorDataset(features, targets) to wrap a pair of tensors
  • The wrapped tensors must share the same first size
  • Use data[index] to fetch one row, and data[start:stop] to fetch a slice
  • TensorDataset serves the same aligned rows a custom class would serve

In this example, TensorDataset wraps the same student rows as a custom class would. The loop takes slices of 2 rows at a time. Each slice returns an inputs part and a targets part. The shapes show batches of 2, 2, and then 1 row. The last batch is short because 5 rows do not split evenly by 2.

Note

TensorDataset only wraps tensors, so data in other forms needs a custom Dataset class.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

When is TensorDataset the right choice?