Benchmark of filter-adjoint implementations

benchmark_filter_adjoint.png

benchmark_filter_adjoint.py

#!/usr/bin/env -S uv run python
# ruff: noqa: TID253
# %%
import matplotlib.pyplot as plt
import pandas as pd
import torch
from torch.utils.benchmark import Timer
from deepinv.physics import functional as dF

torch.manual_seed(0)
torch.set_num_threads(1)


def filter_adjoint_via_transpose(x, y, filter_size):
    hf, wf = filter_size
    full = dF.conv_transpose2d(y, x, padding="circular")
    height, width = x.shape[-2:]
    rows = (torch.arange(hf) + height // 2 - hf // 2) % height
    columns = (torch.arange(wf) + width // 2 - wf // 2) % width
    return full.index_select(-2, rows).index_select(-1, columns)


def runtime_ms(fn):
    measurement = Timer("fn()", globals={"fn": fn}).blocked_autorange(min_run_time=0.1)
    return 1e3 * measurement.median


filter_size = (13, 13)
rows = []
for size in (16, 32, 64, 128, 256):
    x = torch.randn(1, 1, size, size)
    y = torch.randn_like(x)
    specialized = lambda: dF.conv2d_filter_adjoint(x, y, filter_size)
    fft = lambda: dF.conv2d_filter_adjoint_fft(x, y, filter_size)
    via_transpose = lambda: filter_adjoint_via_transpose(x, y, filter_size)
    reference = specialized()
    for name, candidate in (("FFT", fft()), ("transpose", via_transpose())):
        relative_error = torch.linalg.vector_norm(
            reference - candidate
        ) / torch.linalg.vector_norm(reference)
        assert relative_error < 1e-5, f"{name} relative error: {relative_error:.2e}"
    specialized_ms = runtime_ms(specialized)
    fft_ms = runtime_ms(fft)
    transpose_ms = runtime_ms(via_transpose)
    rows.append((f"{size}x{size}", specialized_ms, fft_ms, transpose_ms))
results = pd.DataFrame(
    rows,
    columns=[
        "image size",
        "conv2d_filter_adjoint (ms)",
        "conv2d_filter_adjoint_fft (ms)",
        "conv_transpose2d + crop (ms)",
    ],
)
print(results.to_string(index=False, float_format=lambda value: f"{value:.3f}"))

image_size = 128
x = torch.randn(1, 1, image_size, image_size)
y = torch.randn_like(x)
filter_rows = []
for width in (3, 7, 15, 31, 63):
    current_filter_size = (width, width)
    specialized = lambda: dF.conv2d_filter_adjoint(x, y, current_filter_size)
    fft = lambda: dF.conv2d_filter_adjoint_fft(x, y, current_filter_size)
    via_transpose = lambda: filter_adjoint_via_transpose(x, y, current_filter_size)
    reference = specialized()
    for name, candidate in (("FFT", fft()), ("transpose", via_transpose())):
        relative_error = torch.linalg.vector_norm(
            reference - candidate
        ) / torch.linalg.vector_norm(reference)
        assert relative_error < 1e-5, f"{name} relative error: {relative_error:.2e}"
    filter_rows.append(
        (
            f"{width}x{width}",
            runtime_ms(specialized),
            runtime_ms(fft),
            runtime_ms(via_transpose),
        )
    )
filter_results = pd.DataFrame(
    filter_rows,
    columns=["filter size", *results.columns[1:]],
)
print(filter_results.to_string(index=False, float_format=lambda value: f"{value:.3f}"))

fig, axes = plt.subplots(1, 2, figsize=(14, 5))
results.plot(x="image size", y=results.columns[1:], marker="o", logy=True, ax=axes[0])
filter_results.plot(
    x="filter size", y=filter_results.columns[1:], marker="o", logy=True, ax=axes[1]
)
axes[0].set_title(f"Fixed filter size: {filter_size[0]}x{filter_size[1]}")
axes[1].set_title(f"Fixed image size: {image_size}x{image_size}")
for axis in axes:
    axis.set_ylabel("median runtime (ms, log scale)")
    axis.grid(True, which="both", alpha=0.3)
fig.tight_layout()
fig.savefig("benchmark_filter_adjoint.png", dpi=200, bbox_inches="tight")
plt.show()

# %%