quantlab.model.torch_data.Batch

class quantlab.model.torch_data.Batch(x, y, mask, y_raw, where)

Bases: NamedTuple

What every step hook, _loss and _forward receive.

where and mask share the batch’s sample dimensions: [S_t] for one cross-section, [B] for a batch of samples. The base moves a batch to the model’s device and replaces x by the head’s _transform_feature output before a hook sees it. A NamedTuple, so PyTorch’s default collation and pinning work on it as is.

Attributes:

xtorch.Tensor

Features shaped by the dataset, [S_t, N, F] for a cross-section.

ytorch.Tensor

Training target, mask.shape + (L,), 0 where mask is False.

masktorch.Tensor

True where the sample has a valid training target; a loss counts only these.

y_rawtorch.Tensor

Raw labels, mask.shape + (L,), NaN where missing.

wheretuple[torch.Tensor, torch.Tensor]

The timestamp index and symbol index of every sample in the panel, each shaped like mask.

Examples

>>> batch = Batch(
...     x=torch.zeros(2, 1, 1), y=torch.tensor([[1.0], [0.0]]),
...     mask=torch.tensor([True, False]),
...     y_raw=torch.tensor([[1.0], [float("nan")]]),
...     where=(torch.tensor([4, 4]), torch.tensor([0, 3])),
... )
>>> int(batch.mask.sum()), batch.where[1].tolist()
(1, [0, 3])

mask

Alias for field number 2

pin_memory()

Return the batch with every tensor copied into pinned (page-locked) memory.

Examples

>>> batch.pin_memory().x.is_pinned()
True

tensors()

Every tensor of the batch: x, y, mask, y_raw and both where.

Examples

>>> [tuple(tensor.shape) for tensor in batch.tensors()]
[(2, 1, 1), (2, 1), (2,), (2, 1), (2,), (2,)]

to(device, non_blocking=False)

Return the batch with every tensor on device.

non_blocking is passed to every Tensor.to: from pinned memory to CUDA, the copies are only queued on the current stream.

Examples

>>> batch.to("cpu").x.device.type
'cpu'

where

Alias for field number 4

x

Alias for field number 0

y

Alias for field number 1

y_raw

Alias for field number 3