pagesift_ocr/layout/
decoder.rs

1//! 解码器（`PPDocLayoutV3Decoder`）：6 层（自注意力 + 可变形交叉注意力 + 前馈），每层之后用框头迭代
2//! 修正参考框。类别头、框头与编码器的 `enc_score_head`、`enc_bbox_head` 共用权重（`_tied_weights_keys`，
3//! 权重文件里没有另存）。阅读顺序头（`decoder_order_head` + 全局指针）只算最后一层。逐层掩码输出不算。
4
5use pagesift_infer::{LayerNorm, Linear, Tensor, Weights, ops};
6
7use super::config::Config;
8use super::deform::{DeformAttn, Level};
9use super::layers::{Ffn, Mlp, SelfAttention, add_norm, conv_bn, count};
10use super::query::{Bbox, QuerySelect, inverse_sigmoid, sigmoid};
11use crate::lcnet::ConvBn;
12
13/// 解码器的一层（`PPDocLayoutV3DecoderLayer`，后归一化）。
14#[derive(Debug, Clone)]
15struct Layer {
16    attn: SelfAttention,
17    norm1: LayerNorm,
18    cross: DeformAttn,
19    norm2: LayerNorm,
20    ffn: Ffn,
21    norm3: LayerNorm,
22}
23
24impl Layer {
25    fn load(w: &Weights, p: &str, cfg: &Config) -> Result<Layer, String> {
26        let eps = cfg.layer_norm_eps;
27        Ok(Layer {
28            attn: SelfAttention::load(w, &format!("{p}.self_attn"), cfg.decoder_heads)?,
29            norm1: LayerNorm::load(w, &format!("{p}.self_attn_layer_norm"), eps)?,
30            cross: DeformAttn::load(
31                w,
32                &format!("{p}.encoder_attn"),
33                cfg.decoder_heads,
34                cfg.decoder_points,
35            )?,
36            norm2: LayerNorm::load(w, &format!("{p}.encoder_attn_layer_norm"), eps)?,
37            ffn: Ffn::load(w, p, cfg.decoder_act)?,
38            norm3: LayerNorm::load(w, &format!("{p}.final_layer_norm"), eps)?,
39        })
40    }
41
42    fn forward(
43        &self,
44        x: &Tensor,
45        pos: &Tensor,
46        src: &Tensor,
47        levels: &[Level],
48        refs: &[Bbox],
49    ) -> Tensor {
50        let h = add_norm(x, &self.attn.forward(x, pos), &self.norm1);
51        let cross = self.cross.forward(&ops::add(&h, pos), src, levels, refs);
52        let h = add_norm(&h, &cross, &self.norm2);
53        add_norm(&h, &self.ffn.forward(&h), &self.norm3)
54    }
55}
56
57/// 解码器的原始输出。
58#[derive(Debug, Clone)]
59pub struct Raw {
60    /// 类别 logit `[Q, K]`。
61    pub logits: Tensor,
62    /// 框 (cx, cy, w, h)，相对网络输入 0–1。
63    pub boxes: Vec<Bbox>,
64    /// 阅读顺序 logit `[Q, Q]`，下三角（含对角线）无意义。
65    pub order: Tensor,
66}
67
68/// 整个解码端：输入投影、查询选取、解码层和各个头。
69#[derive(Debug, Clone)]
70pub struct Decoder {
71    input_proj: Vec<ConvBn>,
72    select: QuerySelect,
73    layers: Vec<Layer>,
74    query_pos: Mlp,
75    bbox: Mlp,
76    class: Linear,
77    norm: LayerNorm,
78    order_head: Linear,
79    pointer: Linear,
80    pointer_dim: usize,
81}
82
83impl Decoder {
84    pub fn load(w: &Weights, cfg: &Config) -> Result<Decoder, String> {
85        let n = count(w, |i| {
86            format!("model.decoder.layers.{i}.self_attn.q_proj.weight")
87        });
88        if n == 0 {
89            return Err("no decoder layers in the weights".into());
90        }
91        let layers = (0..n)
92            .map(|i| Layer::load(w, &format!("model.decoder.layers.{i}"), cfg))
93            .collect::<Result<_, _>>()?;
94        let input_proj = (0..cfg.feat_strides.len())
95            .map(|i| {
96                conv_bn(
97                    w,
98                    &format!("model.decoder_input_proj.{i}.0"),
99                    &format!("model.decoder_input_proj.{i}.1"),
100                    1,
101                    0,
102                    false,
103                    pagesift_infer::Activation::Identity,
104                )
105            })
106            .collect::<Result<_, _>>()?;
107        Ok(Decoder {
108            input_proj,
109            select: QuerySelect::load(w, cfg.layer_norm_eps, cfg.num_queries)?,
110            layers,
111            query_pos: Mlp::load(w, "model.decoder.query_pos_head")?,
112            // [协议] 这两个头与编码器共用权重（`_tied_weights_keys`）
113            bbox: Mlp::load(w, "model.enc_bbox_head")?,
114            class: Linear::load(w, "model.enc_score_head")?,
115            norm: LayerNorm::load(w, "model.decoder_norm", cfg.layer_norm_eps)?,
116            order_head: Linear::load(w, &format!("model.decoder_order_head.{}", n - 1))?,
117            pointer: Linear::load(w, "model.decoder_global_pointer.dense")?,
118            pointer_dim: cfg.pointer_dim,
119        })
120    }
121
122    /// `feats` 是编码器各层输出（步长从小到大），`mask_feat` 是掩码特征。
123    pub fn forward(&self, feats: &[Tensor], mask_feat: &Tensor) -> Raw {
124        let mut levels = Vec::with_capacity(feats.len());
125        let mut data = Vec::new();
126        let mut channels = 0;
127        let mut start = 0;
128        for (proj, f) in self.input_proj.iter().zip(feats) {
129            let x = proj.forward(f);
130            let (c, h, w) = x.chw();
131            levels.push(Level { h, w, start });
132            start += h * w;
133            channels = c;
134            data.extend_from_slice(&ops::chw_to_seq(&x).data);
135        }
136        let src = Tensor::new(vec![start, channels], data);
137
138        let (mut h, mut refs) = self.select.forward(&src, &levels, mask_feat);
139        for l in &self.layers {
140            let pos = self.query_pos.forward(&boxes_tensor(&refs));
141            h = l.forward(&h, &pos, &src, &levels, &refs);
142            let delta = self.bbox.forward(&h);
143            for (r, d) in refs.iter_mut().zip(delta.data.as_chunks::<BOX_LEN>().0) {
144                for (v, d) in r.iter_mut().zip(d) {
145                    *v = sigmoid(d + inverse_sigmoid(*v));
146                }
147            }
148        }
149        let q = self.norm.forward(&h);
150        let logits = self.class.forward(&q);
151        let order = pointer_logits(
152            &self.pointer.forward(&self.order_head.forward(&q)),
153            self.pointer_dim,
154        );
155        Raw {
156            logits,
157            boxes: refs,
158            order,
159        }
160    }
161}
162
163/// 一个框的分量数。
164const BOX_LEN: usize = size_of::<Bbox>() / size_of::<f32>();
165
166fn boxes_tensor(refs: &[Bbox]) -> Tensor {
167    Tensor::new(vec![refs.len(), BOX_LEN], refs.concat())
168}
169
170/// [协议] 源码 `logits.masked_fill(tril, -1e4)`
171const MASKED: f32 = -1e4;
172
173/// 全局指针（`PPDocLayoutV3GlobalPointer`）：`dense` 的输出 `[Q, 2·d]` 前半是 q、后半是 k，
174/// logit = qᵢ·kⱼ / √d；下三角（含对角线）填一个很小的数。
175pub(crate) fn pointer_logits(dense: &Tensor, d: usize) -> Tensor {
176    let (q, c) = dense.tc();
177    assert_eq!(c, 2 * d, "全局指针的宽度不符");
178    let mut qs = Vec::with_capacity(q * d);
179    let mut ks = Vec::with_capacity(q * d);
180    for row in dense.data.chunks_exact(c) {
181        qs.extend_from_slice(&row[..d]);
182        ks.extend_from_slice(&row[d..]);
183    }
184    let qs = Tensor::new(vec![q, d], qs);
185    let ks = ops::transpose(&Tensor::new(vec![q, d], ks));
186    let mut out = ops::matmul(&qs, &ks);
187    let scale = (d as f32).sqrt().recip();
188    for (i, row) in out.data.chunks_exact_mut(q).enumerate() {
189        for (j, v) in row.iter_mut().enumerate() {
190            *v = if j <= i { MASKED } else { *v * scale };
191        }
192    }
193    out
194}
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199
200    #[test]
201    fn 全局指针() {
202        // 2 个查询，d = 1（√d = 1）：q = [1, 3]，k = [2, 4]
203        let dense = Tensor::new(vec![2, 2], vec![1.0, 2.0, 3.0, 4.0]);
204        let o = pointer_logits(&dense, 1);
205        assert_eq!(o.data, [MASKED, 4.0, MASKED, MASKED]);
206    }
207}
pagesift_ocr/layout/decoder.rs

pagesift_ocr/layout/
decoder.rs
