Function check_batch

dali/test/python/test_utils.py:252–363

Compare two batches of data, be it dali TensorList or list of numpy arrays. Args: batch1: input batch batch2: input batch batch_size: reference batch size - if None, only equality is enforced eps (float, optional): Used for mean error validation. Defaults to 1e-0

(
    batch1,
    batch2,
    batch_size=None,
    eps=1e-07,
    max_allowed_error=None,
    expected_layout=None,
    compare_layouts=True,
    dump_artifacts=True,
)

# If the `max_allowed_error` is not None, it's checked instead of comparing mean error with `eps`.
def check_batch(
    batch1,
    batch2,
    batch_size=None,
    eps=1e-07,
    max_allowed_error=None,
    expected_layout=None,
    compare_layouts=True,
    dump_artifacts=True,
):
    """Compare two batches of data, be it dali TensorList or list of numpy arrays.

    Args:
        batch1: input batch
        batch2: input batch
        batch_size: reference batch size - if None, only equality is enforced
        eps (float, optional): Used for mean error validation. Defaults to 1e-07.
        max_allowed_error (int or float, optional): If provided the max diff between elements.
        expected_layout (str, optional): If provided, the batches that are DALI types
            will be checked to match this layout. If None, there will be no check
        compare_layouts (bool, optional): Whether to compare layouts between two batches.
            Checked only if both inputs are DALI types. Defaults to True.
    """

    def is_error(mean_err, max_err, eps, max_allowed_error):
        if max_allowed_error is not None:
            if max_err > max_allowed_error:
                return True
        elif mean_err > eps:
            return True
        return False

    import_numpy()
    if isinstance(batch1, dali.backend.TensorListGPU):
        batch1 = batch1.as_cpu()
    if isinstance(batch2, dali.backend.TensorListGPU):
        batch2 = batch2.as_cpu()

    if batch_size is None:
        batch_size = len(batch1)

    def _verify_batch_size(batch):
        if isinstance(batch, dali.backend.TensorListCPU) or isinstance(batch, list):
            tested_batch_size = len(batch)
        else:
            tested_batch_size = batch.shape[0]
        assert (
            tested_batch_size == batch_size
        ), "Incorrect batch size. Expected: {}, actual: {}".format(batch_size, tested_batch_size)

    _verify_batch_size(batch1)
    _verify_batch_size(batch2)

    # Check layouts where possible
    for batch in [batch1, batch2]:
        if expected_layout is not None and isinstance(batch, dali.backend.TensorListCPU):
            assert (
                batch.layout() == expected_layout