read_onnx.cpp 12.8 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
#include <functional>
Scott Thornton's avatar
Scott Thornton committed
9
#include <array>
Paul's avatar
Paul committed
10

Paul's avatar
Paul committed
11
#include <rtg/fallthrough.hpp>
Paul's avatar
Paul committed
12
#include <rtg/program.hpp>
Paul's avatar
Paul committed
13
#include <rtg/operators.hpp>
Paul's avatar
Paul committed
14
15
16
17

struct unknown
{
    std::string op;
Paul's avatar
Paul committed
18
    std::string name() const { return "unknown:" + op; }
Paul's avatar
Paul committed
19
20
    rtg::shape compute_shape(std::vector<rtg::shape> input) const
    {
Paul's avatar
Paul committed
21
22
23
24
        if(input.empty())
            return {};
        else
            return input.front();
Paul's avatar
Paul committed
25
    }
Paul's avatar
Paul committed
26
27
28
29
    rtg::argument compute(rtg::shape, std::vector<rtg::argument>) const
    {
        RTG_THROW("not computable");
    }
Paul's avatar
Paul committed
30
    friend std::ostream& operator<<(std::ostream& os, const unknown& x)
Paul's avatar
Paul committed
31
32
33
34
    {
        os << x.name();
        return os;
    }
Paul's avatar
Paul committed
35
};
Paul's avatar
Paul committed
36

Paul's avatar
Paul committed
37
template <class C, class T>
Paul's avatar
Paul committed
38
39
40
41
42
bool contains(C&& c, T&& x)
{
    return c.find(x) != c.end();
}

Paul's avatar
Paul committed
43
template <class Range, class Iterator>
Paul's avatar
Paul committed
44
45
46
47
48
void copy(Range&& r, Iterator it)
{
    std::copy(r.begin(), r.end(), it);
}

Paul's avatar
Paul committed
49
struct onnx_parser
Paul's avatar
Paul committed
50
{
Paul's avatar
Paul committed
51
    using attribute_map = std::unordered_map<std::string, onnx::AttributeProto>;
Paul's avatar
Paul committed
52
    using node_map      = std::unordered_map<std::string, onnx::NodeProto>;
Paul's avatar
Paul committed
53
54
    using op_func =
        std::function<rtg::instruction_ref(attribute_map, std::vector<rtg::instruction_ref>)>;
Paul's avatar
Paul committed
55
    node_map nodes;
Paul's avatar
Paul committed
56
    std::unordered_map<std::string, rtg::instruction_ref> instructions;
Paul's avatar
Paul committed
57
    rtg::program prog = rtg::program();
Paul's avatar
Paul committed
58

Paul's avatar
Paul committed
59
    std::unordered_map<std::string, op_func> ops;
Paul's avatar
Paul committed
60
61
62

    onnx_parser()
    {
Paul's avatar
Paul committed
63
        add_op("Conv", [this](attribute_map attributes, std::vector<rtg::instruction_ref> args) {
Paul's avatar
Paul committed
64
65
66
67
68
69
70
71
72
73
74
75
76
            rtg::convolution op;
            if(contains(attributes, "pads"))
            {
                copy(attributes["pads"].ints(), op.padding.begin());
            }
            if(contains(attributes, "strides"))
            {
                copy(attributes["strides"].ints(), op.stride.begin());
            }
            if(contains(attributes, "dilations"))
            {
                copy(attributes["dilations"].ints(), op.dilation.begin());
            }
Paul's avatar
Paul committed
77
            return prog.add_instruction(op, args);
Paul's avatar
Paul committed
78
        });
79
80
81
        add_op("MatMul", [this](attribute_map, std::vector<rtg::instruction_ref> args) {
            return prog.add_instruction(rtg::gemm{}, args);
        });
Paul's avatar
Paul committed
82
        add_op("MaxPool", [this](attribute_map attributes, std::vector<rtg::instruction_ref> args) {
Paul's avatar
Paul committed
83
84
85
86
87
88
89
90
91
92
93
94
95
96
            rtg::pooling op{"max"};
            // for(auto&& p:attributes) std::cout << p.first << std::endl;
            if(contains(attributes, "pads"))
            {
                copy(attributes["pads"].ints(), op.padding.begin());
            }
            if(contains(attributes, "strides"))
            {
                copy(attributes["strides"].ints(), op.stride.begin());
            }
            if(contains(attributes, "kernel_shape"))
            {
                copy(attributes["kernel_shape"].ints(), op.lengths.begin());
            }
Paul's avatar
Paul committed
97
            return prog.add_instruction(op, args);
Paul's avatar
Paul committed
98
        });
Paul's avatar
Paul committed
99
        add_op("Relu", [this](attribute_map, std::vector<rtg::instruction_ref> args) {
Paul's avatar
Paul committed
100
            return prog.add_instruction(rtg::activation{"relu"}, args);
Paul's avatar
Paul committed
101
        });
Paul's avatar
Paul committed
102
        add_op("Reshape", [this](attribute_map attributes, std::vector<rtg::instruction_ref> args) {
Paul's avatar
Paul committed
103
104
            rtg::reshape op;
            rtg::literal s = parse_value(attributes.at("shape"));
Paul's avatar
Paul committed
105
            s.visit([&](auto v) { copy(v, std::back_inserter(op.dims)); });
Paul's avatar
Paul committed
106
            return prog.add_instruction(op, args);
Paul's avatar
Paul committed
107
        });
Paul's avatar
Paul committed
108
        add_op("Constant", [this](attribute_map attributes, std::vector<rtg::instruction_ref>) {
Paul's avatar
Paul committed
109
            rtg::literal v = parse_value(attributes.at("value"));
Paul's avatar
Paul committed
110
            return prog.add_literal(v);
Paul's avatar
Paul committed
111
        });
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
        add_op("Add", [this](attribute_map attributes, std::vector<rtg::instruction_ref> args) {
            if (contains(attributes, "broadcast"))
            {
                uint64_t broadcast = parse_value(attributes.at("broadcast")).at<uint64_t>();
                if (broadcast != 0) {
                    uint64_t axis = (contains(attributes, "axis")) ? 
                        parse_value(attributes.at("axis")).at<uint64_t>() : 0;
                    auto l = prog.add_instruction(rtg::broadcast{axis}, args);
                    return prog.add_instruction(rtg::add{}, args[0], l);
                } 
            }
            return prog.add_instruction(rtg::add{}, args);
        });
        add_op("Sub", [this](attribute_map, std::vector<rtg::instruction_ref> args) {
            return prog.add_instruction(rtg::sub{}, args);
        });
        add_op("Mul", [this](attribute_map, std::vector<rtg::instruction_ref> args) {
            return prog.add_instruction(rtg::mul{}, args);
        });
        add_op("Div", [this](attribute_map, std::vector<rtg::instruction_ref> args) {
            return prog.add_instruction(rtg::div{}, args);
        });
Paul's avatar
Paul committed
134
135
    }

Paul's avatar
Paul committed
136
    template <class F>
Paul's avatar
Paul committed
137
138
139
140
141
142
143
144
    void add_op(std::string name, F f)
    {
        ops.emplace(name, f);
    }

    void parse_from(std::istream& is)
    {
        onnx::ModelProto model;
Paul's avatar
Paul committed
145
        if(model.ParseFromIstream(&is))
Paul's avatar
Paul committed
146
        {
Paul's avatar
Paul committed
147
            if(model.has_graph())
Paul's avatar
Paul committed
148
149
150
            {
                this->parse_graph(model.graph());
            }
Paul's avatar
Paul committed
151
152
        }
        else
Paul's avatar
Paul committed
153
154
155
156
157
        {
            throw std::runtime_error("Failed reading");
        }
    }

Paul's avatar
Paul committed
158
    void parse_graph(const onnx::GraphProto& graph)
Paul's avatar
Paul committed
159
    {
Paul's avatar
Paul committed
160
        nodes = get_nodes(graph);
Paul's avatar
Paul committed
161
        for(auto&& input : graph.input())
Paul's avatar
Paul committed
162
        {
Paul's avatar
Paul committed
163
            const std::string& name = input.name();
Paul's avatar
Paul committed
164
            // TODO: Get shape of input parameter
Paul's avatar
Paul committed
165
            rtg::shape s       = parse_type(input.type());
Paul's avatar
Paul committed
166
            instructions[name] = prog.add_parameter(name, s);
Paul's avatar
Paul committed
167
        }
Paul's avatar
Paul committed
168
        for(auto&& p : nodes)
Paul's avatar
Paul committed
169
170
171
        {
            this->parse_node(p.second.name());
        }
Paul's avatar
Paul committed
172
173
    }

Paul's avatar
Paul committed
174
    void parse_node(std::string name)
Paul's avatar
Paul committed
175
    {
Paul's avatar
Paul committed
176
        if(instructions.count(name) == 0)
Paul's avatar
Paul committed
177
178
        {
            auto&& node = nodes.at(name);
Paul's avatar
Paul committed
179
            std::vector<rtg::instruction_ref> args;
Paul's avatar
Paul committed
180
            for(auto&& input : node.input())
Paul's avatar
Paul committed
181
182
183
184
185
186
187
188
189
190
191
192
            {
                if(nodes.count(input) > 0)
                {
                    auto&& iname = nodes.at(input).name();
                    this->parse_node(iname);
                    args.push_back(instructions.at(iname));
                }
                else
                {
                    args.push_back(instructions.at(input));
                }
            }
Paul's avatar
Paul committed
193
            if(ops.count(node.op_type()) == 0)
Paul's avatar
Paul committed
194
            {
Paul's avatar
Paul committed
195
                instructions[name] = prog.add_instruction(unknown{node.op_type()}, args);
Paul's avatar
Paul committed
196
197
198
199
200
            }
            else
            {
                instructions[name] = ops[node.op_type()](get_attributes(node), args);
            }
Paul's avatar
Paul committed
201
        }
Paul's avatar
Paul committed
202
203
    }

Paul's avatar
Paul committed
204
    static attribute_map get_attributes(const onnx::NodeProto& node)
Paul's avatar
Paul committed
205
206
    {
        std::unordered_map<std::string, onnx::AttributeProto> result;
Paul's avatar
Paul committed
207
        for(auto&& attr : node.attribute())
Paul's avatar
Paul committed
208
        {
Paul's avatar
Paul committed
209
210
211
212
213
            result[attr.name()] = attr;
        }
        return result;
    }

Paul's avatar
Paul committed
214
    static node_map get_nodes(const onnx::GraphProto& graph)
Paul's avatar
Paul committed
215
216
    {
        std::unordered_map<std::string, onnx::NodeProto> result;
Paul's avatar
Paul committed
217
        for(auto&& node : graph.node())
Paul's avatar
Paul committed
218
219
        {
            result[node.name()] = node;
Paul's avatar
Paul committed
220
            for(auto&& output : node.output())
Paul's avatar
Paul committed
221
222
223
224
225
226
227
            {
                result[output] = node;
            }
        }
        return result;
    }

Paul's avatar
Paul committed
228
229
230
231
    static rtg::literal parse_value(const onnx::AttributeProto& attr)
    {
        switch(attr.type())
        {
Paul's avatar
Paul committed
232
233
234
235
236
237
238
239
240
241
242
243
244
245
        case onnx::AttributeProto::UNDEFINED: return {};
        case onnx::AttributeProto::FLOAT: return rtg::literal{attr.f()};
        case onnx::AttributeProto::INT: return rtg::literal{attr.i()};
        case onnx::AttributeProto::STRING: return {};
        case onnx::AttributeProto::TENSOR: return parse_tensor(attr.t());
        case onnx::AttributeProto::GRAPH: return {};
        case onnx::AttributeProto::FLOATS:
            return rtg::literal{rtg::shape::float_type, attr.floats().begin(), attr.floats().end()};
        case onnx::AttributeProto::INTS:
            return rtg::literal{rtg::shape::int32_type, attr.ints().begin(), attr.ints().end()};
            ;
        case onnx::AttributeProto::STRINGS: return {};
        case onnx::AttributeProto::TENSORS: return {};
        case onnx::AttributeProto::GRAPHS: return {};
Paul's avatar
Paul committed
246
        }
Paul's avatar
Paul committed
247
        RTG_THROW("Invalid attribute type");
Paul's avatar
Paul committed
248
249
250
251
252
253
254
    }

    static rtg::literal parse_tensor(const onnx::TensorProto& t)
    {
        std::vector<std::size_t> dims(t.dims().begin(), t.dims().end());
        switch(t.data_type())
        {
Paul's avatar
Paul committed
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
        case onnx::TensorProto::UNDEFINED: throw std::runtime_error("");
        case onnx::TensorProto::FLOAT:
            return rtg::literal{
                {rtg::shape::float_type, dims}, t.float_data().begin(), t.float_data().end()};
        case onnx::TensorProto::UINT8: throw std::runtime_error("");
        case onnx::TensorProto::INT8:
            return rtg::literal{
                {rtg::shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
        case onnx::TensorProto::UINT16:
            return rtg::literal{
                {rtg::shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
        case onnx::TensorProto::INT16:
            return rtg::literal{
                {rtg::shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
        case onnx::TensorProto::INT32:
            return rtg::literal{
                {rtg::shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
        case onnx::TensorProto::INT64:
            return rtg::literal{
                {rtg::shape::int64_type, dims}, t.int64_data().begin(), t.int64_data().end()};
        case onnx::TensorProto::STRING: throw std::runtime_error("");
        case onnx::TensorProto::BOOL:
            return rtg::literal{
                {rtg::shape::int32_type, dims}, t.int32_data().begin(), t.int32_data().end()};
        case onnx::TensorProto::FLOAT16: throw std::runtime_error("");
        case onnx::TensorProto::DOUBLE:
            return rtg::literal{
                {rtg::shape::double_type, dims}, t.double_data().begin(), t.double_data().end()};
        case onnx::TensorProto::UINT32: throw std::runtime_error("");
        case onnx::TensorProto::UINT64: throw std::runtime_error("");
        case onnx::TensorProto::COMPLEX64: throw std::runtime_error("");
        case onnx::TensorProto::COMPLEX128: throw std::runtime_error("");
Paul's avatar
Paul committed
287
        }
Paul's avatar
Paul committed
288
        RTG_THROW("Invalid tensor type");
Paul's avatar
Paul committed
289
    }
Paul's avatar
Paul committed
290
291
292

    static rtg::shape parse_type(const onnx::TypeProto& t)
    {
Paul's avatar
Paul committed
293
        rtg::shape::type_t shape_type{};
Paul's avatar
Paul committed
294
295
        switch(t.tensor_type().elem_type())
        {
Paul's avatar
Paul committed
296
297
        case onnx::TensorProto::UNDEFINED:
            break; // throw std::runtime_error("Unsupported type UNDEFINED");
Paul's avatar
Paul committed
298
        case onnx::TensorProto::FLOAT: shape_type = rtg::shape::float_type; break;
Paul's avatar
Paul committed
299
300
        case onnx::TensorProto::UINT8:
            break; // throw std::runtime_error("Unsupported type UINT8");
Paul's avatar
Paul committed
301
302
303
304
305
        case onnx::TensorProto::INT8: shape_type = rtg::shape::int8_type; break;
        case onnx::TensorProto::UINT16: shape_type = rtg::shape::uint16_type; break;
        case onnx::TensorProto::INT16: shape_type = rtg::shape::int16_type; break;
        case onnx::TensorProto::INT32: shape_type = rtg::shape::int32_type; break;
        case onnx::TensorProto::INT64: shape_type = rtg::shape::int64_type; break;
Paul's avatar
Paul committed
306
307
308
309
310
311
        case onnx::TensorProto::STRING:
            break; // throw std::runtime_error("Unsupported type STRING");
        case onnx::TensorProto::BOOL:
            break; // throw std::runtime_error("Unsupported type BOOL");
        case onnx::TensorProto::FLOAT16:
            break; // throw std::runtime_error("Unsupported type FLOAT16");
Paul's avatar
Paul committed
312
313
314
        case onnx::TensorProto::DOUBLE: shape_type = rtg::shape::double_type; break;
        case onnx::TensorProto::UINT32: shape_type = rtg::shape::uint32_type; break;
        case onnx::TensorProto::UINT64: shape_type = rtg::shape::uint64_type; break;
Paul's avatar
Paul committed
315
316
317
318
        case onnx::TensorProto::COMPLEX64:
            break; // throw std::runtime_error("Unsupported type COMPLEX64");
        case onnx::TensorProto::COMPLEX128:
            break; // throw std::runtime_error("Unsupported type COMPLEX128");
Paul's avatar
Paul committed
319
320
321
        }
        std::vector<std::size_t> dims;
        // TODO: USe std::transform
Paul's avatar
Paul committed
322
        for(auto&& d : t.tensor_type().shape().dim())
Paul's avatar
Paul committed
323
324
325
326
327
328
        {
            dims.push_back(d.dim_value());
        }
        return {shape_type, dims};
    }
};
Paul's avatar
Paul committed
329

Paul's avatar
Paul committed
330
int main(int argc, char const* argv[])
Paul's avatar
Paul committed
331
332
333
334
335
{
    if(argc > 1)
    {
        std::string file = argv[1];
        std::fstream input(file.c_str(), std::ios::in | std::ios::binary);
Paul's avatar
Paul committed
336
337
338
339
340
341
342
        onnx_parser parser;
        try
        {
            parser.parse_from(input);
        }
        catch(...)
        {
Paul's avatar
Paul committed
343
            std::cout << parser.prog << std::endl;
Paul's avatar
Paul committed
344
345
            throw;
        }
Paul's avatar
Paul committed
346
        std::cout << parser.prog << std::endl;
Paul's avatar
Paul committed
347
348
    }
}