read_onnx.cpp 1.68 KB
Newer Older
Paul's avatar
Paul committed
1
2
3
4
5
6

#include <google/protobuf/text_format.h>
#include <google/protobuf/io/zero_copy_stream_impl.h>
#include <onnx.pb.h>
#include <iostream>
#include <fstream>
Paul's avatar
Paul committed
7
#include <unordered_map>
Paul's avatar
Paul committed
8
9


Paul's avatar
Paul committed
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
std::unordered_map<std::string, onnx::AttributeProto> get_attributes(const onnx::NodeProto& node)
{
    std::unordered_map<std::string, onnx::AttributeProto> result;
    for(auto&& attr:node.attribute())
    {
        result[attr.name()] = attr;
    }
    return result;
}

void parse_graph(onnx::GraphProto graph)
{
    std::cout << "Graph name: " << graph.name() << std::endl;
    for(onnx::NodeProto node:graph.node()) {
        std::cout << "Layer: " << node.op_type() << std::endl;
        std::cout << "    Name: " << node.name() << std::endl;
        if(node.input_size() > 0)
            std::cout << "    Input: " << node.input(0) << std::endl;
        if(node.output_size() > 0)
            std::cout << "    Output: " << node.output(0) << std::endl;
    }
}

Paul's avatar
Paul committed
33
34
35
36
37
38
39
40
41
42
43
44
45
int main(int argc, char const *argv[])
{
    if(argc > 1)
    {
        std::string file = argv[1];
        std::fstream input(file.c_str(), std::ios::in | std::ios::binary);
        onnx::ModelProto model;
        if(model.ParseFromIstream(&input)) {
            std::cout << "Model version: " << model.model_version() << std::endl;
            std::cout << "Producer name: " << model.producer_name() << std::endl;
            std::cout << "Producer version: " << model.release_producer_version() << std::endl;
            if(model.has_graph()) {
                std::cout << "Model has graph" << std::endl;
Paul's avatar
Paul committed
46
                parse_graph(model.graph());
Paul's avatar
Paul committed
47
48
49
50
51
52
53
            }
        } else {
            std::cout << "Failed reading: " << file << std::endl;
        }

    }
}