diff --git a/burn-book/src/building-blocks/metric.md b/burn-book/src/building-blocks/metric.md index 93c7216bcd5..0f27577805b 100644 --- a/burn-book/src/building-blocks/metric.md +++ b/burn-book/src/building-blocks/metric.md @@ -25,6 +25,7 @@ throughout the training process. We currently offer a restricted range of metric | Vision Metric | Description | | ------------- | ---------------------------------------------------------------------------------------------------- | +| A-FINE | Computes the Adaptive Fidelity-Naturalness Evaluator (A-FINE) full-reference perceptual quality metric built on CLIP ViT-B/32 features | | Dice | Computes the Dice-Sorenson coefficient (DSC) for evaluating overlap between binary masks | | DISTS | Computes the Deep Image Structure and Texture Similarity (DISTS) metric for image quality assessment | | FID | Computes the Frechet Inception Distance (FID) for evaluating generative model quality | diff --git a/crates/burn-train/src/metric/vision/afine/calibrators.rs b/crates/burn-train/src/metric/vision/afine/calibrators.rs new file mode 100644 index 00000000000..b818b1b5b4e --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/calibrators.rs @@ -0,0 +1,250 @@ +//! Calibrators, adapter, and final-score mapping for A-FINE. +//! +//! Four small modules sit between the heads and the final per-sample +//! quality score: +//! +//! - [`NrCalibrator`] — logistic mapping of the naturalness head's raw +//! output into `(-2, 2)`. Two learnable scalars. +//! - [`FrCalibratorWithLimit`] — logistic mapping of the fidelity head's +//! raw output into `(-2, 2)`, with `yita3` clamped to `[0.05, 0.95]` +//! and `yita4` to `[0.01, 0.70]` on every forward. +//! - [`AfineAdapter`] — `D = exp(softplus(k) * (N_ref - N_dis)) * N_dis + F`. +//! Single learnable scalar `k`. +//! - [`scale_finalscore`] — fixed logistic into `(0, 100)` with the +//! paper-reported constants. +//! +//! All three calibrators implement the same logistic shape: +//! `out = (yita1 - yita2) * sigmoid((x - yita3) / (|yita4| + eps)) + yita2`. +//! This is the algebraic equivalent of PyIQA's two-branch +//! `if exp_pow >= 10` formulation, rewritten as a single expression so +//! it batches correctly. PyIQA's branch only works on 0-D scalar +//! tensors. + +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Module, Param}; +use burn::tensor::Tensor; +use burn::tensor::activation::{sigmoid, softplus}; +use burn::tensor::backend::Backend; + +const NR_YITA1: f64 = 2.0; +const NR_YITA2: f64 = -2.0; +const NR_YITA3_INIT: f32 = 4.9592; +const NR_YITA4_INIT: f32 = 21.5968; + +const FR_YITA1: f64 = 2.0; +const FR_YITA2: f64 = -2.0; +const FR_YITA3_INIT: f32 = 0.5; +const FR_YITA4_INIT: f32 = 0.15; +const FR_YITA3_MIN: f32 = 0.05; +const FR_YITA3_MAX: f32 = 0.95; +const FR_YITA4_MIN: f32 = 0.01; +const FR_YITA4_MAX: f32 = 0.70; + +const ADAPTER_K_INIT: f32 = 5.0; + +const SCALE_YITA1: f64 = 100.0; +const SCALE_YITA2: f64 = 0.0; +const SCALE_YITA3: f64 = -1.971_0; +const SCALE_YITA4: f64 = -2.373_4; + +/// Numerical-stability epsilon in the logistic denominator. Matches +/// PyIQA exactly; do not change without coordinating a parity-test +/// re-capture. +const EPS: f64 = 1e-10; + +/// Apply `(yita1 - yita2) * sigmoid((x - yita3) / (|yita4| + eps)) + yita2` +/// element-wise, broadcasting the 1-D scalar parameters over the input. +fn logistic_calibrate( + x: Tensor, + yita3: Tensor, + yita4_abs: Tensor, + yita1: f64, + yita2: f64, +) -> Tensor { + let yita3 = yita3.reshape([1, 1]); + let denom = yita4_abs.reshape([1, 1]).add_scalar(EPS); + let inner = (x - yita3) / denom; + sigmoid(inner).mul_scalar(yita1 - yita2).add_scalar(yita2) +} + +/// Configuration for [`NrCalibrator`]. +#[derive(Config, Debug)] +pub(crate) struct NrCalibratorConfig {} + +impl NrCalibratorConfig { + pub(crate) fn init(&self, device: &B::Device) -> NrCalibrator { + NrCalibrator { + yita3: Param::from_tensor(Tensor::from_floats([NR_YITA3_INIT], device)), + yita4: Param::from_tensor(Tensor::from_floats([NR_YITA4_INIT], device)), + } + } +} + +/// Naturalness logistic calibrator. Maps `[B, 1]` into `(-2, 2)`. +#[derive(Module, Debug)] +pub(crate) struct NrCalibrator { + pub(crate) yita3: Param>, + pub(crate) yita4: Param>, +} + +impl NrCalibrator { + pub(crate) fn forward(&self, x: Tensor) -> Tensor { + logistic_calibrate( + x, + self.yita3.val(), + self.yita4.val().abs(), + NR_YITA1, + NR_YITA2, + ) + } +} + +/// Configuration for [`FrCalibratorWithLimit`]. +#[derive(Config, Debug)] +pub(crate) struct FrCalibratorWithLimitConfig {} + +impl FrCalibratorWithLimitConfig { + pub(crate) fn init(&self, device: &B::Device) -> FrCalibratorWithLimit { + FrCalibratorWithLimit { + yita3: Param::from_tensor(Tensor::from_floats([FR_YITA3_INIT], device)), + yita4: Param::from_tensor(Tensor::from_floats([FR_YITA4_INIT], device)), + } + } +} + +/// Fidelity logistic calibrator with on-forward clamping of `yita3` and +/// `yita4`. PyIQA clamps the values used in the formula on every call; +/// the stored parameter is unchanged. +#[derive(Module, Debug)] +pub(crate) struct FrCalibratorWithLimit { + pub(crate) yita3: Param>, + pub(crate) yita4: Param>, +} + +impl FrCalibratorWithLimit { + pub(crate) fn forward(&self, x: Tensor) -> Tensor { + // Match PyIQA semantics exactly: clamp first, then abs. The + // clamp range is positive so the abs is a no-op for in-range + // values, but for an out-of-range checkpoint or a parameter + // that drifts negative during training the order matters. + let yita3 = self.yita3.val().clamp(FR_YITA3_MIN, FR_YITA3_MAX); + let yita4 = self.yita4.val().clamp(FR_YITA4_MIN, FR_YITA4_MAX); + logistic_calibrate(x, yita3, yita4.abs(), FR_YITA1, FR_YITA2) + } +} + +/// Configuration for [`AfineAdapter`]. +#[derive(Config, Debug)] +pub(crate) struct AfineAdapterConfig {} + +impl AfineAdapterConfig { + pub(crate) fn init(&self, device: &B::Device) -> AfineAdapter { + AfineAdapter { + k: Param::from_tensor(Tensor::from_floats([ADAPTER_K_INIT], device)), + } + } +} + +/// Fuses the calibrated naturalness and fidelity scores into a single +/// raw `D` value. +/// +/// `D = exp(softplus(k) * (N_ref - N_dis)) * N_dis + F`. The `softplus` +/// wrapper enforces `k > 0` without constraining the stored parameter. +#[derive(Module, Debug)] +pub(crate) struct AfineAdapter { + pub(crate) k: Param>, +} + +impl AfineAdapter { + pub(crate) fn forward( + &self, + x_nr: Tensor, + ref_nr: Tensor, + xref_fr: Tensor, + ) -> Tensor { + let k_pos = softplus(self.k.val(), 1.0).reshape([1, 1]); + let weight = (k_pos * (ref_nr - x_nr.clone())).exp(); + weight * x_nr + xref_fr + } +} + +/// Map a raw adapter score into `(0, 100)` via a fixed 4-parameter +/// logistic. Constants are the paper-reported defaults. +pub(crate) fn scale_finalscore(score: Tensor) -> Tensor { + let denom = SCALE_YITA4.abs() + EPS; + let inner = score.sub_scalar(SCALE_YITA3).div_scalar(denom); + sigmoid(inner) + .mul_scalar(SCALE_YITA1 - SCALE_YITA2) + .add_scalar(SCALE_YITA2) +} + +#[cfg(test)] +mod tests { + use super::*; + use burn_flex::Flex; + + type TestBackend = Flex; + + #[test] + fn nr_calibrator_maps_to_bounded_range() { + let device = Default::default(); + let calibrator = NrCalibratorConfig::new().init::(&device); + + let extremes = Tensor::::from_floats([[-1000.0], [0.0], [1000.0]], &device); + let out = calibrator.forward(extremes); + let values = out.into_data().to_vec::().unwrap(); + + for v in &values { + assert!(*v >= -2.0 && *v <= 2.0, "out-of-range value: {v}"); + } + // Monotonic increasing. + assert!(values[0] < values[1]); + assert!(values[1] < values[2]); + } + + #[test] + fn fr_calibrator_clamp_does_not_panic() { + let device = Default::default(); + let calibrator = FrCalibratorWithLimitConfig::new().init::(&device); + + let input = Tensor::::from_floats([[0.5], [1.5], [-0.5]], &device); + let out = calibrator.forward(input); + + assert_eq!(out.dims(), [3, 1]); + } + + #[test] + fn adapter_forward_propagates_shape() { + let device = Default::default(); + let adapter = AfineAdapterConfig::new().init::(&device); + + let nr_dis = Tensor::::from_floats([[0.5], [-0.3]], &device); + let nr_ref = Tensor::::from_floats([[0.7], [-0.1]], &device); + let fr = Tensor::::from_floats([[0.2], [0.4]], &device); + + let out = adapter.forward(nr_dis, nr_ref, fr); + assert_eq!(out.dims(), [2, 1]); + } + + #[test] + fn scale_finalscore_maps_to_0_100_range() { + let device = Default::default(); + + let scores = + Tensor::::from_floats([[-1000.0], [-1.971], [1000.0]], &device); + let out = scale_finalscore(scores); + let values = out.into_data().to_vec::().unwrap(); + + assert!(values[0] >= 0.0 && values[0] <= 100.0); + assert!(values[2] >= 0.0 && values[2] <= 100.0); + // At yita3 = -1.971 the sigmoid argument is 0, so the output is + // 100 * 0.5 = 50. + assert!( + (values[1] - 50.0).abs() < 0.5, + "midpoint should be ~50, got {}", + values[1] + ); + } +} diff --git a/crates/burn-train/src/metric/vision/afine/clip_attention.rs b/crates/burn-train/src/metric/vision/afine/clip_attention.rs new file mode 100644 index 00000000000..0f0807926c0 --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/clip_attention.rs @@ -0,0 +1,137 @@ +//! Self-attention block matching PyTorch's `nn.MultiheadAttention` wire +//! format. +//! +//! PyTorch stores Q/K/V as a single fused `in_proj_weight` of shape +//! `[3 * d_model, d_model]` and `in_proj_bias` of shape `[3 * d_model]`. +//! Burn's [`burn_nn::attention::MultiHeadAttention`] uses three separate +//! Linear layers, so the CLIP checkpoint cannot map to it without +//! pre-splitting the weights at load time. +//! +//! This module keeps the fused layout: a single `qkv_proj` Linear of +//! shape `(d_model -> 3 * d_model)` and a `chunk(3, -1)` at forward, +//! giving a one-to-one mapping with the checkpoint. +//! +//! No attention mask is supported. CLIP's image encoder runs +//! unconditional self-attention; the text encoder (which uses a causal +//! mask) is not ported. + +use burn_core as burn; + +use burn::config::Config; +use burn::module::Module; +use burn::tensor::Tensor; +use burn::tensor::activation::softmax; +use burn::tensor::backend::Backend; +use burn_nn::{Linear, LinearConfig}; + +/// Configuration for [`ClipQkvAttention`]. +#[derive(Config, Debug)] +pub(crate) struct ClipQkvAttentionConfig { + /// Embedding dimension. Must be divisible by `n_heads`. + pub d_model: usize, + /// Number of attention heads. + pub n_heads: usize, +} + +impl ClipQkvAttentionConfig { + /// Initialize a [`ClipQkvAttention`] block. + pub(crate) fn init(&self, device: &B::Device) -> ClipQkvAttention { + assert_eq!( + self.d_model % self.n_heads, + 0, + "d_model ({}) must be divisible by n_heads ({})", + self.d_model, + self.n_heads + ); + let head_dim = self.d_model / self.n_heads; + ClipQkvAttention { + qkv_proj: LinearConfig::new(self.d_model, 3 * self.d_model) + .with_bias(true) + .init(device), + out_proj: LinearConfig::new(self.d_model, self.d_model) + .with_bias(true) + .init(device), + d_model: self.d_model, + n_heads: self.n_heads, + head_dim, + } + } +} + +/// Self-attention with a fused QKV projection, matching CLIP's checkpoint +/// layout one-to-one. +#[derive(Module, Debug)] +pub(crate) struct ClipQkvAttention { + /// Fused projection `d_model -> 3 * d_model` for Q, K, V. + pub(crate) qkv_proj: Linear, + /// Output projection `d_model -> d_model`. + pub(crate) out_proj: Linear, + /// Embedding dimension. + pub(crate) d_model: usize, + /// Number of attention heads. + pub(crate) n_heads: usize, + /// Per-head dimension (`d_model / n_heads`). + pub(crate) head_dim: usize, +} + +impl ClipQkvAttention { + /// Apply self-attention. Input and output shape: `[batch, seq, d_model]`. + pub(crate) fn forward(&self, x: Tensor) -> Tensor { + let [batch, seq, _] = x.dims(); + + let qkv = self.qkv_proj.forward(x); + let mut chunks = qkv.chunk(3, 2); + let value = chunks.remove(2); + let key = chunks.remove(1); + let query = chunks.remove(0); + + let to_heads = |t: Tensor| { + t.reshape([batch, seq, self.n_heads, self.head_dim]) + .swap_dims(1, 2) + }; + let query = to_heads(query); + let key = to_heads(key); + let value = to_heads(value); + + let scale = (self.head_dim as f32).sqrt(); + let scores = query.matmul(key.transpose()).div_scalar(scale); + let weights = softmax(scores, 3); + + let context = weights + .matmul(value) + .swap_dims(1, 2) + .reshape([batch, seq, self.d_model]); + self.out_proj.forward(context) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Distribution; + use burn_flex::Flex; + + type TestBackend = Flex; + + #[test] + fn clip_qkv_attention_preserves_shape() { + let device = Default::default(); + let attn = ClipQkvAttentionConfig::new(768, 12).init::(&device); + + let input = Tensor::::random([1, 50, 768], Distribution::Default, &device); + let output = attn.forward(input); + + assert_eq!(output.dims(), [1, 50, 768]); + } + + #[test] + fn clip_qkv_attention_handles_batch() { + let device = Default::default(); + let attn = ClipQkvAttentionConfig::new(64, 4).init::(&device); + + let input = Tensor::::random([3, 16, 64], Distribution::Default, &device); + let output = attn.forward(input); + + assert_eq!(output.dims(), [3, 16, 64]); + } +} diff --git a/crates/burn-train/src/metric/vision/afine/clip_vit.rs b/crates/burn-train/src/metric/vision/afine/clip_vit.rs new file mode 100644 index 00000000000..4f984fda562 --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/clip_vit.rs @@ -0,0 +1,367 @@ +//! CLIP ViT-B/32 image encoder. +//! +//! Structurally a port of OpenAI CLIP's `VisionTransformer` minus the +//! text-side and joint-embedding projection. A-FINE consumes the +//! per-layer patch features (twelve `[batch, num_patches, 768]` tensors), +//! not the final CLIP embedding, so [`Self::forward_with_features`] is +//! the entry point used by the metric. [`Self::forward`] returns just the +//! `ln_post`-normalized class-token embedding and is provided for +//! completeness. +//! +//! Layout: NLD (`[batch, seq, embed]`) end-to-end. The PyTorch reference +//! permutes to LND (`[seq, batch, embed]`) inside the transformer stack +//! and back again before `ln_post`; attention is invariant to that swap, +//! so we stay in NLD throughout. + +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Initializer, Module, Param}; +use burn::tensor::Tensor; +use burn::tensor::backend::Backend; +use burn_nn::conv::{Conv2d, Conv2dConfig}; +use burn_nn::{LayerNorm, LayerNormConfig, Linear, LinearConfig}; + +use super::clip_attention::{ClipQkvAttention, ClipQkvAttentionConfig}; +use super::quick_gelu::QuickGelu; + +/// Configuration for [`ClipVisualEncoder`]. +#[derive(Config, Debug)] +pub struct ClipVisualEncoderConfig { + /// Input image channels. Defaults to 3 (RGB). + #[config(default = "3")] + pub in_channels: usize, + /// Token embedding dimension. Defaults to 768 (CLIP ViT-B/32). + #[config(default = "768")] + pub embed_dim: usize, + /// Patch (and conv stride) size. Defaults to 32 (CLIP ViT-B/32). + #[config(default = "32")] + pub patch_size: usize, + /// Number of transformer blocks. Defaults to 12. + #[config(default = "12")] + pub num_layers: usize, + /// Number of attention heads per block. Defaults to 12. + #[config(default = "12")] + pub num_heads: usize, + /// Hidden dimension of the per-block MLP. Defaults to 3072 (4 * 768). + #[config(default = "3072")] + pub mlp_dim: usize, + /// Expected input resolution. Defaults to 256. Must be a multiple of `patch_size`. + #[config(default = "256")] + pub image_size: usize, +} + +impl ClipVisualEncoderConfig { + /// Initialize a [`ClipVisualEncoder`] with random weights. + pub fn init(&self, device: &B::Device) -> ClipVisualEncoder { + assert_eq!( + self.image_size % self.patch_size, + 0, + "image_size ({}) must be a multiple of patch_size ({})", + self.image_size, + self.patch_size + ); + + let num_patches = (self.image_size / self.patch_size).pow(2); + let seq_len = num_patches + 1; // +1 for the prepended CLS token. + + let patch_embed = Conv2dConfig::new( + [self.in_channels, self.embed_dim], + [self.patch_size, self.patch_size], + ) + .with_stride([self.patch_size, self.patch_size]) + .with_bias(false) + .init(device); + + // Standard transformer init magnitude. Final values get overwritten + // when pretrained weights load; this just keeps random-init forward + // passes well-conditioned for tests. + let init = Initializer::Normal { + mean: 0.0, + std: 0.02, + }; + let class_token = init.init([self.embed_dim], device); + let positional_embedding = init.init([seq_len, self.embed_dim], device); + + let blocks = (0..self.num_layers) + .map(|_| { + TransformerBlockConfig::new(self.embed_dim, self.num_heads, self.mlp_dim) + .init(device) + }) + .collect(); + + ClipVisualEncoder { + patch_embed, + class_token, + positional_embedding, + ln_pre: LayerNormConfig::new(self.embed_dim).init(device), + blocks, + ln_post: LayerNormConfig::new(self.embed_dim).init(device), + embed_dim: self.embed_dim, + patch_size: self.patch_size, + image_size: self.image_size, + } + } +} + +/// Output of [`ClipVisualEncoder::forward_with_features`]. +/// +/// `features` are the per-block patch tokens A-FINE consumes. `cls` is +/// the post-`ln_post` class-token embedding, present only when the +/// caller requests it; the metric path skips the slice + LayerNorm to +/// avoid wasted work. +#[derive(Debug)] +pub struct ClipOutput { + /// Per-block patch tokens, length `num_layers`, each + /// `[batch, num_patches, embed_dim]`. + pub features: Vec>, + /// Post-`ln_post` class-token embedding, `[batch, embed_dim]`. `Some` + /// when `forward_with_features` is called with `return_cls = true`. + pub cls: Option>, +} + +/// CLIP ViT-B/32 image encoder. +/// +/// Consumes a normalized RGB image and returns either the class-token +/// embedding ([`Self::forward`]) or the per-layer patch-token features +/// used by A-FINE ([`Self::forward_with_features`]). +#[derive(Module, Debug)] +pub struct ClipVisualEncoder { + /// Patch-embedding convolution: `(in_channels, embed_dim)`, kernel and + /// stride both `patch_size`, no bias (matches CLIP). + pub(crate) patch_embed: Conv2d, + /// Learnable class token, shape `[embed_dim]`. Stored 1-D to match the + /// CLIP checkpoint key `class_embedding`; reshaped to `[1, 1, embed]` + /// at forward time. + pub(crate) class_token: Param>, + /// Learnable positional embedding, shape `[seq_len, embed_dim]`. + pub(crate) positional_embedding: Param>, + /// Pre-stack LayerNorm. + pub(crate) ln_pre: LayerNorm, + /// Stack of transformer blocks. + pub(crate) blocks: Vec>, + /// Post-stack LayerNorm applied to the class token only. + pub(crate) ln_post: LayerNorm, + + pub(crate) embed_dim: usize, + pub(crate) patch_size: usize, + pub(crate) image_size: usize, +} + +impl ClipVisualEncoder { + /// Encode an image to its class-token embedding. + /// + /// # Shapes + /// - input: `[batch, in_channels, height, width]` + /// - output: `[batch, embed_dim]` + pub fn forward(&self, image: Tensor) -> Tensor { + self.forward_with_features(image, true) + .cls + .expect("cls requested") + } + + /// Encode an image and return the patch-token features after each + /// transformer block. The class token is stripped from each level. + /// The post-`ln_post` cls embedding is computed only when + /// `return_cls` is true; A-FINE itself does not use it. + /// + /// # Shapes + /// - input: `[batch, in_channels, height, width]` + /// - features: `Vec` of length `num_layers`, each + /// `[batch, num_patches, embed_dim]` + /// - cls (when present): `[batch, embed_dim]` + pub fn forward_with_features(&self, image: Tensor, return_cls: bool) -> ClipOutput { + let [batch, _, height, width] = image.dims(); + assert_eq!( + height % self.patch_size, + 0, + "image height ({}) must be a multiple of patch_size ({})", + height, + self.patch_size + ); + assert_eq!( + width % self.patch_size, + 0, + "image width ({}) must be a multiple of patch_size ({})", + width, + self.patch_size + ); + + let embed = self.embed_dim; + + // Patch-embed: [B, C, H, W] -> [B, embed, H/p, W/p]. + let x = self.patch_embed.forward(image); + let [_, _, h_out, w_out] = x.dims(); + let num_patches = h_out * w_out; + let seq_len = num_patches + 1; + + // Flatten spatial axes, then put the patch axis before the channel + // axis: [B, embed, N] -> [B, N, embed]. + let x = x.reshape([batch, embed, num_patches]).swap_dims(1, 2); + + // Prepend the CLS token. Reshape [embed] -> [1, 1, embed], then + // broadcast-expand to [B, 1, embed] before concatenating along the + // sequence axis. + let cls = self + .class_token + .val() + .reshape([1, 1, embed]) + .expand([batch, 1, embed]); + let x = Tensor::cat(vec![cls, x], 1); + + // Add positional embedding. The pos-embed param is initialised to + // match `seq_len` for the configured image_size; runtime mismatches + // are caught here. Bicubic resize of a checkpointed pos-embed (e.g. + // 50 -> 65 when loading a 224-trained checkpoint at 256) is handled + // in the weights loader, not here. + let pos_seq = self.positional_embedding.dims()[0]; + assert_eq!( + pos_seq, seq_len, + "positional_embedding length {} does not match runtime sequence length {}", + pos_seq, seq_len + ); + let pos = self.positional_embedding.val().reshape([1, seq_len, embed]); + let x = x + pos; + + let mut x = self.ln_pre.forward(x); + + // Run the transformer stack, collecting patch features (CLS + // dropped) after each block. A-FINE consumes all twelve. + let mut features = Vec::with_capacity(self.blocks.len()); + for block in &self.blocks { + x = block.forward(x); + let patch_features = x.clone().slice([0..batch, 1..seq_len, 0..embed]); + features.push(patch_features); + } + + // ln_post is applied only to the class-token output. Skip both + // the slice and the LayerNorm when the caller doesn't need it. + let cls = return_cls.then(|| { + let cls = x.slice([0..batch, 0..1, 0..embed]).reshape([batch, embed]); + self.ln_post.forward(cls) + }); + + ClipOutput { features, cls } + } +} + +/// Configuration for a single transformer block. +#[derive(Config, Debug)] +pub(crate) struct TransformerBlockConfig { + pub d_model: usize, + pub n_heads: usize, + pub mlp_dim: usize, +} + +impl TransformerBlockConfig { + pub(crate) fn init(&self, device: &B::Device) -> TransformerBlock { + TransformerBlock { + ln_1: LayerNormConfig::new(self.d_model).init(device), + attn: ClipQkvAttentionConfig::new(self.d_model, self.n_heads).init(device), + ln_2: LayerNormConfig::new(self.d_model).init(device), + mlp: TransformerMlpConfig::new(self.d_model, self.mlp_dim).init(device), + } + } +} + +/// Pre-norm transformer block: `x = x + attn(ln_1(x))`, then +/// `x = x + mlp(ln_2(x))`. Matches CLIP's encoder block exactly. +#[derive(Module, Debug)] +pub(crate) struct TransformerBlock { + pub(crate) ln_1: LayerNorm, + pub(crate) attn: ClipQkvAttention, + pub(crate) ln_2: LayerNorm, + pub(crate) mlp: TransformerMlp, +} + +impl TransformerBlock { + pub(crate) fn forward(&self, x: Tensor) -> Tensor { + let attn_out = self.attn.forward(self.ln_1.forward(x.clone())); + let x = x + attn_out; + let mlp_out = self.mlp.forward(self.ln_2.forward(x.clone())); + x + mlp_out + } +} + +/// Configuration for [`TransformerMlp`]. +#[derive(Config, Debug)] +pub(crate) struct TransformerMlpConfig { + pub d_model: usize, + pub mlp_dim: usize, +} + +impl TransformerMlpConfig { + pub(crate) fn init(&self, device: &B::Device) -> TransformerMlp { + TransformerMlp { + c_fc: LinearConfig::new(self.d_model, self.mlp_dim) + .with_bias(true) + .init(device), + activation: QuickGelu, + c_proj: LinearConfig::new(self.mlp_dim, self.d_model) + .with_bias(true) + .init(device), + } + } +} + +/// Two-layer MLP with [`QuickGelu`] in between. CLIP's `c_fc` and `c_proj` +/// names are preserved so the checkpoint maps without remapping. +#[derive(Module, Debug)] +pub(crate) struct TransformerMlp { + pub(crate) c_fc: Linear, + pub(crate) activation: QuickGelu, + pub(crate) c_proj: Linear, +} + +impl TransformerMlp { + pub(crate) fn forward(&self, x: Tensor) -> Tensor { + let x = self.c_fc.forward(x); + let x = self.activation.forward(x); + self.c_proj.forward(x) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Distribution; + use burn_flex::Flex; + + type TestBackend = Flex; + + #[test] + fn clip_visual_encoder_forward_shape() { + let device = Default::default(); + let encoder = ClipVisualEncoderConfig::new() + .with_image_size(64) + .with_num_layers(2) + .init::(&device); + + let image = + Tensor::::random([1, 3, 64, 64], Distribution::Default, &device); + let cls = encoder.forward(image); + + assert_eq!(cls.dims(), [1, 768]); + } + + #[test] + fn clip_visual_encoder_forward_with_features_shape() { + let device = Default::default(); + let encoder = ClipVisualEncoderConfig::new() + .with_image_size(64) + .with_num_layers(3) + .init::(&device); + + let image = + Tensor::::random([2, 3, 64, 64], Distribution::Default, &device); + let ClipOutput { features, cls } = encoder.forward_with_features(image, true); + let cls = cls.expect("cls requested"); + + assert_eq!(cls.dims(), [2, 768]); + assert_eq!(features.len(), 3); + // 64x64 image at patch_size=32 -> 2x2 = 4 patches. + for level in &features { + assert_eq!(level.dims(), [2, 4, 768]); + } + } +} diff --git a/crates/burn-train/src/metric/vision/afine/heads.rs b/crates/burn-train/src/metric/vision/afine/heads.rs new file mode 100644 index 00000000000..6c8e9d59e4a --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/heads.rs @@ -0,0 +1,383 @@ +//! Naturalness and fidelity heads. +//! +//! Both heads consume the 12 per-block CLIP feature maps plus the raw +//! RGB image (re-introduced as a "level 0" feature). The heads' `mean` +//! and `std` buffers are CLIP's normalization constants, used here as +//! the **inverse** of preprocessing — they de-normalize the input back +//! into pixel space before computing raw-RGB statistics. +//! +//! - [`AfineQHead`] — naturalness, single-image, returns `[B, 1]`. +//! - [`AfineDHead`] — fidelity, two-image, returns `[B, 1]`. +//! +//! The `chns` tuple has 13 entries: `[3, 768, 768, ..., 768]`. The +//! leading 3 is the raw image; the remaining twelve 768s are the CLIP +//! transformer outputs after each block. + +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Module, Param}; +use burn::tensor::Tensor; +use burn::tensor::activation::{relu, softplus}; +use burn::tensor::backend::Backend; +use burn_nn::{Gelu, Linear, LinearConfig}; + +/// CLIP RGB normalization mean (per-channel). Used as a de-normalization +/// constant inside the heads. +const CLIP_MEAN: [f32; 3] = [0.481_454_66, 0.457_827_5, 0.408_210_73]; +/// CLIP RGB normalization std. +const CLIP_STD: [f32; 3] = [0.268_629_54, 0.261_302_58, 0.275_777_11]; + +/// Per-level channel counts: raw RGB (3) plus 12 CLIP layers (768 each). +const CHNS: [usize; 13] = [ + 3, 768, 768, 768, 768, 768, 768, 768, 768, 768, 768, 768, 768, +]; +const NUM_LEVELS: usize = 13; + +/// Hidden dim of the shared per-CLIP-level projection (`proj_feat`). +const Q_HIDDEN_DIM: usize = 128; +/// Output dim of the first MLP layer in `proj_head`. +const Q_PROJ_HEAD_OUT: usize = 768; +/// Input dim of `proj_head`: 6 (raw, k=0) + 128 * 12 (projected, k=1..=12). +const Q_PROJ_HEAD_IN: usize = 6 + Q_HIDDEN_DIM * 12; + +/// Sum of `CHNS`. `alpha` and `beta` are length `D_CHNS_SUM` along the +/// last axis. +const D_CHNS_SUM: usize = 9219; + +/// Numerical-stability epsilon used in the SSIM-like ratios. **Must be +/// 1e-10**, not the larger 1e-6 that DISTS uses; the feature magnitudes +/// after CLIP + ReLU are unit-scale-ish, so DISTS's eps would overwhelm +/// the signal. +const EPS: f64 = 1e-10; + +fn build_clip_mean_std(device: &B::Device) -> (Tensor, Tensor) { + let mean = Tensor::from_floats( + [[[[CLIP_MEAN[0]]], [[CLIP_MEAN[1]]], [[CLIP_MEAN[2]]]]], + device, + ); + let std = Tensor::from_floats( + [[[[CLIP_STD[0]]], [[CLIP_STD[1]]], [[CLIP_STD[2]]]]], + device, + ); + (mean, std) +} + +/// Mean and biased variance over the token axis, returned as a 2-D +/// tensor `[B, 2 * channels]` (mean followed by var, both flattened). +fn level_mean_var(feat: Tensor) -> (Tensor, Tensor, Tensor) { + let mean = feat.clone().mean_dim(1); + let centered = feat - mean.clone(); + let var = centered.powi_scalar(2).mean_dim(1); + let descriptor = Tensor::cat( + vec![ + mean.clone().flatten::<2>(1, 2), + var.clone().flatten::<2>(1, 2), + ], + 1, + ); + (mean, var, descriptor) +} + +/// Configuration for [`AfineQHead`]. +#[derive(Config, Debug)] +pub struct AfineQHeadConfig {} + +impl AfineQHeadConfig { + /// Initialize the naturalness head with random weights. + pub fn init(&self, device: &B::Device) -> AfineQHead { + let (mean, std) = build_clip_mean_std(device); + AfineQHead { + mean, + std, + proj_feat: LinearConfig::new(2 * 768, Q_HIDDEN_DIM) + .with_bias(true) + .init(device), + proj_head_fc1: LinearConfig::new(Q_PROJ_HEAD_IN, Q_PROJ_HEAD_OUT) + .with_bias(true) + .init(device), + proj_head_fc2: LinearConfig::new(Q_PROJ_HEAD_OUT, 1) + .with_bias(true) + .init(device), + activation: Gelu::new(), + } + } +} + +/// A-FINE naturalness head. Per-level mean+variance over CLIP tokens, +/// shared projection on the 12 CLIP levels, then a small MLP to a +/// scalar. The level-0 raw-image descriptor (`[B, 6]`) is concatenated +/// directly without going through `proj_feat`. +/// +/// Activation is the erf-based [`burn_nn::Gelu`], not the QuickGELU used +/// elsewhere in this crate. PyIQA's reference is explicit on this. +#[derive(Module, Debug)] +pub struct AfineQHead { + pub(crate) mean: Tensor, + pub(crate) std: Tensor, + pub(crate) proj_feat: Linear, + pub(crate) proj_head_fc1: Linear, + pub(crate) proj_head_fc2: Linear, + pub(crate) activation: Gelu, +} + +impl AfineQHead { + /// Compute the naturalness score for one image plus its CLIP + /// features. + /// + /// # Shapes + /// - `image`: `[B, 3, H, W]`, **already CLIP-normalized**. + /// - `clip_features`: 12 levels of `[B, num_patches, 768]`. + /// - returns: `[B, 1]`. + pub fn forward(&self, image: Tensor, clip_features: &[Tensor]) -> Tensor { + assert_eq!( + clip_features.len(), + 12, + "AfineQHead expects 12 CLIP feature maps, got {}", + clip_features.len() + ); + let [batch, channels, height, width] = image.dims(); + + // De-normalize: x = x * std + mean. + let img = image * self.std.clone() + self.mean.clone(); + + // [B, 3, H, W] -> [B, H*W, 3]: token-major raw-image features. + let img_feat = img + .reshape([batch, channels, height * width]) + .swap_dims(1, 2); + + let mut level_descriptors: Vec> = Vec::with_capacity(NUM_LEVELS); + + // Level 0: raw image, no projection. + let (_, _, raw_descriptor) = level_mean_var(img_feat); + level_descriptors.push(raw_descriptor); + + // Levels 1..=12: ReLU'd CLIP features, then shared `proj_feat`. + for h in clip_features { + let activated = relu(h.clone()); + let (_, _, descriptor) = level_mean_var(activated); + level_descriptors.push(self.proj_feat.forward(descriptor)); + } + + let concat_all = Tensor::cat(level_descriptors, 1); + let hidden = self + .activation + .forward(self.proj_head_fc1.forward(concat_all)); + self.proj_head_fc2.forward(hidden) + } +} + +/// Configuration for [`AfineDHead`]. +#[derive(Config, Debug)] +pub struct AfineDHeadConfig {} + +impl AfineDHeadConfig { + /// Initialize the fidelity head with random weights. + pub fn init(&self, device: &B::Device) -> AfineDHead { + let (mean, std) = build_clip_mean_std(device); + // PyIQA initializes alpha/beta with `.normal_(0.1, 0.01)`. That + // distribution doesn't really matter under random init — values + // are always overwritten by the checkpoint at `init_pretrained` + // time — but we mirror the magnitude here so the random-init + // forward stays well-conditioned. + let alpha = Tensor::random( + [1, 1, D_CHNS_SUM], + burn::tensor::Distribution::Normal(0.1, 0.01), + device, + ); + let beta = Tensor::random( + [1, 1, D_CHNS_SUM], + burn::tensor::Distribution::Normal(0.1, 0.01), + device, + ); + AfineDHead { + mean, + std, + alpha: Param::from_tensor(alpha), + beta: Param::from_tensor(beta), + } + } +} + +/// A-FINE fidelity head. SSIM-like luminance and contrast statistics +/// across 13 levels, weighted by globally-normalized softplus(alpha) and +/// softplus(beta). +#[derive(Module, Debug)] +pub struct AfineDHead { + pub(crate) mean: Tensor, + pub(crate) std: Tensor, + /// Luminance weights, shape `[1, 1, sum(chns)] = [1, 1, 9219]`. + pub(crate) alpha: Param>, + /// Contrast weights, same shape as `alpha`. + pub(crate) beta: Param>, +} + +impl AfineDHead { + /// Compute the fidelity score between distorted and reference + /// images. + /// + /// # Shapes + /// - `distorted`, `reference`: `[B, 3, H, W]`, both CLIP-normalized. + /// - `feat_dis`, `feat_ref`: 12 levels of `[B, num_patches, 768]`. + /// - returns: `[B, 1]`. + pub fn forward( + &self, + distorted: Tensor, + reference: Tensor, + feat_dis: &[Tensor], + feat_ref: &[Tensor], + ) -> Tensor { + assert_eq!(feat_dis.len(), 12); + assert_eq!(feat_ref.len(), 12); + let [batch, channels, height, width] = distorted.dims(); + + // De-normalize both images and lay them out as token sequences. + let raw_x = (distorted * self.std.clone() + self.mean.clone()) + .reshape([batch, channels, height * width]) + .swap_dims(1, 2); + let raw_y = (reference * self.std.clone() + self.mean.clone()) + .reshape([batch, channels, height * width]) + .swap_dims(1, 2); + + // 13-entry feature lists: raw RGB at level 0, ReLU'd CLIP + // features at levels 1..=12. + let mut feat_x: Vec> = Vec::with_capacity(NUM_LEVELS); + let mut feat_y: Vec> = Vec::with_capacity(NUM_LEVELS); + feat_x.push(raw_x); + feat_y.push(raw_y); + for h in feat_dis { + feat_x.push(relu(h.clone())); + } + for h in feat_ref { + feat_y.push(relu(h.clone())); + } + + // Global softplus normalization. `alpha_/w_sum + beta_/w_sum` + // sums to ~1 across all 13 levels and both terms. We keep + // `w_sum` as a tensor and broadcast-divide so the autodiff graph + // stays intact for any future training-mode caller. + let alpha_sp = softplus(self.alpha.val(), 1.0); + let beta_sp = softplus(self.beta.val(), 1.0); + let w_sum = (alpha_sp.clone().sum() + beta_sp.clone().sum()) + .add_scalar(EPS) + .reshape([1, 1, 1]); + let alpha_norm = alpha_sp / w_sum.clone(); + let beta_norm = beta_sp / w_sum; + + let mut dist1: Option> = None; + let mut dist2: Option> = None; + let mut offset: usize = 0; + + for k in 0..NUM_LEVELS { + let cn = CHNS[k]; + let alpha_k = alpha_norm.clone().slice([0..1, 0..1, offset..offset + cn]); + let beta_k = beta_norm.clone().slice([0..1, 0..1, offset..offset + cn]); + + let xm = feat_x[k].clone().mean_dim(1); // [B, 1, cn] + let ym = feat_y[k].clone().mean_dim(1); + + // S1 (luminance): (2 * x_mean * y_mean + eps) + // / (x_mean^2 + y_mean^2 + eps). + let s1_num = (xm.clone() * ym.clone()).mul_scalar(2.0).add_scalar(EPS); + let s1_den = xm.clone().powi_scalar(2) + ym.clone().powi_scalar(2); + let s1 = s1_num / s1_den.add_scalar(EPS); + let term1 = (alpha_k * s1).sum_dim(2); // [B, 1, 1] + dist1 = Some(match dist1 { + None => term1, + Some(d) => d + term1, + }); + + // S2 (contrast/structure): + // var_x = mean((x - x_mean)^2) (biased), + // var_y = mean((y - y_mean)^2), + // cov_xy = mean(x*y) - x_mean * y_mean, + // S2 = (2 * cov_xy + eps) / (var_x + var_y + eps). + let xc = feat_x[k].clone() - xm.clone(); + let yc = feat_y[k].clone() - ym.clone(); + let x_var = xc.powi_scalar(2).mean_dim(1); + let y_var = yc.powi_scalar(2).mean_dim(1); + let xy_mean = (feat_x[k].clone() * feat_y[k].clone()).mean_dim(1); + let xy_cov = xy_mean - xm * ym; + let s2_num = xy_cov.mul_scalar(2.0).add_scalar(EPS); + let s2_den = (x_var + y_var).add_scalar(EPS); + let s2 = s2_num / s2_den; + let term2 = (beta_k * s2).sum_dim(2); + dist2 = Some(match dist2 { + None => term2, + Some(d) => d + term2, + }); + + offset += cn; + } + + let total = dist1.unwrap() + dist2.unwrap(); // [B, 1, 1] + let score = total.ones_like() - total; // 1 - (dist1 + dist2) + score.squeeze_dim::<2>(2) // [B, 1] + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Distribution; + use burn_flex::Flex; + + type TestBackend = Flex; + + #[test] + fn afine_q_head_forward_shape() { + let device = Default::default(); + let head = AfineQHeadConfig::new().init::(&device); + + let image = + Tensor::::random([2, 3, 64, 64], Distribution::Default, &device); + let features: Vec<_> = (0..12) + .map(|_| Tensor::::random([2, 4, 768], Distribution::Default, &device)) + .collect(); + + let out = head.forward(image, &features); + assert_eq!(out.dims(), [2, 1]); + } + + #[test] + fn afine_d_head_forward_shape() { + let device = Default::default(); + let head = AfineDHeadConfig::new().init::(&device); + + let dis = Tensor::::random([2, 3, 64, 64], Distribution::Default, &device); + let reference = + Tensor::::random([2, 3, 64, 64], Distribution::Default, &device); + let feat_dis: Vec<_> = (0..12) + .map(|_| Tensor::::random([2, 4, 768], Distribution::Default, &device)) + .collect(); + let feat_ref: Vec<_> = (0..12) + .map(|_| Tensor::::random([2, 4, 768], Distribution::Default, &device)) + .collect(); + + let out = head.forward(dis, reference, &feat_dis, &feat_ref); + assert_eq!(out.dims(), [2, 1]); + } + + #[test] + fn afine_d_head_identical_inputs_unit_score() { + // When dis == ref and CLIP features match, S1 and S2 should both + // be ≈ 1, so dist1 + dist2 ≈ 1 and the final score 1 - (...) ≈ 0. + // Random init alpha/beta sum to 1 so this is a property-test + // sanity check on the global normalization. + let device = Default::default(); + let head = AfineDHeadConfig::new().init::(&device); + + let image = + Tensor::::random([1, 3, 32, 32], Distribution::Default, &device); + let features: Vec<_> = (0..12) + .map(|_| Tensor::::random([1, 1, 768], Distribution::Default, &device)) + .collect(); + + let out = head.forward(image.clone(), image, &features.clone(), &features); + let value = out.into_data().to_vec::().unwrap()[0]; + assert!( + value.abs() < 0.1, + "fidelity head on identical inputs should yield ~0, got {value}" + ); + } +} diff --git a/crates/burn-train/src/metric/vision/afine/metric.rs b/crates/burn-train/src/metric/vision/afine/metric.rs new file mode 100644 index 00000000000..a8ab3ce4c1e --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/metric.rs @@ -0,0 +1,370 @@ +//! A-FINE metric: top-level module, configuration, end-to-end forward +//! pass, `ModuleDisplay`, and property tests. + +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Module, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::tensor::backend::Backend; + +use super::calibrators::{ + AfineAdapter, AfineAdapterConfig, FrCalibratorWithLimit, FrCalibratorWithLimitConfig, + NrCalibrator, NrCalibratorConfig, scale_finalscore, +}; +use super::clip_vit::{ClipVisualEncoder, ClipVisualEncoderConfig}; +use super::heads::{AfineDHead, AfineDHeadConfig, AfineQHead, AfineQHeadConfig}; + +/// CLIP RGB normalization mean (per-channel). Applied at the top of +/// `forward` when `normalize_input` is true. +const CLIP_MEAN: [f32; 3] = [0.481_454_66, 0.457_827_5, 0.408_210_73]; +/// CLIP RGB normalization std. +const CLIP_STD: [f32; 3] = [0.268_629_54, 0.261_302_58, 0.275_777_11]; + +/// Configuration for the A-FINE metric. +#[derive(Config, Debug)] +pub struct AfineConfig { + /// Expected input resolution. Must be a multiple of 32. Defaults to 256. + #[config(default = "256")] + pub image_size: usize, + + /// Apply CLIP RGB normalization inside the forward pass. Set to false + /// if the caller has already normalized the input. Defaults to true. + #[config(default = "true")] + pub normalize_input: bool, +} + +impl AfineConfig { + /// Initialize an A-FINE module with default (random) weights. + /// + /// All six learnable submodules are initialized fresh; useful for + /// shape/property tests but produces meaningless quality scores + /// until [`Self::init_pretrained`] is called. + pub fn init(&self, device: &B::Device) -> Afine { + assert_eq!( + self.image_size % 32, + 0, + "A-FINE image_size must be a multiple of 32, got {}", + self.image_size + ); + + Afine { + clip_visual: ClipVisualEncoderConfig::new() + .with_image_size(self.image_size) + .init(device), + qhead: AfineQHeadConfig::new().init(device), + dhead: AfineDHeadConfig::new().init(device), + nr_calibrator: NrCalibratorConfig::new().init(device), + fr_calibrator: FrCalibratorWithLimitConfig::new().init(device), + adapter: AfineAdapterConfig::new().init(device), + image_size: self.image_size, + normalize_input: self.normalize_input, + } + } + + /// Initialize an A-FINE module and load pretrained PyIQA weights. + /// + /// Downloads `afine.pth` (~600 MB) from the PyIQA Hugging Face + /// mirror on first call, caches it under + /// `~/.cache/burn-dataset/afine/`, and loads all six shards into the + /// matching submodules. + pub fn init_pretrained(&self, device: &B::Device) -> Afine { + let afine = self.init(device); + super::weights::load_pretrained_weights(afine) + } +} + +/// A-FINE (Adaptive Fidelity-Naturalness Evaluator) full-reference image +/// quality metric built on CLIP ViT-B/32 features. +/// +/// `forward(distorted, reference)` returns a per-sample quality score in +/// `(0, 100)` — higher means worse perceptual quality. The metric is +/// **not symmetric**: `forward(a, b) != forward(b, a)` in general, +/// because the adapter term `exp(k * (N_ref - N_dis))` weights the two +/// inputs differently. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Afine { + pub(crate) clip_visual: ClipVisualEncoder, + pub(crate) qhead: AfineQHead, + pub(crate) dhead: AfineDHead, + pub(crate) nr_calibrator: NrCalibrator, + pub(crate) fr_calibrator: FrCalibratorWithLimit, + pub(crate) adapter: AfineAdapter, + + pub(crate) image_size: usize, + pub(crate) normalize_input: bool, +} + +impl Afine { + /// Compute the A-FINE quality score. + /// + /// # Shapes + /// - `distorted`, `reference`: `[batch, 3, H, W]` with `H` and `W` + /// both multiples of 32. Values in `[0, 1]` (RGB) when + /// `normalize_input` is true; already-normalized when false. + /// - returns: `[batch]` per-sample score in `(0, 100)`. + pub fn forward(&self, distorted: Tensor, reference: Tensor) -> Tensor { + let [_, _, height, width] = distorted.dims(); + assert_eq!( + height % 32, + 0, + "A-FINE input height ({height}) must be a multiple of 32" + ); + assert_eq!( + width % 32, + 0, + "A-FINE input width ({width}) must be a multiple of 32" + ); + + let device = distorted.device(); + let (dis_norm, ref_norm) = if self.normalize_input { + let mean = Tensor::::from_floats( + [[[[CLIP_MEAN[0]]], [[CLIP_MEAN[1]]], [[CLIP_MEAN[2]]]]], + &device, + ); + let std = Tensor::::from_floats( + [[[[CLIP_STD[0]]], [[CLIP_STD[1]]], [[CLIP_STD[2]]]]], + &device, + ); + ( + (distorted - mean.clone()) / std.clone(), + (reference - mean) / std, + ) + } else { + (distorted, reference) + }; + + // CLIP gives us the 12 per-block patch-feature maps the heads + // consume; the class-token vector is unused by A-FINE so we ask + // the encoder to skip it. + let feat_dis = self + .clip_visual + .forward_with_features(dis_norm.clone(), false) + .features; + let feat_ref = self + .clip_visual + .forward_with_features(ref_norm.clone(), false) + .features; + + let natural_dis = self.qhead.forward(dis_norm.clone(), &feat_dis); + let natural_ref = self.qhead.forward(ref_norm.clone(), &feat_ref); + let natural_dis_scaled = self.nr_calibrator.forward(natural_dis); + let natural_ref_scaled = self.nr_calibrator.forward(natural_ref); + + let fidelity = self.dhead.forward(dis_norm, ref_norm, &feat_dis, &feat_ref); + let fidelity_scaled = self.fr_calibrator.forward(fidelity); + + let d = self + .adapter + .forward(natural_dis_scaled, natural_ref_scaled, fidelity_scaled); + let score = scale_finalscore(d); // [B, 1] + score.squeeze_dim::<1>(1) // [B] + } +} + +impl ModuleDisplay for Afine { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("backbone", &"CLIP ViT-B/32".to_string()) + .add("image_size", &self.image_size.to_string()) + .add("normalize_input", &self.normalize_input.to_string()) + .optional() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Distribution; + use burn_flex::Flex; + + type TestBackend = Flex; + + fn small_metric() -> Afine { + // A small image_size keeps tests fast — 12-layer CLIP at the + // default 256x256 takes a noticeable fraction of a second per + // forward. + let device = Default::default(); + AfineConfig::new() + .with_image_size(64) + .init::(&device) + } + + #[test] + fn afine_forward_shape() { + let device = Default::default(); + let metric = small_metric(); + let dis = Tensor::::random([2, 3, 64, 64], Distribution::Default, &device); + let reference = + Tensor::::random([2, 3, 64, 64], Distribution::Default, &device); + let score = metric.forward(dis, reference); + assert_eq!(score.dims(), [2]); + } + + #[test] + fn afine_batch_processing() { + let device = Default::default(); + let metric = small_metric(); + let dis = Tensor::::random([4, 3, 64, 64], Distribution::Default, &device); + let reference = + Tensor::::random([4, 3, 64, 64], Distribution::Default, &device); + let score = metric.forward(dis, reference); + let values = score.into_data().to_vec::().unwrap(); + assert_eq!(values.len(), 4); + for v in values { + assert!(v.is_finite(), "A-FINE produced non-finite value: {v}"); + } + } + + #[test] + fn afine_finite_on_constant_inputs() { + // Guards against the c1=c2=1e-10 epsilon failing to keep the + // SSIM ratios finite when feature variance is zero. + let device = Default::default(); + let metric = small_metric(); + let zeros = Tensor::::zeros([1, 3, 64, 64], &device); + let ones = Tensor::::ones([1, 3, 64, 64], &device); + let score = metric.forward(zeros, ones); + let value = score.into_data().to_vec::().unwrap()[0]; + assert!(value.is_finite(), "got non-finite value {value}"); + } + + #[test] + fn afine_asymmetry() { + // D(a, b) != D(b, a) — the adapter weights N_dis by an + // exponential of (N_ref - N_dis), so swapping inputs flips the + // sign in the exponent. + let device = Default::default(); + let metric = small_metric(); + let a = Tensor::::random([1, 3, 64, 64], Distribution::Default, &device); + let b = Tensor::::random([1, 3, 64, 64], Distribution::Default, &device); + let forward = metric + .forward(a.clone(), b.clone()) + .into_data() + .to_vec::() + .unwrap()[0]; + let reverse = metric.forward(b, a).into_data().to_vec::().unwrap()[0]; + assert!( + (forward - reverse).abs() > 1e-6, + "expected asymmetric output, got fwd={forward}, rev={reverse}" + ); + } + + #[test] + fn afine_image_size_must_be_multiple_of_32() { + let result = std::panic::catch_unwind(|| { + AfineConfig::new() + .with_image_size(33) + .init::(&Default::default()); + }); + assert!(result.is_err(), "expected init to panic on bad image_size"); + } + + #[test] + fn display_afine() { + let metric = small_metric(); + let formatted = format!("{metric}"); + assert!( + formatted.contains("CLIP ViT-B/32"), + "expected backbone name in display output: {formatted}" + ); + } + + #[test] + #[ignore = "downloads pre-trained weights"] + fn test_afine_pretrained_parity() { + // Numerical parity against PyIQA on a deterministic input pair. + // Catches silent QuickGELU-vs-GELU swaps, fused-QKV transpose- + // direction bugs, channel-order mistakes, `c1=c2=1e-10` epsilon + // drift, and any partially loaded shard — all of which the + // property tests miss. + // + // Expected value captured 2026-04-28 via + // `Notes/capture_afine_fixtures.py` (torch=2.11.0, + // pyiqa=0.1.15.post2). Input pair: `arange/(N-1)` vs `arange/N` + // reshaped to `[1, 3, 224, 224]`. Re-capture if the upstream + // PyIQA checkpoint or normalization conventions change. + const D_EXPECTED: f32 = 43.15711594_f32; + + let device = Default::default(); + let metric = AfineConfig::new() + .with_image_size(224) + .init_pretrained::(&device); + + let total: i64 = 1 * 3 * 224 * 224; + let dis = Tensor::::arange(0..total, &device) + .float() + .div_scalar((total - 1) as f64) + .reshape([1, 3, 224, 224]); + let reference = Tensor::::arange(0..total, &device) + .float() + .div_scalar(total as f64) + .reshape([1, 3, 224, 224]); + + let value = metric + .forward(dis, reference) + .into_data() + .to_vec::() + .unwrap()[0]; + + // Tolerance: relative 5e-3 covers fp32 attention drift across + // 12 transformer layers. Absolute floor of 5e-3 prevents a + // captured value near zero from forcing absurd precision. In + // practice the burn output matches the captured PyIQA scalar to + // within ~5e-5 absolute, well inside this budget. + let tolerance = D_EXPECTED.abs() * 5e-3 + 5e-3; + assert!( + (value - D_EXPECTED).abs() < tolerance, + "expected {D_EXPECTED}, got {value} (tolerance {tolerance})" + ); + } + + #[test] + #[ignore = "downloads pre-trained weights"] + fn test_afine_pretrained() { + // Loads the real PyIQA checkpoint (~600 MB on first run, cached + // afterwards) and verifies the forward pass produces a finite + // score that meaningfully differs from the random-init metric. + // The differs-from-random check guards against a silent + // `allow_partial(true)` skip — if a regex remap is wrong the + // weights stay at random init and the property tests still + // pass, but the pretrained output should land in a very + // different region of the score range. + let device = Default::default(); + let dis = + Tensor::::random([1, 3, 256, 256], Distribution::Default, &device); + let reference = + Tensor::::random([1, 3, 256, 256], Distribution::Default, &device); + + let random_metric = AfineConfig::new().init::(&device); + let random_score = random_metric + .forward(dis.clone(), reference.clone()) + .into_data() + .to_vec::() + .unwrap()[0]; + + let pretrained_metric = AfineConfig::new().init_pretrained::(&device); + let pretrained_score = pretrained_metric + .forward(dis, reference) + .into_data() + .to_vec::() + .unwrap()[0]; + + assert!( + pretrained_score.is_finite(), + "pretrained A-FINE produced non-finite output: {pretrained_score}" + ); + assert!( + (pretrained_score - random_score).abs() > 1e-3, + "pretrained output ({pretrained_score}) too close to random-init ({random_score}) — \ + a shard may have failed to load silently" + ); + } +} diff --git a/crates/burn-train/src/metric/vision/afine/mod.rs b/crates/burn-train/src/metric/vision/afine/mod.rs new file mode 100644 index 00000000000..dab543d1ffd --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/mod.rs @@ -0,0 +1,19 @@ +//! A-FINE (Adaptive Fidelity-Naturalness Evaluator) metric. +//! +//! Full-reference perceptual image quality metric that combines a naturalness +//! branch and a fidelity branch over CLIP ViT-B/32 features. +//! +//! Reference: "A Novel Fidelity-Naturalness Evaluator for Image Quality Assessment" +//! + +mod calibrators; +mod clip_attention; +mod clip_vit; +mod heads; +mod metric; +mod quick_gelu; +mod weights; + +pub use clip_vit::{ClipOutput, ClipVisualEncoder, ClipVisualEncoderConfig}; +pub use heads::{AfineDHead, AfineDHeadConfig, AfineQHead, AfineQHeadConfig}; +pub use metric::{Afine, AfineConfig}; diff --git a/crates/burn-train/src/metric/vision/afine/quick_gelu.rs b/crates/burn-train/src/metric/vision/afine/quick_gelu.rs new file mode 100644 index 00000000000..601a5636da3 --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/quick_gelu.rs @@ -0,0 +1,71 @@ +//! QuickGELU activation used by OpenAI's CLIP. +//! +//! Defined as `x * sigmoid(1.702 * x)`. Distinct from the erf-based +//! [`burn_nn::Gelu`]: CLIP was trained with QuickGELU and substituting +//! standard GELU silently shifts every transformer-block MLP output. + +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Tensor; +use burn::tensor::activation::sigmoid; +use burn::tensor::backend::Backend; + +/// Approximation coefficient. Must be exactly 1.702; CLIP weights are +/// trained with this value and changing it breaks numerical parity. +const COEFFICIENT: f64 = 1.702; + +/// QuickGELU activation: `x * sigmoid(1.702 * x)`. +#[derive(Module, Clone, Debug, Default)] +pub(crate) struct QuickGelu; + +impl QuickGelu { + /// Apply the activation element-wise. Shape-preserving. + pub(crate) fn forward(&self, x: Tensor) -> Tensor { + let scaled = x.clone().mul_scalar(COEFFICIENT); + x * sigmoid(scaled) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn_core::tensor::{Tolerance, ops::FloatElem}; + use burn_flex::Flex; + + type TestBackend = Flex; + type FT = FloatElem; + + #[test] + fn quick_gelu_matches_formula() { + let device = Default::default(); + let activation = QuickGelu; + + let input = Tensor::::from_floats([[-1.0, 0.0, 1.0, 2.0]], &device); + let output = activation.forward(input); + + // Hand-computed: x * sigmoid(1.702 * x). + // sigmoid(-1.702) = 1 / (1 + exp(1.702)) ≈ 0.154160 + // sigmoid( 1.702) ≈ 0.845840 + // sigmoid( 3.404) ≈ 0.967812 + let expected = Tensor::::from_floats( + [[-0.154_160, 0.0, 0.845_840, 1.935_624]], + &device, + ); + + output + .into_data() + .assert_approx_eq::(&expected.into_data(), Tolerance::default()); + } + + #[test] + fn quick_gelu_preserves_shape() { + let device = Default::default(); + let activation = QuickGelu; + + let input = Tensor::::zeros([2, 5, 8], &device); + let output = activation.forward(input); + + assert_eq!(output.dims(), [2, 5, 8]); + } +} diff --git a/crates/burn-train/src/metric/vision/afine/weights.rs b/crates/burn-train/src/metric/vision/afine/weights.rs new file mode 100644 index 00000000000..70aa7c7c1e8 --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/weights.rs @@ -0,0 +1,264 @@ +//! Pretrained weights loading for A-FINE. +//! +//! `afine.pth` is published by the PyIQA project on Hugging Face and +//! bundles all six A-FINE shards into a single file: a fine-tuned CLIP +//! ViT-B/32 (`finetuned_clip`), the naturalness head (`natural`), the +//! fidelity head (`fidelity`), the two per-head logistic calibrators +//! (`natural_scale`, `fidelity_scale`), and the adapter (`adapter`). +//! Each shard is loaded into the corresponding submodule of [`Afine`] +//! via a separate [`PytorchStore`] with `with_top_level_key`. +//! +//! The OpenAI CLIP `ViT-B-32.pt` checkpoint is **not** used. That file +//! is a TorchScript archive that `burn-store` cannot read, and PyIQA's +//! `finetuned_clip` shard is the variant A-FINE was actually trained +//! with — substituting stock CLIP would silently produce wrong scores. +//! +//! ## Hosting +//! +//! The URL points at the PyIQA author's personal HF account, which is +//! the canonical source for these weights. There is no organization- +//! hosted mirror anywhere (HF orgs, Zenodo, figshare, ModelScope, OSF, +//! Kaggle, GitHub releases — all checked). The hosting question is +//! flagged in the PR description for the maintainer; if a different +//! mirror is preferred, swap the URL constant below. + +use burn_core as burn; + +use burn::module::Param; +use burn::tensor::Tensor; +use burn::tensor::backend::Backend; +use burn_std::network::downloader::download_file_as_bytes; +use burn_store::pytorch::PytorchReader; +use burn_store::{ModuleSnapshot, PytorchStore}; +use std::fs::{File, create_dir_all}; +use std::io::Write; +use std::path::Path; +use std::path::PathBuf; + +use super::calibrators::{AfineAdapter, FrCalibratorWithLimit, NrCalibrator}; +use super::metric::Afine; + +/// Source URL for `afine.pth` on Hugging Face. +const AFINE_URL: &str = + "https://huggingface.co/chaofengc/IQA-PyTorch-Weights/resolve/main/afine.pth"; + +/// Cached filename inside `~/.cache/burn-dataset/afine/`. +const CACHE_FILENAME: &str = "afine.pth"; + +fn get_cache_dir() -> PathBuf { + let cache_dir = dirs::cache_dir() + .expect("Could not get cache directory") + .join("burn-dataset") + .join("afine"); + if !cache_dir.exists() { + create_dir_all(&cache_dir).expect("Failed to create cache directory"); + } + cache_dir +} + +fn download_if_needed(url: &str, cache_path: &PathBuf, message: &str) { + if !cache_path.exists() { + let bytes = download_file_as_bytes(url, message); + let mut file = File::create(cache_path).expect("Failed to create cache file"); + file.write_all(&bytes).expect("Failed to write weights"); + } +} + +/// Download `afine.pth` (if not cached) and load all six shards into +/// the matching submodules of an `Afine` previously produced by +/// `AfineConfig::init`. +/// +/// Errors during loading are logged via `log::warn!` and the function +/// returns the module unconditionally — `allow_partial(true)` plus +/// per-shard regex remapping mean unmapped or unknown checkpoint keys +/// are silently dropped, matching the behaviour of LPIPS, DISTS, and +/// FID's pretrained loaders. +pub(crate) fn load_pretrained_weights(mut afine: Afine) -> Afine { + let cache_dir = get_cache_dir(); + let cache_path = cache_dir.join(CACHE_FILENAME); + download_if_needed(AFINE_URL, &cache_path, "Downloading A-FINE weights..."); + + afine.clip_visual = load_clip_shard(afine.clip_visual, &cache_path); + afine.qhead = load_qhead_shard(afine.qhead, &cache_path); + afine.dhead = load_simple_shard(afine.dhead, &cache_path, "fidelity", "fidelity head"); + + // The calibrator and adapter shards store yita3, yita4, and k as + // 0-D scalar tensors (`shape=()`). Burn's `Param>` of + // length 1 expects shape `(1,)`, and PytorchStore drops keys whose + // shape doesn't match the destination — silently leaving the + // params at random init. Load these manually via PytorchReader so + // the 0-D scalars become 1-element 1-D tensors. + let device = afine.adapter.k.val().device(); + afine.nr_calibrator = load_nr_calibrator_scalars(&cache_path, &device); + afine.fr_calibrator = load_fr_calibrator_scalars(&cache_path, &device); + afine.adapter = load_adapter_scalar(&cache_path, &device); + + afine +} + +/// Read a 0-D scalar value from a checkpoint shard, returning it as an +/// `f32`. Logs a warning and returns `None` on failure so the caller can +/// fall back to a default. +fn read_scalar_value>( + cache_path: P, + top_level_key: &str, + field_name: &str, +) -> Option { + let reader = PytorchReader::with_top_level_key(cache_path.as_ref(), top_level_key) + .map_err(|e| log::warn!("Failed to open shard '{top_level_key}': {e:?}")) + .ok()?; + let snapshot = reader.get(field_name)?; + let data = snapshot + .to_data() + .map_err(|e| log::warn!("Failed to read '{top_level_key}.{field_name}' tensor data: {e:?}")) + .ok()?; + let values = data + .to_vec::() + .map_err(|e| log::warn!("Failed to convert '{top_level_key}.{field_name}' to f32: {e:?}")) + .ok()?; + values.first().copied() +} + +fn scalar_param(value: f32, device: &B::Device) -> Param> { + Param::from_tensor(Tensor::from_floats([value], device)) +} + +fn load_nr_calibrator_scalars( + cache_path: &PathBuf, + device: &B::Device, +) -> NrCalibrator { + // PyIQA defaults if reading fails — preserves random-init fallback. + const NR_YITA3_FALLBACK: f32 = 4.9592; + const NR_YITA4_FALLBACK: f32 = 21.5968; + let yita3 = + read_scalar_value(cache_path, "natural_scale", "yita3").unwrap_or(NR_YITA3_FALLBACK); + let yita4 = + read_scalar_value(cache_path, "natural_scale", "yita4").unwrap_or(NR_YITA4_FALLBACK); + NrCalibrator { + yita3: scalar_param(yita3, device), + yita4: scalar_param(yita4, device), + } +} + +fn load_fr_calibrator_scalars( + cache_path: &PathBuf, + device: &B::Device, +) -> FrCalibratorWithLimit { + const FR_YITA3_FALLBACK: f32 = 0.5; + const FR_YITA4_FALLBACK: f32 = 0.15; + let yita3 = + read_scalar_value(cache_path, "fidelity_scale", "yita3").unwrap_or(FR_YITA3_FALLBACK); + let yita4 = + read_scalar_value(cache_path, "fidelity_scale", "yita4").unwrap_or(FR_YITA4_FALLBACK); + FrCalibratorWithLimit { + yita3: scalar_param(yita3, device), + yita4: scalar_param(yita4, device), + } +} + +fn load_adapter_scalar(cache_path: &PathBuf, device: &B::Device) -> AfineAdapter { + const K_FALLBACK: f32 = 5.0; + let k = read_scalar_value(cache_path, "adapter", "k").unwrap_or(K_FALLBACK); + AfineAdapter { + k: scalar_param(k, device), + } +} + +/// Load the `finetuned_clip` shard into the visual encoder. +/// +/// Key-remap rules strip the `visual.` prefix used in the CLIP state +/// dict and rename the few fields where the burn module diverges from +/// PyTorch's CLIP layout: `conv1` → `patch_embed`, +/// `class_embedding` → `class_token`, `transformer.resblocks` → `blocks`, +/// and the fused `attn.in_proj_*` → `attn.qkv_proj.*`. The +/// `weight`/`bias` ↔ `gamma`/`beta` rename for LayerNorm is handled +/// automatically by `PyTorchToBurnAdapter`. +/// +/// `with_regex(r"^visual\.")` filters the load to only the visual +/// encoder; the CLIP text encoder (~150 keys) is dropped at the filter +/// stage. `visual.proj` (the unused joint-embedding projection) has +/// no corresponding burn field and is silently dropped via +/// `allow_partial(true)`. +fn load_clip_shard( + mut clip: super::clip_vit::ClipVisualEncoder, + cache_path: &PathBuf, +) -> super::clip_vit::ClipVisualEncoder { + let mut store = PytorchStore::from_file(cache_path) + .with_top_level_key("finetuned_clip") + .allow_partial(true) + .skip_enum_variants(true) + // No `with_regex` filter: PathFilter matches against burn-side + // destination paths, not the source PyTorch keys, so a regex on + // `^visual\.` rejects every remapped key. Text-encoder keys + // (token_embedding, transformer.resblocks.*) have no + // corresponding burn fields and are dropped by + // `allow_partial(true)`, which is the same end result. + // + // Special case: the text encoder also has a top-level + // `positional_embedding` key (shape `[77, 512]`) which collides + // with the visual one after our `^visual\.positional_embedding$` + // rename targets `positional_embedding`. Without this first + // rule, HashMap iteration order decides which one wins, giving + // a flaky load. Rename the text encoder's first to a key that + // matches no burn field. + .with_key_remapping(r"^positional_embedding$", "_text_positional_embedding_drop") + .with_key_remapping(r"^visual\.conv1\.", "patch_embed.") + .with_key_remapping(r"^visual\.class_embedding$", "class_token") + .with_key_remapping(r"^visual\.positional_embedding$", "positional_embedding") + .with_key_remapping(r"^visual\.ln_pre\.", "ln_pre.") + .with_key_remapping(r"^visual\.ln_post\.", "ln_post.") + .with_key_remapping(r"^visual\.transformer\.resblocks\.", "blocks.") + .with_key_remapping(r"\.attn\.in_proj_weight$", ".attn.qkv_proj.weight") + .with_key_remapping(r"\.attn\.in_proj_bias$", ".attn.qkv_proj.bias"); + if let Err(e) = clip.load_from(&mut store) { + log::warn!( + "Some CLIP visual encoder weights could not be loaded: {:?}", + e + ); + } + clip +} + +/// Load the `natural` shard into the naturalness head. +/// +/// PyIQA stores the proj_head as a `nn.Sequential` with index-2 GELU, +/// so the FC layers come out as `proj_head.0.*` and `proj_head.2.*`. +/// We named them explicitly to avoid the GELU-index gap, hence the +/// remap. +fn load_qhead_shard( + mut qhead: super::heads::AfineQHead, + cache_path: &PathBuf, +) -> super::heads::AfineQHead { + let mut store = PytorchStore::from_file(cache_path) + .with_top_level_key("natural") + .allow_partial(true) + .skip_enum_variants(true) + .with_key_remapping(r"^proj_head\.0\.", "proj_head_fc1.") + .with_key_remapping(r"^proj_head\.2\.", "proj_head_fc2."); + if let Err(e) = qhead.load_from(&mut store) { + log::warn!("Some naturalness head weights could not be loaded: {:?}", e); + } + qhead +} + +/// Load any shard whose keys map directly to burn-side field names. +/// Used for the fidelity head, both calibrators, and the adapter. +fn load_simple_shard( + mut module: M, + cache_path: &PathBuf, + top_level_key: &'static str, + description: &str, +) -> M +where + B: Backend, + M: ModuleSnapshot, +{ + let mut store = PytorchStore::from_file(cache_path) + .with_top_level_key(top_level_key) + .allow_partial(true) + .skip_enum_variants(true); + if let Err(e) = module.load_from(&mut store) { + log::warn!("Some {} weights could not be loaded: {:?}", description, e); + } + module +} diff --git a/crates/burn-train/src/metric/vision/mod.rs b/crates/burn-train/src/metric/vision/mod.rs index 9af07620f36..9fbd4253f73 100644 --- a/crates/burn-train/src/metric/vision/mod.rs +++ b/crates/burn-train/src/metric/vision/mod.rs @@ -1,3 +1,4 @@ +mod afine; mod dice; mod dists; mod fid; @@ -6,6 +7,7 @@ mod ms_ssim; mod psnr; mod ssim; +pub use afine::*; pub use dice::*; pub use dists::*; pub use fid::*;