calibration

calibration

Calibration data provider for static quantization.

CalibrationDataProvider

CalibrationDataProvider(dataloader, observation_keys, num_calibration_steps=128, device=None)

Yields observation tuples from a VersatIL dataloader for calibration.

Extracts observations from dataloader batches and yields them as positional tensor tuples matching the ExportablePolicy's key order.

The VersatIL dataloader already normalizes and tokenizes observations, so no additional preprocessing is needed here.

Initialize calibration data provider.

Parameters:

Name	Type	Description	Default
dataloader	DataLoader	VersatIL dataloader yielding normalized, tokenized samples.	required
observation_keys	list[str]	Keys defining positional argument order.	required
num_calibration_steps	int	Maximum number of calibration batches.	128
device	device | None	Device for calibration tensors. Defaults to CPU.	None

Source code in src/versatil/quantization/calibration.py

def __init__(
    self,
    dataloader: torch.utils.data.DataLoader,
    observation_keys: list[str],
    num_calibration_steps: int = 128,
    device: torch.device | None = None,
) -> None:
    """Initialize calibration data provider.

    Args:
        dataloader: VersatIL dataloader yielding normalized, tokenized samples.
        observation_keys: Keys defining positional argument order.
        num_calibration_steps: Maximum number of calibration batches.
        device: Device for calibration tensors. Defaults to CPU.
    """
    self._dataloader = dataloader
    self._observation_keys = observation_keys
    self._num_calibration_steps = num_calibration_steps
    if device is None:
        device = torch.device("cpu")
    self.device = device

__iter__

__iter__()

Yield calibration batches as positional tensor tuples on self.device.

Source code in src/versatil/quantization/calibration.py

def __iter__(self) -> Iterator[tuple[torch.Tensor, ...]]:
    """Yield calibration batches as positional tensor tuples on self.device."""
    for step, batch in enumerate(self._dataloader):
        if step >= self._num_calibration_steps:
            break
        observation = batch[SampleKey.OBSERVATION.value]
        yield tuple(
            observation[key].to(self.device) for key in self._observation_keys
        )