File size: 4,652 Bytes
e199518 4ee800d e199518 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | ---
tags:
- kernel
---
# kernel_image_resize
A pure-Triton Hub kernel that fuses the **resize + rescale + normalize** preprocessing
pipeline run by ~150 `transformers` fast image processors (`TorchvisionBackend`: resize →
fold(rescale, normalize)) into a single GPU pass. It takes raw CHW `uint8` images and
returns the normalized `(N, C, out_h, out_w)` float tensor with no intermediate
full-resolution float buffer.
On a ragged SigLIP-so400m batch (A100, N=32, inputs 384–1024², out 384², bicubic+antialias)
the default backend runs in **1.29 ms/iter vs 3.90 ms for the fast processor (~3× faster)**
and 2.89 ms for torchvision's own per-image loop, at parity ≤1e-4 vs torchvision-float.
It ships as a `kernels` universal build variant (no compiled extension, just Triton), so it
loads on any CUDA PyTorch build via `get_kernel`.
## Usage
```python
import torch
from kernels import get_kernel
kir = get_kernel("Molbap/kernel_image_resize", revision="main", trust_remote_code=True)
# a list of different-H×W uint8 CHW images (the ragged case torchvision loops over)
images = [torch.randint(0, 256, (3, h, w), dtype=torch.uint8, device="cuda")
for h, w in [(640, 480), (800, 600), (384, 1024)]]
pixel_values = kir.resize_normalize(
images,
size=384, # int (square), (H, W), or {"height", "width"}
image_mean=[0.5, 0.5, 0.5],
image_std=[0.5, 0.5, 0.5],
rescale_factor=1 / 255,
resample="bicubic", # or "bilinear", or a PIL resample int
antialias=True, # match the ViT/CLIP/SigLIP default
)
# -> (3, 3, 384, 384) float32, ready for the model
```
Requires `kernels >= 0.15` (published as a `kernel` repo type). `trust_remote_code=True` is needed
because `Molbap` is a personal namespace, not the auto-trusted `kernels-community` org.
`resize_normalize` accepts a stacked `(N, C, H, W)` tensor or a ragged list of CHW
tensors. `resize_normalize_ragged` is the same kernel, list-only.
## With a transformers processor
There is no `use_kernels=True` hook for image processors — that machinery swaps `nn.Module`
layer forwards inside the model, not processor code. Use the kernel directly with the
processor's config (see `example_transformers.py` for a runnable version):
```python
from kernels import get_kernel
kir = get_kernel("Molbap/kernel_image_resize", revision="main", trust_remote_code=True)
_PIL_RESAMPLE = {0: "bilinear", 2: "bilinear", 3: "bicubic"}
def preprocess_with_kernel(processor, images):
size = processor.size # must be fixed {"height", "width"}; no crop/pad
return kir.resize_normalize(
images, (size["height"], size["width"]),
processor.image_mean, processor.image_std,
rescale_factor=float(processor.rescale_factor),
resample=_PIL_RESAMPLE[int(processor.resample)],
antialias=bool(getattr(processor, "antialias", True)),
)
```
## Backends
- `backend="separable"` (default): two-pass `uint8` kernel doing `taps+taps` loads —
torchvision's own separable algorithm. Fastest (~3× the fast processor on the batch
above); parity ≤1e-4 vs torchvision-float. The float intermediate makes it more accurate
than, but not bit-identical to, torchvision's fixed-point `uint8` intermediate.
- `backend="fused"`: a single 2D launch, `taps×taps` loads per output pixel. Same parity,
kept as the reference path but ~9× slower than separable (the 2D float load count is the
reason a separable pass wins — see `benchmarks/benchmark.py`).
## Parity notes
The resampling weights match PyTorch aten `UpSampleKernel`. Antialiased bicubic uses the
PIL cubic coefficient `a=-0.5`; non-antialiased bicubic uses Keys `a=-0.75`. The
antialias renormalize-truncate window applies on every axis, including upsampling dims.
## Center crop / shortest-edge
Pass `crop_size` to resize then center-crop in one pass (the crop is folded into the
output-coordinate mapping, no extra buffer). `resize_mode="shortest_edge"` does
aspect-preserving resize (short side = `size`) then crop — the CLIP / DINOv2 pipeline.
```python
# CLIP/DINOv2-style: resize shortest edge to 256, center-crop 224
pv = kir.resize_normalize(images, 256, mean, std, resample="bicubic", antialias=True,
crop_size=224, resize_mode="shortest_edge")
```
`example_transformers.py` derives all of this from a processor's config automatically.
## Scope
Resize (+ optional center crop) + rescale + normalize. It does **not** pad — padding
processors (many detection models) run a different pipeline. The `fused` backend is
resize-only; crop is handled by the `separable` backend.
|