read_onnx.cpp 972 Bytes
Newer Older
Paul's avatar
Paul committed
1

Paul's avatar
Paul committed
2
#include <rtg/onnx.hpp>
Paul's avatar
Paul committed
3

Paul's avatar
Paul committed
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
#include <rtg/cpu/cpu_target.hpp>
#include <random>

// TODO: Move this to a seperate header
std::vector<float> get_tensor_data(rtg::shape s)
{
    std::vector<float> result(s.elements());
    std::mt19937 engine{0};
    std::uniform_real_distribution<> dist;
    std::generate(result.begin(), result.end(), [&] { return dist(engine); });
    return result;
}

rtg::argument get_tensor_argument(rtg::shape s)
{
    auto v = get_tensor_data(s);
    return {s, [v]() mutable { return reinterpret_cast<char*>(v.data()); }};
}

Paul's avatar
Paul committed
23
int main(int argc, char const* argv[])
Paul's avatar
Paul committed
24
25
26
27
{
    if(argc > 1)
    {
        std::string file = argv[1];
Paul's avatar
Paul committed
28
        auto prog        = rtg::parse_onnx(file);
Paul's avatar
Paul committed
29
30
31
32
33
34
        prog.compile(rtg::cpu::cpu_target{});
        auto s      = prog.get_parameter_shape("Input3");
        auto input3 = get_tensor_argument(s);
        auto out    = prog.eval({{"Input3", input3}});
        (void)out;
        std::cout << prog << std::endl;
Paul's avatar
Paul committed
35
36
    }
}