SenseCraft

Perceptual and signal losses for image and video models, with our own feature nets, a compose layer that returns the gradient at the prediction directly, and a measured precision and speed record for every term.

Work in progress.

Install

pip install sensecraft            # torch, safetensors, pytorch-msssim
pip install "sensecraft[dinov3]"  # adds huggingface_hub for the DINOv3 terms

Python 3.13. Pretrained weights are converted once to safetensors under ~/.cache/sensecraft (SENSECRAFT_CACHE moves it). VGG, AlexNet and SqueezeNet come from the torchvision URLs; the LPIPS heads ship in the package. DINOv3 checkpoints are gated on Hugging Face, so their first use needs a token.

Use

import torch
from sensecraft import Loss, Term

loss = Loss(
    [
        Term("mse"),
        Term("lpips", 0.5, net="vgg", shift=16),  # random 16 px shift per call
        Term("dino_vit", 250, model="base", use_gram=False),
        Term("ssim", monitor=True),          # logged, not in the total
    ],
    dims=2,                                  # 3 for (B, T, C, H, W) video
    input_range=(-1, 1),
).prepare("cuda", torch.bfloat16, torch.channels_last, compile="term")

out = loss(fake, real)                       # out.total, out.terms
out = loss.value_and_seed(fake, real)        # also out.seed = d total / d fake, no graph kept

• Each term reads its images in the range it was trained for; Loss converts from input_range.
• shift=s crops the pair together by s pixels at a random offset on every call, so a net's pooling grid lands on a new phase each step; monitor terms see the whole image. See Shift.
• Image terms run per frame on video; video terms see the clip.
• prepare puts the feature nets in the compute dtype and layout and compiles by scope: off, term or loss. Distances stay in fp32.
• Plain terms (pixel, edge, FFT, SSIM, video) run in strict fp32: autocast off, fp32 inputs, and IEEE fp32 convs and matmuls instead of TF32.
• Feature nets run their convs, matmuls and attention in the compute dtype. The layer that reads the image keeps the image's low bits, and the residual stream, norms, activations and layer scales run in fp32; see Precision.
• Nothing reads the host: totals and terms stay on the device.

Metrics on the same units, without gradient:

from sensecraft.metrics import lpips, ms_ssim, psnr, ssim

psnr(fake, real, data_range=2.0)

Terms

A new term is an nn.Module with forward(input, target) -> scalar, registered with @register(name, Range.UNIT, dims="2d", fp32=False). fp32=True runs it in strict fp32.

Layout

src/sensecraft/
  nets/        feature nets, features only: VGG16, AlexNet, SqueezeNet, ConvNeXt, DINOv3 ViT and ConvNeXt
  weights/     upstream sources, conversion to cached safetensors, the LPIPS heads
  distance/    LPIPS and feature distances, in fp32
  perceptual/  terms: a net and a distance
  pixel/       pixel regression terms
  signal/      FFT, edge, SSIM and video terms
  kernels/     fused ops with plain PyTorch fallbacks, the split-input first layer
  compose.py   Term, Loss, value_and_seed
  precision.py strict fp32 for the plain terms
  prepare.py   dtype, layout and compile scope
  metrics.py   metrics on device tensors
measure/       the precision and speed chain, see measure/README.md
checks/        behavior checks run by hand, see checks/README.md

Precision

Every perceptual term matches upstream at fp64 (S1, error at the level of fp64 rounding), and every fixed term matches 0.3.13 exactly at fp64 (S5). Against fp64, at the same dtype as upstream:

The LPIPS distance runs in fp32: in fp16 the channel sum of squares overflows, and in bf16 the difference of two close unit vectors cancels. Value error drops 3 to 5x in bf16 and from 28% to 0.02% for SqueezeNet in fp16. Its unit vectors take the eps inside the root, so an all-zero feature vector has a finite gradient. DINOv3 ConvNeXt rows are measured on random stand-in weights.

Near-identical images

A perceptual term's gradient at the prediction comes from the difference between the features of the prediction and of the target. When the two images differ by less than the spacing of the compute dtype, that difference is rounded away and the gradient keeps little of its direction. This is not fp16 only: bf16 keeps 8 mantissa bits, fp16 and TF32 keep 10.

S7 (measure/s7_near.py) runs each term on 8 ImageNet validation images against the same images plus Gaussian noise, and compares the gradient with an fp64 run. Rounding one point at a time to fp16 in an otherwise fp64 net showed where the bits go: casting the image to the compute dtype before the first layer loses almost all of it, and in DINOv3 ViT the residual stream loses most of the rest. Rounding the backward pass at a training loss scale costs nothing.

So the layer that reads the image runs on the image split in two, the image rounded to the compute dtype and the scaled remainder, and adds the two outputs in fp32. In DINOv3 ViT and the ConvNeXts the residual stream, norms, rope, activations and layer scales run in fp32. Convs, matmuls and attention stay in the compute dtype. Cosine to fp64, before and after:

Strict fp32 stays above 0.99 at every noise. Noise 1e-3 on [-1, 1] images is about 66 dB PSNR, so this matters where a reconstruction is already near exact, and for DINOv3 ViT in bf16 or fp16 well before that. DINOv3 ViT-B at larger differences, mean of 4 batches of 8, before and after:

Every term and mode is in measure/figures/s7_near.svg; dashed lines are before.

Changed on purpose against 0.3.13: tssim reads [-1, 1] clips, which its formulas assume; ssim3d and stssim read [0, 1] clips instead of guessing the range from the input's minimum.

Shift

A net with pooling sends its gradient through a fixed grid: LPIPS VGG's deepest tap, behind four max-pools, carries about 70% of its gradient norm, and each pool passes the gradient to one input of four. The same pixels get the same push at every step, and a pattern locked to that grid can build up. Term(..., shift=s) crops the pair by s pixels at a random offset per call, drawn on the host and kept out of compiled graphs; the window size is fixed, so nothing recompiles.

S8 (measure/s8_shift.py) takes the gradient at a blurred and noised copy of 8 ImageNet validation images over 16 calls. Persistence is |mean gradient|^2 over mean |gradient|^2: 1 when every call pushes the same way.

About half of LPIPS VGG's push moves with the shift and half stays with the image content; the mean gradient keeps about two thirds of its power above a quarter of the sampling rate. 16 px, the stride of VGG's deepest tap, gets most of the effect. A shifted LPIPS VGG call at batch 32 and 256 px takes 24.7 ms against 25.8 ms, on the smaller window.

Speed

Batch 32, 256 px, one RTX 5090, bf16, channels-last, compiled per term, forward and the gradient at the prediction; the 0.4.0 and current columns are measured in the same session:

The split first layer and the fp32 stream cost 2.2 ms on LPIPS VGG, whose first layer works at full resolution, and under 1 ms on the others.

LPIPS runs every conv, bias and ReLU as one cuDNN op with an input-only backward, since feature net weights never take gradients; its heads are a channel sum inside the fused distance, and VGG's pools keep no indices (peak memory 2.8 to 1.8 GB). LPIPS runs channels-last inside whatever layout it is given; from NCHW input that is 36 ms otherwise. DINOv3 embeds patches with one matmul. The full numbers, per-layer floor ratios and the changes that were tried and dropped are in measure/.

License

Apache-2.0. The LPIPS heads and the pretrained nets keep their own licenses; see LICENSE-NOTICES.