Custom Dataset Class

mediumFREEDatasetsPyTorch DatasetsPytorch

Problem Statement

Create a custom PyTorch Dataset that wraps in-memory features and labels.

Background

PyTorch's Dataset interface is how training loops and DataLoader read samples. A map-style dataset implements three methods:

init: store (or prepare) the data
len: return the number of samples
getitem(idx): return the sample at position idx

Your Task

Write a class NumberDataset(features, labels) that subclasses torch.utils.data.Dataset:

Store features as a float32 tensor and labels as an int64 (long) tensor.
len(dataset) returns the number of samples.
dataset[i] returns the tuple (feature_tensor, label_tensor) for sample i.

Example

Input

features=[[1.0, 2.0], [3.0, 4.0]], labels=[0, 1]

Output

Dataset with 2 samples, dataset[0] returns (tensor([1., 2.]), tensor(0))

Reasoning

Inputs: features is a list of $N = 2$ samples, each a vector of 2 real numbers: $x_0 = (1.0, 2.0)$ , $x_1 = (3.0, 4.0)$ . labels holds one integer class per sample: $y_0 = 0$ , $y_1 = 1$ . Row $i$ of the features pairs with entry $i$ of the labels.
Length: the dataset size is the number of samples $N$ (rows of features), so it is 2, not 4 (the total count of numbers).
Indexing: sample $0$ is the pair $(x_0, y_0)$ : a feature tensor holding [1., 2.] and a label tensor holding 0. Sample $1$ would give $(3., 4.)$ with label $1$ .
Output types: the feature is a 1-D tensor of length 2 with 32-bit floats, which is why it prints 1. and 2.. The label is a 0-dimensional (scalar) tensor with a 64-bit integer, so it prints with no decimal point; classification losses expect integer class indices.
Notice that indexing drops the sample dimension: the stored features form a $2 \times 2$ grid, but one sample is just a length-2 vector.

Constraints

Features and labels lists will have the same length
Features can be 1D or 2D lists
Labels are integers
Return tensors (not lists) from getitem

Learn this in

Python FoundationsFoundations PyTorch FundamentalsFoundations

More from Pytorch

Apply Softmax for Predictionseasy Count Model Parametersmedium Build with nn.Sequentialeasy Use nn.Flatten Layereasy Create a Simple Neural Network Classeasy

Concept primer

1. Background Knowledge

PyTorch's torch.utils.data.Dataset is the interface every training pipeline reads samples through. A map-style dataset implements three methods:

init(self, ...): store or prepare the data once.
len(self): return the number of samples. DataLoader uses it to build index ranges.
getitem(self, idx): return sample idx, usually as a (features, label) tuple.

Wrapping data in a Dataset lets DataLoader handle batching, shuffling and parallel loading for you.

2. Algorithm Approach

NumberDataset(features, labels) holds in-memory data, so the work happens once, up front:

Convert features to a float32 tensor (shape (N, d)): models expect floating-point inputs.
Convert labels to an int64 (long) tensor (shape (N,)): nn.CrossEntropyLoss requires integer class labels.
len returns N, and getitem(idx) returns row idx of each tensor.

3. Step-by-Step Strategy

import torch
from torch.utils.data import Dataset

class NumberDataset(Dataset):
    def __init__(self, features, labels):
        # 1. convert once, with the right dtypes
        ...
    def __len__(self):
        # 2. number of samples
        ...
    def __getitem__(self, idx):
        # 3. one (feature, label) pair
        ...

Example: NumberDataset([[1.0, 2.0], [3.0, 4.0]], [0, 1]) has length 2, and dataset[0] is (tensor([1., 2.]), tensor(0)).

4. Common Pitfalls

**Converting inside **getitem****: calling torch.tensor per sample repeats work on every access. Convert once in init.
Wrong label dtype: float labels make CrossEntropyLoss fail. Use dtype=torch.long.
len returning the feature width: len(self.features) is the number of rows (samples), not self.features.shape[1].
Returning a list instead of a tuple: x, y = dataset[i] and the default DataLoader collation both expect a (feature, label) pair.

5. Time & Space Complexity

Construction: $O(N \cdot d)$ to copy the data into tensors.
len and getitem: $O(1)$ . Indexing a tensor row returns a view without copying.
Space: $O(N \cdot d)$ for the stored tensors.