environment/training-source/source-snapshot/src/decoder.rs

//! Reconstructing RGB from DINOv3 patch features.
//!
//! This is the roundtrip the project is named for: encode an image to
//! features, decode those features back to pixels, and look at what
//! survived.
//!
//! # What to expect
//!
//! At 224² the encoder produces a 14×14 grid of 384-dimensional features —
//! about 75k numbers standing in for a 150k-value image. That alone would
//! permit a decent reconstruction, but DINOv3 features are trained for
//! *semantic invariance*: two views of the same object should land in the
//! same place regardless of colour, lighting, or fine texture, which means
//! precisely that information is discarded on purpose. So the reconstruction
//! recovers layout and dominant colour well and fine detail poorly. That is
//! the interesting result, not a defect in the decoder.
//!
//! # Shape
//!
//! A plain convolutional upsampler, four stages of ×2 from 14 to 224:
//!
//! ```text
//! [384, 14, 14] → 256@28 → 128@56 → 64@112 → 32@224 → 3@224
//! ```
//!
//! 2.01M parameters and about 1.14 GMAC per image. The early two stages use
//! a second convolution to blend patch seams; the later stages do not.
//!
//! No skip connections and nothing fancy: the point is to measure what the
//! features carry, and a decoder with its own access to the input would
//! confuse that question entirely.

use meganeura::{Graph, NodeId};

use crate::dinov3::Config;

/// Channel widths after each upsampling stage, from the feature grid up.
pub const STAGES: [usize; 4] = [256, 128, 64, 32];
/// How many of the early stages get a second convolution. Seams originate
/// at the patch grid, so blending earns its cost at low resolution and
/// mostly smooths detail away at high resolution.
pub const BLEND_STAGES: usize = 2;
/// Channels per group in the group norms.
const GROUP_SIZE: usize = 16;
const EPS: f32 = 1e-5;

/// Total parameter count, for reporting and for sizing a weights file.
pub fn parameter_count(config: &Config) -> usize {
    let mut total = 0;
    let mut in_c = config.hidden_size;
    for (i, &out_c) in STAGES.iter().enumerate() {
        total += out_c * in_c * 9 + out_c + out_c * 2;
        if i < BLEND_STAGES {
            total += out_c * out_c * 9 + out_c + out_c * 2;
        }
        in_c = out_c;
    }
    total += 3 * in_c * 9 + 3; // final projection to RGB
    total
}

/// Multiply-accumulates in the decoder's convolution kernels at batch one.
/// Bias, normalization, activations, and upsampling are intentionally not
/// folded into this number; timing still includes every operation.
pub fn forward_macs(config: &Config) -> u64 {
    let mut total = 0u64;
    let mut in_c = config.hidden_size;
    let mut hw = config.grid();
    for (i, &out_c) in STAGES.iter().enumerate() {
        total += (out_c * in_c * 9 * hw * hw) as u64;
        if i < BLEND_STAGES {
            total += (out_c * out_c * 9 * hw * hw) as u64;
        }
        hw *= 2;
        in_c = out_c;
    }
    total + (3 * in_c * 9 * hw * hw) as u64
}

/// One `conv3x3 → bias → group-norm → SiLU` block at a fixed resolution.
fn block(
    g: &mut Graph,
    x: NodeId,
    name: &str,
    batch: usize,
    in_c: usize,
    out_c: usize,
    hw: usize,
) -> NodeId {
    let kernel = g.parameter(&format!("{name}.weight"), &[out_c, in_c, 3, 3]);
    let x = g.conv2d(
        x,
        kernel,
        batch as u32,
        in_c as u32,
        hw as u32,
        hw as u32,
        out_c as u32,
        3,
        3,
        1,
        1,
    );
    let bias = g.parameter(&format!("{name}.bias"), &[out_c]);
    let x = g.add_per_channel(x, bias, out_c as u32, (hw * hw) as u32);

    let gn_w = g.parameter(&format!("{name}.norm.weight"), &[out_c]);
    let gn_b = g.parameter(&format!("{name}.norm.bias"), &[out_c]);
    let x = g.group_norm(
        x,
        gn_w,
        gn_b,
        batch as u32,
        out_c as u32,
        (hw * hw) as u32,
        (out_c / GROUP_SIZE) as u32,
        EPS,
    );
    g.silu(x)
}

/// Build the decoder over an existing feature node.
///
/// `features` must be `[batch, hidden, grid, grid]` flattened NCHW — patch
/// tokens only, transposed out of the encoder's token-major layout. See
/// [`patch_features_to_nchw`].
///
/// Returns `[batch, 3, image_size, image_size]` in `[0, 1]`.
pub fn build_decoder(g: &mut Graph, config: &Config, features: NodeId, batch: usize) -> NodeId {
    let mut x = features;
    let mut in_c = config.hidden_size;
    let mut hw = config.grid();

    for (i, &out_c) in STAGES.iter().enumerate() {
        x = block(g, x, &format!("dec.{i}"), batch, in_c, out_c, hw);
        // A second convolution, but only where it pays for itself.
        //
        // The seams come from the input being piecewise constant — one
        // feature vector per 16×16 patch — so hiding them needs a receptive
        // field wide enough to mix across cells, and one 3×3 per scale
        // reaches only a single neighbour. Doing that at *every* scale
        // removed the tiling but doubled the decoder, and the decoder is
        // ~40% of the frame on device.
        //
        // Each stage halves the channels and doubles the resolution, so all
        // four second convolutions cost the same. The early ones are where
        // patch boundaries actually live; by 56² and beyond the grid has
        // already been blended and the extra pass mostly smooths detail
        // away. Keeping the first two buys the seam repair for half the
        // price.
        if i < BLEND_STAGES {
            x = block(g, x, &format!("dec.{i}b"), batch, out_c, out_c, hw);
        }
        x = g.upsample_2x(x, batch as u32, out_c as u32, hw as u32, hw as u32);
        hw *= 2;
        in_c = out_c;
    }
    assert_eq!(
        hw, config.image_size,
        "stage count does not reach the image resolution"
    );

    let kernel = g.parameter("dec.out.weight", &[3, in_c, 3, 3]);
    let x = g.conv2d(
        x,
        kernel,
        batch as u32,
        in_c as u32,
        hw as u32,
        hw as u32,
        3,
        3,
        3,
        1,
        1,
    );
    let bias = g.parameter("dec.out.bias", &[3]);
    let x = g.add_per_channel(x, bias, 3, (hw * hw) as u32);
    // Sigmoid rather than a clamp: pixels live in [0, 1] and a hard clamp
    // has zero gradient outside the range, which strands any channel that
    // starts saturated.
    g.sigmoid(x)
}

/// Attach the decoder directly to a live encoder output.
///
/// Does in the graph what [`patch_features_to_nchw`] does on the CPU: drop
/// the prefix tokens and transpose from the encoder's token-major layout to
/// the channel-major one convolution needs. Batch 1 only — that is what
/// inference runs.
pub fn attach_to_encoder(g: &mut Graph, config: &Config, encoder_out: NodeId) -> NodeId {
    let hidden = config.hidden_size;
    let patches = config.num_patches();
    let prefix = config.num_prefix_tokens();

    // Slice off CLS and the register tokens. They carry global rather than
    // spatial information and have no place on a feature map.
    let patch_tokens = g.split_b(
        encoder_out,
        1,
        (prefix * hidden) as u32,
        (patches * hidden) as u32,
        1,
    );
    let patch_tokens = g.reshape(patch_tokens, &[patches, hidden]);
    let planes = g.transpose(patch_tokens);
    let planes = g.reshape(planes, &[hidden * patches]);
    build_decoder(g, config, planes, 1)
}

/// Load parameters written by `examples/train_decoder.rs`.
///
/// The file is raw f32 in graph declaration order, so it is only valid for
/// the graph shape that produced it; a length mismatch means the decoder
/// architecture changed and the weights are stale.
pub fn load_parameters(
    session: &mut meganeura::Session,
    graph: &Graph,
    path: &std::path::Path,
) -> Result<(), Box<dyn std::error::Error>> {
    let bytes = std::fs::read(path)?;
    let values: Vec<f32> = bytes
        .chunks_exact(4)
        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
        .collect();

    let mut offset = 0;
    let mut loaded = 0;
    for node in graph.nodes() {
        let meganeura::graph::Op::Parameter { name } = &node.op else {
            continue;
        };
        // Only decoder parameters live in this file; the encoder's come
        // from the checkpoint.
        if !name.starts_with("dec.") {
            continue;
        }
        let n = node.ty.num_elements();
        if offset + n > values.len() {
            return Err(format!(
                "{} is too short: needed {} values by parameter '{name}', have {}",
                path.display(),
                offset + n,
                values.len()
            )
            .into());
        }
        session.set_parameter(name, &values[offset..offset + n]);
        offset += n;
        loaded += 1;
    }
    if offset != values.len() {
        return Err(format!(
            "{} has {} values but the graph consumed {offset}; the decoder \
             architecture and the weights disagree",
            path.display(),
            values.len()
        )
        .into());
    }
    log::info!("loaded {loaded} decoder parameters from {}", path.display());
    Ok(())
}

/// Rearrange the encoder's `[tokens, hidden]` output into the `[hidden,
/// grid, grid]` NCHW block the decoder consumes, dropping the CLS and
/// register tokens.
///
/// The encoder is token-major (each row one token); convolution wants
/// channel-major. Done on the CPU here because it happens once per image
/// during dataset preparation; the live path folds it into the graph.
pub fn patch_features_to_nchw(features: &[f32], config: &Config) -> Vec<f32> {
    let hidden = config.hidden_size;
    let patches = config.num_patches();
    let skip = config.num_prefix_tokens();
    assert_eq!(features.len(), config.num_tokens() * hidden);

    let mut out = vec![0.0f32; hidden * patches];
    for p in 0..patches {
        let src = (skip + p) * hidden;
        for c in 0..hidden {
            out[c * patches + p] = features[src + c];
        }
    }
    out
}

/// Peak signal-to-noise ratio in dB between two `[0, 1]` images.
///
/// The number to quote for reconstruction quality. Roughly: below 15 dB is
/// unrecognizable, 20 dB is a recognizable blur, 30 dB is visually close.
pub fn psnr(a: &[f32], b: &[f32]) -> f32 {
    assert_eq!(a.len(), b.len());
    let mse: f64 = a
        .iter()
        .zip(b)
        .map(|(x, y)| {
            let d = (x - y) as f64;
            d * d
        })
        .sum::<f64>()
        / a.len() as f64;
    if mse <= f64::EPSILON {
        return f32::INFINITY;
    }
    (10.0 * (1.0 / mse).log10()) as f32
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn stages_reach_the_image_resolution() {
        let c = Config::vits16();
        assert_eq!(c.grid() * 2usize.pow(STAGES.len() as u32), c.image_size);
    }

    #[test]
    fn every_stage_divides_into_groups() {
        for &c in &STAGES {
            assert_eq!(c % GROUP_SIZE, 0, "{c} channels do not group evenly");
        }
    }

    #[test]
    fn parameter_count_is_modest() {
        let n = parameter_count(&Config::vits16());
        assert!(
            (1_000_000..4_000_000).contains(&n),
            "unexpected decoder size: {n}"
        );
    }

    #[test]
    fn forward_macs_match_the_deployed_decoder() {
        let c = Config::vits16().at_resolution(224).with_layers(3);
        assert_eq!(forward_macs(&c), 1_141_604_352);
    }

    #[test]
    fn graph_builds_with_the_right_output_shape() {
        let c = Config::vits16();
        let mut g = Graph::new();
        let feat = g.input("feat", &[c.hidden_size * c.num_patches()]);
        let out = build_decoder(&mut g, &c, feat, 1);
        assert_eq!(
            g.node(out).ty.num_elements(),
            3 * c.image_size * c.image_size
        );
    }

    /// The NCHW rearrangement must move each token's channel `c` to plane
    /// `c` at the token's grid position, and must skip the prefix tokens.
    #[test]
    fn nchw_rearrangement_is_a_transpose_past_the_prefix() {
        let c = Config::vits16();
        let h = c.hidden_size;
        // Encode each value as token*1000 + channel so the mapping is
        // checkable by arithmetic.
        let features: Vec<f32> = (0..c.num_tokens() * h)
            .map(|i| ((i / h) * 1000 + (i % h)) as f32)
            .collect();
        let nchw = patch_features_to_nchw(&features, &c);
        assert_eq!(nchw.len(), h * c.num_patches());

        for &(p, ch) in &[(0usize, 0usize), (37, 5), (195, 383)] {
            let token = c.num_prefix_tokens() + p;
            assert_eq!(
                nchw[ch * c.num_patches() + p],
                (token * 1000 + ch) as f32,
                "patch {p} channel {ch} came from the wrong token"
            );
        }
    }

    #[test]
    fn psnr_behaves() {
        let a = vec![0.5f32; 100];
        assert!(psnr(&a, &a).is_infinite(), "identical images are lossless");
        // A uniform 0.1 error is MSE 0.01, i.e. 20 dB.
        let b: Vec<f32> = a.iter().map(|v| v + 0.1).collect();
        assert!((psnr(&a, &b) - 20.0).abs() < 0.1);
    }
}