pagesift_ocr/layout/
encoder.rs

1//! 混合编码器（`PPDocLayoutV3HybridEncoder`，RT-DETR 式）：骨干后三段各投影到同一通道 → 最高层做
2//! 层内自注意力（AIFI）→ 自顶向下 FPN + 自底向上 PAN 融合；另外用融合结果和骨干第一段算掩码特征。
3//! RepVGG 块的两路卷积在载入时并成一个 3×3，结果只差浮点舍入。
4
5use pagesift_infer::{Activation, Conv2d, LayerNorm, Tensor, Weights, ops};
6
7use super::config::Config;
8use super::layers::{
9    Ffn, SelfAttention, add_norm, conv_bn, conv_layer, conv_norm, count, resize_bilinear,
10    sine_pos_2d,
11};
12use crate::lcnet::ConvBn;
13
14/// AIFI 的一层 Transformer（`PPDocLayoutV3EncoderLayer`，后归一化）。
15#[derive(Debug, Clone)]
16struct AifiLayer {
17    attn: SelfAttention,
18    norm1: LayerNorm,
19    ffn: Ffn,
20    norm2: LayerNorm,
21    pos_temperature: f64,
22}
23
24impl AifiLayer {
25    fn load(w: &Weights, p: &str, cfg: &Config) -> Result<AifiLayer, String> {
26        let eps = cfg.layer_norm_eps;
27        Ok(AifiLayer {
28            attn: SelfAttention::load(w, &format!("{p}.self_attn"), cfg.encoder_heads)?,
29            norm1: LayerNorm::load(w, &format!("{p}.self_attn_layer_norm"), eps)?,
30            ffn: Ffn::load(w, p, cfg.encoder_act)?,
31            norm2: LayerNorm::load(w, &format!("{p}.final_layer_norm"), eps)?,
32            pos_temperature: cfg.pos_temperature,
33        })
34    }
35
36    fn forward(&self, x: &Tensor) -> Tensor {
37        let (c, h, w) = x.chw();
38        let seq = ops::chw_to_seq(x);
39        let pos = sine_pos_2d(h, w, c, self.pos_temperature);
40        let y = add_norm(&seq, &self.attn.forward(&seq, &pos), &self.norm1);
41        let z = add_norm(&y, &self.ffn.forward(&y), &self.norm2);
42        ops::seq_to_chw(&z, h, w)
43    }
44}
45
46/// 把 1×1 卷积的权重并进 3×3 卷积的中心、偏置相加：`a(x) + b(x)` 变成一次卷积（RepVGG 重参数化）。
47fn merge_rep(mut a: ConvBn, b: &ConvBn) -> Result<ConvBn, String> {
48    let (o, c, kh, kw) = (
49        a.conv.weight.shape[0],
50        a.conv.weight.shape[1],
51        a.conv.weight.shape[2],
52        a.conv.weight.shape[3],
53    );
54    if b.conv.weight.shape != [o, c, 1, 1] {
55        return Err("the two RepVGG convolutions have different channels".into());
56    }
57    let center = (kh / 2) * kw + kw / 2;
58    let per = kh * kw;
59    for (i, v) in b.conv.weight.data.iter().enumerate() {
60        a.conv.weight.data[i * per + center] += v;
61    }
62    if let (Some(x), Some(y)) = (&mut a.conv.bias, &b.conv.bias) {
63        for (x, y) in x.iter_mut().zip(y) {
64            *x += y;
65        }
66    }
67    Ok(a)
68}
69
70/// `PPDocLayoutV3CSPRepLayer`：两个 1×1 分流，一路过若干 RepVGG 块，再相加。
71#[derive(Debug, Clone)]
72struct CspRep {
73    conv1: ConvBn,
74    conv2: ConvBn,
75    blocks: Vec<ConvBn>,
76}
77
78impl CspRep {
79    fn load(w: &Weights, p: &str, act: Activation) -> Result<CspRep, String> {
80        let n = count(w, |i| format!("{p}.bottlenecks.{i}.conv1.conv.weight"));
81        let mut blocks = Vec::with_capacity(n);
82        for i in 0..n {
83            let b = format!("{p}.bottlenecks.{i}");
84            // RepVggBlock：两路都不带激活，相加后再激活
85            let c3 = conv_norm(w, &format!("{b}.conv1"), 1, act)?;
86            let c1 = conv_norm(w, &format!("{b}.conv2"), 1, Activation::Identity)?;
87            blocks.push(merge_rep(c3, &c1)?);
88        }
89        Ok(CspRep {
90            conv1: conv_norm(w, &format!("{p}.conv1"), 1, act)?,
91            conv2: conv_norm(w, &format!("{p}.conv2"), 1, act)?,
92            blocks,
93        })
94    }
95
96    fn forward(&self, x: &Tensor) -> Tensor {
97        let mut a = self.conv1.forward(x);
98        for b in &self.blocks {
99            a = b.forward(&a);
100        }
101        ops::add(&a, &self.conv2.forward(x))
102    }
103}
104
105/// 掩码特征头（`PPDocLayoutV3MaskFeatFPN` 及其后的上采样、横向连接、输出卷积）。
106#[derive(Debug, Clone)]
107struct MaskHead {
108    /// 每层特征一组卷积；除了步长最小的那层，每个卷积后面都上采样两倍。
109    scale_heads: Vec<Vec<ConvBn>>,
110    output_conv: ConvBn,
111    lateral: ConvBn,
112    base_conv: ConvBn,
113    conv: Conv2d,
114}
115
116impl MaskHead {
117    fn load(w: &Weights, p: &str, cfg: &Config) -> Result<MaskHead, String> {
118        let act = Activation::Silu; // [协议] 这几层在源码里写死 activation="silu"
119        let mut scale_heads = Vec::with_capacity(cfg.feat_strides.len());
120        let base = cfg.feat_strides[0];
121        for (i, &stride) in cfg.feat_strides.iter().enumerate() {
122            // head_length = max(1, log2(步长 / 基础步长))；ModuleList 里卷积与 Upsample 交替，
123            // 所以有上采样时卷积在偶数下标
124            let n = (stride / base).ilog2().max(1) as usize;
125            let step = if stride != base { 2 } else { 1 };
126            let head = (0..n)
127                .map(|k| {
128                    conv_layer(
129                        w,
130                        &format!("{p}.mask_feature_head.scale_heads.{i}.layers.{}", k * step),
131                        1,
132                        false,
133                        act,
134                    )
135                })
136                .collect::<Result<_, _>>()?;
137            scale_heads.push(head);
138        }
139        Ok(MaskHead {
140            scale_heads,
141            output_conv: conv_layer(
142                w,
143                &format!("{p}.mask_feature_head.output_conv"),
144                1,
145                false,
146                act,
147            )?,
148            lateral: conv_layer(w, &format!("{p}.encoder_mask_lateral"), 1, false, act)?,
149            base_conv: conv_layer(
150                w,
151                &format!("{p}.encoder_mask_output.base_conv"),
152                1,
153                false,
154                act,
155            )?,
156            conv: Conv2d::load(
157                w,
158                &format!("{p}.encoder_mask_output.conv"),
159                (1, 1),
160                (0, 0),
161                (1, 1),
162                1,
163            )?,
164        })
165    }
166
167    /// `pan` 按步长从小到大；`x4` 是骨干第一段。结果 `[原型数, H/4, W/4]`。
168    fn forward(&self, pan: &[Tensor], x4: &Tensor) -> Tensor {
169        let run = |i: usize, x: &Tensor| {
170            let mut h = x.clone();
171            for c in &self.scale_heads[i] {
172                h = c.forward(&h);
173                if i != 0 {
174                    let (_, ph, pw) = h.chw();
175                    h = resize_bilinear(&h, ph * 2, pw * 2);
176                }
177            }
178            h
179        };
180        let mut out = run(0, &pan[0]);
181        let (_, oh, ow) = out.chw();
182        for (i, p) in pan.iter().enumerate().skip(1) {
183            out = ops::add(&out, &resize_bilinear(&run(i, p), oh, ow));
184        }
185        let out = self.output_conv.forward(&out);
186        let up = resize_bilinear(&out, oh * 2, ow * 2);
187        let fused = ops::add(&up, &self.lateral.forward(x4));
188        self.conv.forward(&self.base_conv.forward(&fused))
189    }
190}
191
192/// 整个混合编码器。
193#[derive(Debug, Clone)]
194pub struct Encoder {
195    input_proj: Vec<ConvBn>,
196    aifi: AifiLayer,
197    aifi_level: usize,
198    lateral: Vec<ConvBn>,
199    fpn: Vec<CspRep>,
200    down: Vec<ConvBn>,
201    pan: Vec<CspRep>,
202    mask: MaskHead,
203}
204
205/// 编码器的输出。
206#[derive(Debug, Clone)]
207pub struct EncoderOut {
208    /// 融合后的各层特征，步长从小到大。
209    pub feats: Vec<Tensor>,
210    /// 掩码特征 `[原型数, H/4, W/4]`。
211    pub mask_feat: Tensor,
212}
213
214/// [协议] `PPDocLayoutV3ConvNormLayer` 下采样卷积的步长
215const PAN_STRIDE: usize = 2;
216
217impl Encoder {
218    pub fn load(w: &Weights, cfg: &Config) -> Result<Encoder, String> {
219        let levels = cfg.feat_strides.len();
220        let act = cfg.act;
221        let p = "model.encoder";
222        let input_proj = (0..levels)
223            .map(|i| {
224                conv_bn(
225                    w,
226                    &format!("model.encoder_input_proj.{i}.0"),
227                    &format!("model.encoder_input_proj.{i}.1"),
228                    1,
229                    0,
230                    false,
231                    Activation::Identity,
232                )
233            })
234            .collect::<Result<_, _>>()?;
235        let many = |name: &str, stride: usize| -> Result<Vec<ConvBn>, String> {
236            (0..levels - 1)
237                .map(|i| conv_norm(w, &format!("{p}.{name}.{i}"), stride, act))
238                .collect()
239        };
240        let csp = |name: &str| -> Result<Vec<CspRep>, String> {
241            (0..levels - 1)
242                .map(|i| CspRep::load(w, &format!("{p}.{name}.{i}"), act))
243                .collect()
244        };
245        Ok(Encoder {
246            input_proj,
247            // [协议] 发布权重里 AIFI 叫 encoder.0.layers.0（transformers 载入时才改名 aifi）
248            aifi: AifiLayer::load(w, &format!("{p}.encoder.0.layers.0"), cfg)?,
249            aifi_level: cfg.aifi_level,
250            lateral: many("lateral_convs", 1)?,
251            fpn: csp("fpn_blocks")?,
252            down: many("downsample_convs", PAN_STRIDE)?,
253            pan: csp("pan_blocks")?,
254            mask: MaskHead::load(w, p, cfg)?,
255        })
256    }
257
258    /// `backbone` 是骨干各段的输出：第一段给掩码特征，其余进 FPN。
259    pub fn forward(&self, backbone: &[Tensor]) -> EncoderOut {
260        let (x4, rest) = backbone.split_first().expect("骨干没有输出");
261        let mut f: Vec<Tensor> = self
262            .input_proj
263            .iter()
264            .zip(rest)
265            .map(|(c, x)| c.forward(x))
266            .collect();
267        f[self.aifi_level] = self.aifi.forward(&f[self.aifi_level]);
268
269        // 自顶向下：上一层经横向卷积后上采样，与下一层拼接
270        let stages = self.lateral.len();
271        let mut fpn = vec![f[stages].clone()];
272        for (idx, (lat, block)) in self.lateral.iter().zip(&self.fpn).enumerate() {
273            let top = lat.forward(fpn.last().expect("fpn 不空"));
274            let (_, h, w) = top.chw();
275            let up = ops::interpolate_nearest(&top, h * 2, w * 2);
276            *fpn.last_mut().expect("fpn 不空") = top;
277            let cat = ops::concat_channels(&[&up, &f[stages - idx - 1]]);
278            fpn.push(block.forward(&cat));
279        }
280        fpn.reverse();
281
282        // 自底向上：下一层经下采样后与上一层拼接
283        let mut pan = vec![fpn[0].clone()];
284        for (idx, (down, block)) in self.down.iter().zip(&self.pan).enumerate() {
285            let d = down.forward(pan.last().expect("pan 不空"));
286            let cat = ops::concat_channels(&[&d, &fpn[idx + 1]]);
287            pan.push(block.forward(&cat));
288        }
289        let mask_feat = self.mask.forward(&pan, x4);
290        EncoderOut {
291            feats: pan,
292            mask_feat,
293        }
294    }
295}
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300
301    #[test]
302    fn repvgg_合并等于两路相加() {
303        // 3×3 + 1×1 并成一个 3×3 后，前向结果应等于两路分别卷积再相加
304        let wa: Vec<f32> = (0..2 * 2 * 9).map(|i| (i as f32 * 0.37).sin()).collect();
305        let wb: Vec<f32> = (0..2 * 2).map(|i| (i as f32 * 0.91).cos()).collect();
306        let conv = |shape: Vec<usize>, data: Vec<f32>, pad: usize, bias: Vec<f32>| ConvBn {
307            conv: Conv2d {
308                weight: Tensor::new(shape, data),
309                bias: Some(bias),
310                stride: (1, 1),
311                pad: (pad, pad),
312                dilation: (1, 1),
313                groups: 1,
314            },
315            act: Activation::Identity,
316        };
317        let a = conv(vec![2, 2, 3, 3], wa, 1, vec![0.1, -0.2]);
318        let b = conv(vec![2, 2, 1, 1], wb, 0, vec![0.3, 0.4]);
319        let x = Tensor::new(
320            vec![2, 4, 5],
321            (0..40).map(|i| (i as f32 * 0.13).sin()).collect(),
322        );
323        let want = ops::add(&a.forward(&x), &b.forward(&x));
324        let got = merge_rep(a, &b).unwrap().forward(&x);
325        for (p, q) in got.data.iter().zip(&want.data) {
326            assert!((p - q).abs() < 1e-5, "{p} {q}");
327        }
328    }
329}
pagesift_ocr/layout/encoder.rs

pagesift_ocr/layout/
encoder.rs
