resnet18.cpp 3.76 KB
Newer Older
1
2
3
4
5
6
7
8
#include <cstdio>
#include <string>
#include <fstream>
#include <numeric>
#include <stdexcept>

#include <migraph/onnx.hpp>

wsttiger's avatar
wsttiger committed
9
#include <migraph/cpu/cpu_target.hpp>
wsttiger's avatar
wsttiger committed
10
11
#include <migraph/gpu/target.hpp>
#include <migraph/gpu/hip.hpp>
12
13
#include <migraph/generate.hpp>

wsttiger's avatar
wsttiger committed
14
15
auto read_cifar10_images(std::string full_path)
{
16
17
18

    std::ifstream file(full_path, std::ios::binary);

wsttiger's avatar
wsttiger committed
19
    const size_t nimages          = 10;
20
    const size_t nbytes_per_image = 3072;
21
22
    std::vector<uint8_t> raw_data(nimages * (nbytes_per_image + 1));
    std::vector<uint8_t> labels(nimages);
wsttiger's avatar
wsttiger committed
23
24
    std::vector<float> data(nimages * nbytes_per_image);
    if(file.is_open())
25
    {
wsttiger's avatar
wsttiger committed
26
        file.read(reinterpret_cast<char*>(raw_data.data()),
27
28
                  (nbytes_per_image + 1) * nimages * sizeof(uint8_t));
        uint8_t* pimage = raw_data.data();
wsttiger's avatar
wsttiger committed
29
        for(size_t i = 0; i < nimages; i++, pimage += nbytes_per_image)
30
31
        {
            labels[i] = *pimage++;
wsttiger's avatar
wsttiger committed
32
            for(size_t j = 0; j < nbytes_per_image; j++)
33
            {
wsttiger's avatar
wsttiger committed
34
35
                float v                        = *(pimage + j) / 255.0f;
                data[i * nbytes_per_image + j] = v;
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
            }
        }
        return std::make_pair(labels, data);
    }
    else
    {
        throw std::runtime_error("Cannot open file `" + full_path + "`!");
    }
}

std::vector<float> softmax(std::vector<float> p)
{
    size_t n = p.size();
    std::vector<float> result(n);
    std::transform(p.begin(), p.end(), result.begin(), [](auto x) { return std::exp(x); });
    float s = std::accumulate(result.begin(), result.end(), 0.0f, std::plus<float>());
    std::transform(result.begin(), result.end(), result.begin(), [=](auto x) { return x / s; });
    return result;
}

56
57
int main(int argc, char const* argv[])
{
wsttiger's avatar
wsttiger committed
58
    std::string file     = argv[1];
59
    std::string datafile = argv[2];
wsttiger's avatar
wsttiger committed
60
    auto prog            = migraph::parse_onnx(file);
61
    std::cout << prog << std::endl;
62
63
    auto imageset = read_cifar10_images(datafile);

64
65
66
67
68
69
70
71
    // // GPU target
    // prog.compile(migraph::gpu::target{});
    // migraph::program::parameter_map m;
    // auto s = migraph::shape{migraph::shape::float_type, {1, 3, 32, 32}};
    // for(auto&& x : prog.get_parameter_shapes())
    // {
    //     m[x.first] = migraph::gpu::to_gpu(migraph::generate_argument(x.second));
    // }
wsttiger's avatar
wsttiger committed
72
73
74
75
76
77
    // auto labels = imageset.first;
    // auto input  = imageset.second;
    // auto ptr    = input.data();
    // for(int i = 0; i < 10; i++)
    // {
    //     std::cout << "label: " << (uint32_t)labels[i] << "  ---->  ";
78
79
    //     m["0"]      = migraph::gpu::to_gpu(migraph::argument{s, &ptr[3072 * i]});
    //     auto result = migraph::gpu::from_gpu(prog.eval(m));
wsttiger's avatar
wsttiger committed
80
81
82
83
    //     std::vector<float> logits;
    //     result.visit([&](auto output) { logits.assign(output.begin(), output.end()); });
    //     std::vector<float> probs = softmax(logits);
    //     for(auto x : logits)
84
85
86
87
    //         //std::cout << x << "  ";
    //         //std::cout << x << std::endl;
    //         printf("%10.5e    ", x);
    //     std::cout << std::endl;
wsttiger's avatar
wsttiger committed
88
89
    //     std::cout << std::endl;
    // }
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

    // // CPU target
    prog.compile(migraph::cpu::cpu_target{});
    auto s      = migraph::shape{migraph::shape::float_type, {1, 3, 32, 32}};
    auto labels = imageset.first;
    auto input  = imageset.second;
    auto ptr    = input.data();
    for(int i = 0; i < 10; i++)
    {
        std::cout << "label: " << (uint32_t)labels[i] << "  ---->  ";
        auto input3 = migraph::argument{s, &ptr[3072 * i]};
        auto result = prog.eval({{"0", input3}});
        std::vector<float> logits;
        result.visit([&](auto output) { logits.assign(output.begin(), output.end()); });
        std::vector<float> probs = softmax(logits);
        for(auto x : logits)
            printf("%10.5e    ", x);
        std::cout << std::endl;
    }
wsttiger's avatar
wsttiger committed
109
}