Overview

PyTorch 1.11, released on February 11, 2022, introduces TorchData for data loading and integrates functorch for functional transformations.

Main Features

TorchData

TorchData provides composable DataPipes for building flexible and performant data loading pipelines, progressively replacing the older DataLoader.

python 
from torchdata.datapipes.iter import IterableWrapper

# Composable data pipeline
dp = IterableWrapper(range(100))
dp = dp.shuffle(buffer_size=20)
dp = dp.batch(batch_size=8)
dp = dp.map(lambda batch: [x * 2 for x in batch])

for batch in dp:
    print(batch)  # [list of 8 doubled elements]
    break

functorch

functorch brings functional transformations like vmap (automatic vectorization) and grad (differentiation), inspired by JAX.

python 
import torch
from functorch import vmap, grad

# Automatic vectorization with vmap
def compute(x):
    return torch.sum(x ** 2)

batch = torch.randn(10, 3)
results = vmap(compute)(batch)  # applied to each row
print(results.shape)  # (10,)

# Functional gradient
grad_fn = grad(compute)
print(grad_fn(torch.tensor([1.0, 2.0, 3.0])))  # [2., 4., 6.]

Sources