• Docs >
  • Using tensorclasses for datasets
Shortcuts

Using tensorclasses for datasets

In this tutorial we demonstrate how tensorclasses can be used to efficiently and transparently load and manage data inside a training pipeline. The tutorial is based heavily on the PyTorch Quickstart Tutorial, but modified to demonstrate use of tensorclass. See the related tutorial using TensorDict.

import torch
import torch.nn as nn

from tensordict import MemoryMappedTensor
from tensordict.prototype import tensorclass
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
Using device: cpu

The torchvision.datasets module contains a number of convenient pre-prepared datasets. In this tutorial we’ll use the relatively simple FashionMNIST dataset. Each image is an item of clothing, the objective is to classify the type of clothing in the image (e.g. “Bag”, “Sneaker” etc.).

training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)
Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz
Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz to data/FashionMNIST/raw/train-images-idx3-ubyte.gz

  0%|          | 0/26421880 [00:00<?, ?it/s]
  0%|          | 65536/26421880 [00:00<01:13, 359306.72it/s]
  1%|          | 229376/26421880 [00:00<00:38, 681592.87it/s]
  4%|▎         | 950272/26421880 [00:00<00:11, 2188118.27it/s]
 15%|█▍        | 3833856/26421880 [00:00<00:02, 7610021.59it/s]
 38%|███▊      | 9961472/26421880 [00:00<00:00, 17052972.27it/s]
 61%|██████    | 16154624/26421880 [00:01<00:00, 22864733.55it/s]
 81%|████████▏ | 21495808/26421880 [00:01<00:00, 29310534.27it/s]
 96%|█████████▌| 25395200/26421880 [00:01<00:00, 27082701.46it/s]
100%|██████████| 26421880/26421880 [00:01<00:00, 19331657.95it/s]
Extracting data/FashionMNIST/raw/train-images-idx3-ubyte.gz to data/FashionMNIST/raw

Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz
Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz to data/FashionMNIST/raw/train-labels-idx1-ubyte.gz

  0%|          | 0/29515 [00:00<?, ?it/s]
100%|██████████| 29515/29515 [00:00<00:00, 325233.38it/s]
Extracting data/FashionMNIST/raw/train-labels-idx1-ubyte.gz to data/FashionMNIST/raw

Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz
Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz to data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz

  0%|          | 0/4422102 [00:00<?, ?it/s]
  1%|▏         | 65536/4422102 [00:00<00:12, 356748.93it/s]
  5%|▌         | 229376/4422102 [00:00<00:06, 673443.06it/s]
 21%|██        | 917504/4422102 [00:00<00:01, 2072200.28it/s]
 73%|███████▎  | 3211264/4422102 [00:00<00:00, 6192624.39it/s]
100%|██████████| 4422102/4422102 [00:00<00:00, 5994995.27it/s]
Extracting data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz to data/FashionMNIST/raw

Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz
Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz to data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz

  0%|          | 0/5148 [00:00<?, ?it/s]
100%|██████████| 5148/5148 [00:00<00:00, 69428543.38it/s]
Extracting data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz to data/FashionMNIST/raw

Tensorclasses are dataclasses that expose dedicated tensor methods over its contents much like TensorDict. They are a good choice when the structure of the data you want to store is fixed and predictable.

As well as specifying the contents, we can also encapsulate related logic as custom methods when defining the class. In this case we’ll write a from_dataset classmethod that takes a dataset as input and creates a tensorclass containing the data from the dataset. We create memory-mapped tensors to hold the data. This will allow us to efficiently load batches of transformed data from disk rather than repeatedly load and transform individual images.

@tensorclass
class FashionMNISTData:
    images: torch.Tensor
    targets: torch.Tensor

    @classmethod
    def from_dataset(cls, dataset, device=None):
        data = cls(
            images=MemoryMappedTensor.empty(
                (len(dataset), *dataset[0][0].squeeze().shape), dtype=torch.float32
            ),
            targets=MemoryMappedTensor.empty((len(dataset),), dtype=torch.int64),
            batch_size=[len(dataset)],
            device=device,
        )
        for i, (image, target) in enumerate(dataset):
            data[i] = cls(images=image, targets=torch.tensor(target), batch_size=[])
        return data

We will create two tensorclasses, one each for the training and test data. Note that we incur some overhead here as we are looping over the entire dataset, transforming and saving to disk.

training_data_tc = FashionMNISTData.from_dataset(training_data, device=device)
test_data_tc = FashionMNISTData.from_dataset(test_data, device=device)

DataLoaders

We’ll create DataLoaders from the torchvision-provided Datasets, as well as from our memory-mapped TensorDicts.

Since TensorDict implements __len__ and __getitem__ (and also __getitems__) we can use it like a map-style Dataset and create a DataLoader directly from it. Note that because TensorDict can already handle batched indices, there is no need for collation, so we pass the identity function as collate_fn.

batch_size = 64

train_dataloader = DataLoader(training_data, batch_size=batch_size)  # noqa: TOR401
test_dataloader = DataLoader(test_data, batch_size=batch_size)  # noqa: TOR401

train_dataloader_tc = DataLoader(  # noqa: TOR401
    training_data_tc, batch_size=batch_size, collate_fn=lambda x: x
)
test_dataloader_tc = DataLoader(  # noqa: TOR401
    test_data_tc, batch_size=batch_size, collate_fn=lambda x: x
)

Model

We use the same model from the Quickstart Tutorial.

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28 * 28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits


model = Net().to(device)
model_tc = Net().to(device)
model, model_tc
(Net(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
), Net(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
))

Optimizing the parameters

We’ll optimise the parameters of the model using stochastic gradient descent and cross-entropy loss.

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
optimizer_tc = torch.optim.SGD(model_tc.parameters(), lr=1e-3)


def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()

    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)

        pred = model(X)
        loss = loss_fn(pred, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")

The training loop for our tensorclass-based DataLoader is very similar, we just adjust how we unpack the data to the more explicit attribute-based retrieval offered by the tensorclass. The .contiguous() method loads the data stored in the memmap tensor.

def train_tc(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()

    for batch, data in enumerate(dataloader):
        X, y = data.images.contiguous(), data.targets.contiguous()

        pred = model(X)
        loss = loss_fn(pred, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")


def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)

            pred = model(X)

            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= num_batches
    correct /= size

    print(
        f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n"
    )


def test_tc(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for batch in dataloader:
            X, y = batch.images.contiguous(), batch.targets.contiguous()

            pred = model(X)

            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= num_batches
    correct /= size

    print(
        f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n"
    )


for d in train_dataloader_tc:
    print(d)
    break

import time

t0 = time.time()
epochs = 5
for t in range(epochs):
    print(f"Epoch {t + 1}\n-------------------------")
    train_tc(train_dataloader_tc, model_tc, loss_fn, optimizer_tc)
    test_tc(test_dataloader_tc, model_tc, loss_fn)
print(f"Tensorclass training done! time: {time.time() - t0: 4.4f} s")

t0 = time.time()
epochs = 5
for t in range(epochs):
    print(f"Epoch {t + 1}\n-------------------------")
    train(train_dataloader, model, loss_fn, optimizer)
    test(test_dataloader, model, loss_fn)
print(f"Training done! time: {time.time() - t0: 4.4f} s")
FashionMNISTData(
    images=Tensor(shape=torch.Size([64, 28, 28]), device=cpu, dtype=torch.float32, is_shared=False),
    targets=Tensor(shape=torch.Size([64]), device=cpu, dtype=torch.int64, is_shared=False),
    batch_size=torch.Size([64]),
    device=cpu,
    is_shared=False)
Epoch 1
-------------------------
loss: 2.295250 [    0/60000]
loss: 2.287709 [ 6400/60000]
loss: 2.264283 [12800/60000]
loss: 2.261056 [19200/60000]
loss: 2.249043 [25600/60000]
loss: 2.214912 [32000/60000]
loss: 2.223378 [38400/60000]
loss: 2.193253 [44800/60000]
loss: 2.188197 [51200/60000]
loss: 2.152597 [57600/60000]
Test Error:
 Accuracy: 51.1%, Avg loss: 2.148445

Epoch 2
-------------------------
loss: 2.150964 [    0/60000]
loss: 2.141803 [ 6400/60000]
loss: 2.083223 [12800/60000]
loss: 2.101235 [19200/60000]
loss: 2.054458 [25600/60000]
loss: 1.990291 [32000/60000]
loss: 2.011465 [38400/60000]
loss: 1.941634 [44800/60000]
loss: 1.947067 [51200/60000]
loss: 1.859987 [57600/60000]
Test Error:
 Accuracy: 61.3%, Avg loss: 1.865036

Epoch 3
-------------------------
loss: 1.889201 [    0/60000]
loss: 1.855546 [ 6400/60000]
loss: 1.746735 [12800/60000]
loss: 1.788072 [19200/60000]
loss: 1.681532 [25600/60000]
loss: 1.632135 [32000/60000]
loss: 1.645971 [38400/60000]
loss: 1.565696 [44800/60000]
loss: 1.592574 [51200/60000]
loss: 1.468814 [57600/60000]
Test Error:
 Accuracy: 64.0%, Avg loss: 1.498094

Epoch 4
-------------------------
loss: 1.554366 [    0/60000]
loss: 1.520199 [ 6400/60000]
loss: 1.381839 [12800/60000]
loss: 1.457058 [19200/60000]
loss: 1.337932 [25600/60000]
loss: 1.330665 [32000/60000]
loss: 1.346015 [38400/60000]
loss: 1.285760 [44800/60000]
loss: 1.320316 [51200/60000]
loss: 1.209877 [57600/60000]
Test Error:
 Accuracy: 64.8%, Avg loss: 1.238957

Epoch 5
-------------------------
loss: 1.303831 [    0/60000]
loss: 1.287427 [ 6400/60000]
loss: 1.129778 [12800/60000]
loss: 1.242445 [19200/60000]
loss: 1.114481 [25600/60000]
loss: 1.132760 [32000/60000]
loss: 1.160073 [38400/60000]
loss: 1.108973 [44800/60000]
loss: 1.146243 [51200/60000]
loss: 1.056648 [57600/60000]
Test Error:
 Accuracy: 65.5%, Avg loss: 1.076286

Tensorclass training done! time:  9.1115 s
Epoch 1
-------------------------
loss: 2.299600 [    0/60000]
loss: 2.295847 [ 6400/60000]
loss: 2.272360 [12800/60000]
loss: 2.264007 [19200/60000]
loss: 2.249290 [25600/60000]
loss: 2.209519 [32000/60000]
loss: 2.223882 [38400/60000]
loss: 2.185236 [44800/60000]
loss: 2.179446 [51200/60000]
loss: 2.147614 [57600/60000]
Test Error:
 Accuracy: 38.2%, Avg loss: 2.147300

Epoch 2
-------------------------
loss: 2.154238 [    0/60000]
loss: 2.150466 [ 6400/60000]
loss: 2.089001 [12800/60000]
loss: 2.101484 [19200/60000]
loss: 2.059916 [25600/60000]
loss: 1.991209 [32000/60000]
loss: 2.022341 [38400/60000]
loss: 1.945592 [44800/60000]
loss: 1.943485 [51200/60000]
loss: 1.870688 [57600/60000]
Test Error:
 Accuracy: 59.3%, Avg loss: 1.873784

Epoch 3
-------------------------
loss: 1.906193 [    0/60000]
loss: 1.877595 [ 6400/60000]
loss: 1.758397 [12800/60000]
loss: 1.793871 [19200/60000]
loss: 1.694673 [25600/60000]
loss: 1.642398 [32000/60000]
loss: 1.659138 [38400/60000]
loss: 1.571022 [44800/60000]
loss: 1.588730 [51200/60000]
loss: 1.481123 [57600/60000]
Test Error:
 Accuracy: 61.8%, Avg loss: 1.505067

Epoch 4
-------------------------
loss: 1.572078 [    0/60000]
loss: 1.539070 [ 6400/60000]
loss: 1.390134 [12800/60000]
loss: 1.459417 [19200/60000]
loss: 1.340869 [25600/60000]
loss: 1.341498 [32000/60000]
loss: 1.349597 [38400/60000]
loss: 1.288462 [44800/60000]
loss: 1.318691 [51200/60000]
loss: 1.215316 [57600/60000]
Test Error:
 Accuracy: 63.1%, Avg loss: 1.245432

Epoch 5
-------------------------
loss: 1.323515 [    0/60000]
loss: 1.305851 [ 6400/60000]
loss: 1.142176 [12800/60000]
loss: 1.244354 [19200/60000]
loss: 1.114421 [25600/60000]
loss: 1.149229 [32000/60000]
loss: 1.164128 [38400/60000]
loss: 1.115985 [44800/60000]
loss: 1.151084 [51200/60000]
loss: 1.061394 [57600/60000]
Test Error:
 Accuracy: 64.2%, Avg loss: 1.085556

Training done! time:  33.4731 s

Total running time of the script: (1 minutes 0.071 seconds)

Gallery generated by Sphinx-Gallery

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources