onnx-ir 0.15.0

Library for parsing ONNX models
Documentation 
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

use super::ir::{AttributeValue, Node, NodeType};

/// Remap node type using kernel shape
pub fn remap_node_with_kernel_shape<F>(node: &mut Node, new_node_type: F)
where
    F: FnOnce(&Vec<i64>) -> NodeType,
{
    if let AttributeValue::Int64s(ints) = node.attrs.get("kernel_shape").unwrap() {
        node.node_type = new_node_type(ints);
    } else {
        panic!("kernel_shape is not an int64s");
    }
}

/// Remap node type to a more specific one
pub fn remap_node_type(node: &mut Node) {
    match node.node_type {
        NodeType::Conv => remap_node_with_kernel_shape(node, |ints| match ints.len() {
            1 => NodeType::Conv1d,
            2 => NodeType::Conv2d,
            3 => NodeType::Conv3d,
            _ => panic!("Only conv 1d, 2d and 3d are supported"),
        }),
        NodeType::ConvTranspose => remap_node_with_kernel_shape(node, |ints| match ints.len() {
            1 => NodeType::ConvTranspose1d,
            2 => NodeType::ConvTranspose2d,
            3 => NodeType::ConvTranspose3d,
            _ => panic!("Only conv_transpose 1d, 2d and 3d are supported"),
        }),
        NodeType::MaxPool => remap_node_with_kernel_shape(node, |ints| match ints.len() {
            1 => NodeType::MaxPool1d,
            2 => NodeType::MaxPool2d,
            _ => panic!("Only max_pool 1d and 2d are supported"),
        }),
        NodeType::AveragePool => remap_node_with_kernel_shape(node, |ints| match ints.len() {
            1 => NodeType::AveragePool1d,
            2 => NodeType::AveragePool2d,
            _ => panic!("Only avg_pool 1d and 2d are supported"),
        }),
        _ => (),
    }
}