PyTorch The Dataset Class
The Dataset class is a base class for serving rows one at a time. A subclass writes __len__ and __getitem__, the two methods PyTorch calls.
Rows served by hand
Plain tensors can serve a row with an index. A class can take over this job and hide the tensors.
print(features[2])
print(targets[2])The two methods
The class stores the two tensors at build time. PyTorch calls __len__ for the count and __getitem__ for a row.
class Students(Dataset):
def __init__(self, features, targets):
self.features = features
self.targets = targets
def __len__(self):
return len(self.features)
def __getitem__(self, index):
return self.features[index], self.targets[index]Asking the class for a row
Square brackets call __getitem__ under the hood. One call returns one input and its target.
data = Students(features, targets)
x_row, y_row = data[2]
print(x_row)
print(y_row)Example
import torch
from torch.utils.data import Dataset
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])
class Students(Dataset):
def __init__(self, features, targets):
self.features = features
self.targets = targets
def __len__(self):
return len(self.features)
def __getitem__(self, index):
return self.features[index], self.targets[index]
data = Students(features, targets)
print("Rows:", len(data))
print(data[2])Key points
- Dataset: a base class for objects that serve rows one at a time
- __len__: the method that returns how many rows the dataset holds
- __getitem__: the method that returns the row for one index
- Use data[index] to fetch one row from a Dataset
- A Dataset subclass serves one row each time PyTorch asks for an index
In this example, the Students class stores the features and the targets. Its __len__ method returns 5, the number of rows. Its __getitem__ method returns one pair for one index. The last line fetches row 2 and prints the pair. Row 2 is the input 4 hours with its target score 71.
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
Which two methods must a Dataset subclass write?