[Feat] Refactoring type assertions

🚀 Feature

Refactoring or removing type and shape assertions like in https://github.com/kornia/kornia/blob/65777029ff430f841f19aa5cd08ddb3c4ca36338/kornia/losses/ssim.py#L44-L65

Motivation

These assertions are useful for quick prototyping and debugging. However, after that, they are just slowing down the components, especially if they are used repetitively.

Moreover, one of the goal of the roadmap (#700) is to use PyTorch JIT and a scripted function already has type assertion built-in as it transformed to C++.

Alternatives

I thought of multiple solutions:

1. Removing the type and shape assertions. Probably the easiest to implement, but not the most secure or user friendly.
2. Adding a assert_type: bool = True argument to every component and the type assertion would only be performed if assert_type == True:. It would also be best to centralize the type assertions (there is a lot of duplicated code currently).
3. Similarly to the previous idea, but instead of using an argument, use a global variable assert_type, that can be modified with a context manager (inspired by torch.set_grad_enabled()).

Additional context

I’m implementing a Image Quality Assessment (IQA) library (piqa). I’m considering merging it into kornia, but I have put a big effort into making it fast (and supporting JIT) and I don’t want to ruin my efforts with unnecessary type assertions.

[francois-rozet]

I am, but I will be very busy until the end of July (finishing my thesis). After that, I should come back to you to know how I can help. You can close the issue until then 😉

[francois-rozet]

Hi @edgarriba,

I didn’t knew the assert could be disabled, it is a great solution! I’ve used it to implement assertion in PIQA for PSNR and SSIM and benchmark (see benchmark.py) the difference.

Conclusions:

1. Adding assertions does bottleneck (by a factor 8x the PSNR and 2x the SSIM) the implementations.
2. Using the -O flag (with my implementation) completely removes the bottleneck!
3. My implementation of the assertions is compliant with JIT even if __debug__ is a global variable!

---

First, here are the runtimes currently (kornia with if ...: raise ... and piqa without any assertions).

PSNR
----
        method      value  time-sync  time-async
1    piqa.psnr  27.567236   0.042983    0.043069
2    piqa.PSNR  27.567236   0.055748    0.055898
3  kornia.PSNR  27.567261   0.121797    0.121839

SSIM
----
        method     value  time-sync  time-async
1    piqa.SSIM  0.724984   0.470641    0.282218
2  kornia.SSIM  0.731942   1.211249    1.209685

N.B. The value of kornia.SSIM is different because the inputs are not padded in piqa.SSIM.

Then I called the following function in both PSNR and SSIM in piqa.

def _assert(
    tensors: List[torch.Tensor],
    dim: Tuple[int, int] = (0, -1),
    n_channels: int = 0,
    value_range: Tuple[float, float] = (0., -1.),
    debug: bool = __debug__,
) -> None:
    r""""""

    if not debug:
        return

    ref_size = tensors[0].size()
    ref_device = tensors[0].device

    for t in tensors:
        assert t.size() == ref_size, f'Expected tensors to be of the same shape, got {ref_size} and {t.size()}'

        assert t.device == ref_device, f'Expected tensors to be on the same device, got {ref_device} and {t.device}'

        if dim[0] == dim[1]:
            assert t.dim() == dim[0], f'Expected number of dimensions to be {dim[0]}, got {t.dim()}'
        elif dim[0] < dim[1]:
            assert dim[0] <= t.dim() <= dim[1], f'Expected number of dimensions to be between {dim[0]} and {dim[1]}, got {t.dim()}'
        elif dim[0] > 0:
            assert dim[0] <= t.dim(), f'Expected number of dimensions to be greater or equal to {dim[0]}, got {t.dim()}'

        if n_channels > 0:
            assert t.size(1) == n_channels, f'Expected number of channels ot be {n_channels}, got {t.size(1)}'

        if value_range[1] > value_range[0]:
            assert value_range[0] <= t.min(), f'Expected values to be greater or equal to {value_range[0]}, got {t.min()}'
            assert t.max() <= value_range[1], f'Expected values to be smaller or equal to {value_range[1]}, got {t.max()}'

which gave, without -O flag,

PSNR
----
        method      value  time-sync  time-async
2    piqa.psnr  27.567236   0.294254    0.294326
3    piqa.PSNR  27.567236   0.325271    0.325338

SSIM
----
        method     value  time-sync  time-async
1    piqa.SSIM  0.724984   0.906343    0.905503

and, with the -O flag,

PSNR
----
        method      value  time-sync  time-async
1    piqa.psnr  27.567236   0.044030    0.044139
2    piqa.PSNR  27.567236   0.078526    0.079137

SSIM
----
        method     value  time-sync  time-async
1    piqa.SSIM  0.724984   0.472842    0.286530