Skip to main content

backbone_smoke/
backbone_smoke.rs

1//! 冒烟:载入模型目录的骨干,输入一张常数图,打印各段输出的形状和耗时。
2//!     cargo run --profile quick -p pagesift-ocr --example backbone_smoke -- <模型目录> <高> <宽>
3use std::time::Instant;
4fn main() {
5    let a: Vec<String> = std::env::args().skip(1).collect();
6    let dir = std::path::Path::new(&a[0]);
7    let (h, w): (usize, usize) = (a[1].parse().unwrap(), a[2].parse().unwrap());
8    let cfg: serde_json::Value =
9        serde_json::from_slice(&std::fs::read(dir.join("config.json")).unwrap()).unwrap();
10    let wts =
11        pagesift_infer::Weights::from_bytes(&std::fs::read(dir.join("model.safetensors")).unwrap())
12            .unwrap();
13    let c = pagesift_ocr::lcnet::Config::from_json(&cfg).unwrap();
14    let b = pagesift_ocr::lcnet::Backbone::load(&wts, "model.backbone.encoder", &c).unwrap();
15    let x = pagesift_infer::Tensor::new(
16        vec![3, h, w],
17        (0..3 * h * w)
18            .map(|i| ((i % 97) as f32 / 97.0) - 0.5)
19            .collect(),
20    );
21    let t = Instant::now();
22    let outs = b.forward(&x);
23    for o in &outs {
24        println!(
25            "{:?} 均值 {:.5}",
26            o.shape,
27            o.data.iter().sum::<f32>() / o.data.len() as f32
28        );
29    }
30    println!("{:.1} ms", t.elapsed().as_secs_f64() * 1e3);
31}