Data Databatch

databatch

CLASS	DESCRIPTION
DataBatch	Data class to hold a batch of data from

Classes

DataBatch(data: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) dataclass

Data class to hold a batch of data from DatasetWeighted.

ATTRIBUTE	DESCRIPTION
data	Input data tensor for the batch. TYPE: Tensor
weights	Weights tensor for the batch, if provided. TYPE: Tensor | None
target	Target tensor for the batch, if provided. TYPE: Tensor | None

Examples:

Creating a Batch instance

>>> data = torch.randn(32, 3, 224, 224)
>>> weights = torch.rand(32)
>>> target = torch.randint(0, 2, (32,))
>>> batch = DataBatch(data=data, weights=weights, target=target)

METHOD	DESCRIPTION
to	Move all tensors in the batch to the specified device.
device	Return the device of the batch.
cpu	Move all tensors in the batch to CPU.
cuda	Move all tensors in the batch to CUDA (GPU).
pin_memory	Pin memory for all tensors in the batch for faster GPU transfers.
has_target	Check if the batch has target values.
has_weights	Check if the batch has weights.
check_shape	Check if data, weights, and target (if present) have compatible shapes.
check_type	Check if data, weights, and target (if present) are of type torch.Tensor.
check_dim	Check if data, weights, and target (if present) have valid dimensions.
check	Run all checks on the batch.

Functions

to(device: torch.device) -> DataBatch

Move all tensors in the batch to the specified device.

Source code in spectre/data/databatch.py

def to(self, device: torch.device) -> "DataBatch":
    """Move all tensors in the batch to the specified device."""
    data = self.data.to(device)
    weights = self.weights.to(device) if self.weights is not None else None
    target = self.target.to(device) if self.target is not None else None

    return DataBatch(data=data, weights=weights, target=target)

device() -> torch.device

Return the device of the batch.

Source code in spectre/data/databatch.py

def device(self) -> torch.device:
    """Return the device of the batch."""
    return self.data.device

cpu() -> DataBatch

Move all tensors in the batch to CPU.

Source code in spectre/data/databatch.py

def cpu(self) -> "DataBatch":
    """Move all tensors in the batch to CPU."""
    return self.to(torch.device("cpu"))

cuda() -> DataBatch

Move all tensors in the batch to CUDA (GPU).

Source code in spectre/data/databatch.py

def cuda(self) -> "DataBatch":
    """Move all tensors in the batch to CUDA (GPU)."""
    return self.to(torch.device("cuda"))

pin_memory() -> DataBatch

Pin memory for all tensors in the batch for faster GPU transfers.

Note: Pinned memory only applies to CPU tensors. If tensors are already on GPU, they are returned as-is.

Source code in spectre/data/databatch.py

def pin_memory(self) -> "DataBatch":
    """
    Pin memory for all tensors in the batch for faster GPU transfers.

    Note: Pinned memory only applies to CPU tensors. If tensors are
    already on GPU, they are returned as-is.
    """
    # Already on GPU, return as-is
    if self.data.is_cuda:
        return self

    data = self.data if self.data.is_pinned() else self.data.pin_memory()
    weights = None
    if self.weights is not None:
        weights = (
            self.weights if self.weights.is_pinned() else self.weights.pin_memory()
        )
    target = None
    if self.target is not None:
        target = (
            self.target if self.target.is_pinned() else self.target.pin_memory()
        )
    return DataBatch(data=data, weights=weights, target=target)

has_target() -> bool

Check if the batch has target values.

Source code in spectre/data/databatch.py

def has_target(self) -> bool:
    """Check if the batch has target values."""
    return self.target is not None

has_weights() -> bool

Check if the batch has weights.

Source code in spectre/data/databatch.py

def has_weights(self) -> bool:
    """Check if the batch has weights."""
    return self.weights is not None

check_shape() -> bool

Check if data, weights, and target (if present) have compatible shapes.

Source code in spectre/data/databatch.py

def check_shape(self) -> bool:
    """Check if data, weights, and target (if present) have compatible shapes."""
    if self.weights is not None:
        check_same_len(self.data, self.weights)
    if self.target is not None:
        check_same_len(self.data, self.target)
    return True

check_type() -> bool

Check if data, weights, and target (if present) are of type torch.Tensor.

Source code in spectre/data/databatch.py

def check_type(self) -> bool:
    """Check if data, weights, and target (if present) are of type torch.Tensor."""
    check_is_tensor(self.data)
    if self.weights is not None:
        check_is_tensor(self.weights)
    if self.target is not None:
        check_is_tensor(self.target)
    return True

check_dim() -> bool

Check if data, weights, and target (if present) have valid dimensions.

Data must be 2D, weights must be 1D, targets can be 1D or 2D.

Source code in spectre/data/databatch.py

def check_dim(self) -> bool:
    """
    Check if data, weights, and target (if present) have valid dimensions.

    Data must be 2D, weights must be 1D, targets can be 1D or 2D.
    """
    check_2d(self.data)
    if self.weights is not None:
        check_1d(self.weights)
    if self.target is not None:
        # Target can be 1D or 2D (DatasetWeighted converts 1D to 2D)
        if self.target.ndim not in (1, 2):
            raise ValueError(f"Target must be 1D or 2D, got {self.target.ndim}D")
    return True

check() -> bool

Run all checks on the batch.

Source code in spectre/data/databatch.py

def check(self) -> bool:
    """Run all checks on the batch."""
    return self.check_type() and self.check_dim() and self.check_shape()

Functions

collate_databatch(batch: list[DataBatch]) -> DataBatch

Collate function for batching DataBatch objects.

PARAMETER	DESCRIPTION
batch	List of DataBatch objects to collate. TYPE: list[DataBatch]

RETURNS	DESCRIPTION
DataBatch	Collated batch with stacked tensors.

Source code in spectre/data/databatch.py

def collate_databatch(batch: list[DataBatch]) -> DataBatch:
    """
    Collate function for batching DataBatch objects.

    Parameters
    ----------
    batch : list[DataBatch]
        List of DataBatch objects to collate.

    Returns
    -------
    DataBatch
        Collated batch with stacked tensors.
    """
    data = torch.stack([item.data for item in batch])

    if batch[0].weights is not None:
        weights = torch.stack([item.weights for item in batch])
    else:
        weights = None

    if batch[0].target is not None:
        target = torch.stack([item.target for item in batch])
    else:
        target = None

    return DataBatch(data=data, weights=weights, target=target)