wonnx-cli 0.5.1

CLI for WONNX. WONNX is an ONNX runtime based on wgpu aimed at being a universal GPU runtime, written in Rust.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

use std::{collections::HashMap, fmt::Display};

use wonnx::onnx::{GraphProto, ModelProto};

use crate::types::TraceOptions;

enum Source {
    NodeOutput {
        node_name: String,
        node_index: usize,
        op_type: String,
        output_name: String,
        output_index: usize,
    },
    Initializer {
        name: String,
    },
    Input {
        name: String,
    },
}

impl Display for Source {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Source::NodeOutput {
                node_name,
                node_index,
                op_type,
                output_name,
                output_index,
            } => write!(
                f,
                "{output_name}: output #{output_index} of node #{node_index} '{node_name}' ({op_type})"
            ),
            Source::Initializer { name } => write!(f, "{name}: initializer"),
            Source::Input { name } => write!(f, "{name}: input"),
        }
    }
}

pub fn trace_command(model: &ModelProto, opt: &TraceOptions) {
    // Build a map of all outputs
    let mut outputs: HashMap<String, Source> = HashMap::new();
    let graph = model.get_graph();

    for (node_index, node) in graph.node.iter().enumerate() {
        let node_name = node.get_name();
        for (output_index, output_name) in node.get_output().iter().enumerate() {
            outputs.insert(
                output_name.clone(),
                Source::NodeOutput {
                    node_name: node_name.to_string(),
                    node_index,
                    op_type: node.get_op_type().to_string(),
                    output_name: output_name.clone(),
                    output_index,
                },
            );
        }
    }

    for input in graph.input.iter() {
        outputs.insert(
            input.get_name().to_string(),
            Source::Input {
                name: input.get_name().to_string(),
            },
        );
    }

    for initializer in graph.initializer.iter() {
        outputs.insert(
            initializer.get_name().to_string(),
            Source::Initializer {
                name: initializer.get_name().to_string(),
            },
        );
    }

    // Print race
    if let Some(source) = outputs.get(&opt.output_name) {
        source.trace(&outputs, graph, 0, opt.maximum_depth);
    }
}

impl Source {
    fn trace(
        &self,
        map: &HashMap<String, Source>,
        graph: &GraphProto,
        depth: usize,
        max_depth: Option<usize>,
    ) {
        if let Some(md) = max_depth {
            if depth >= md {
                return;
            }
        }
        let tab = "| ";
        let tabs = tab.repeat(depth);
        println!("{tabs}+ {}", self);

        match self {
            Source::NodeOutput { node_index, .. } => {
                let node = &graph.get_node()[*node_index];
                for input in &node.input {
                    if let Some(source) = map.get(input) {
                        source.trace(map, graph, depth + 1, max_depth);
                    } else {
                        println!("{tabs}{tab}* '{}' NOT FOUND", input);
                    }
                }
            }
            Source::Initializer { .. } | Source::Input { .. } => (),
        }
    }
}