mnist.cpp 4.4 KB
Newer Older
Scott Thornton's avatar
Scott Thornton committed
1
2
3
#include <cstdio>
#include <string>
#include <fstream>
Scott Thornton's avatar
Scott Thornton committed
4
#include <numeric>
Scott Thornton's avatar
Scott Thornton committed
5
6
#include <stdexcept>

7
#include <migraph/onnx.hpp>
Scott Thornton's avatar
Scott Thornton committed
8

9
10
#include <migraph/gpu/target.hpp>
#include <migraph/gpu/hip.hpp>
11
#include <migraph/generate.hpp>
Scott Thornton's avatar
Scott Thornton committed
12

13
14
#include "softmax.h"

15
auto reverse_int(unsigned int i)
Scott Thornton's avatar
Scott Thornton committed
16
{
17
18
19
20
21
22
23
24
    unsigned char c1, c2, c3, c4;
    c1 = i & 255u;
    c2 = (i >> 8u) & 255u;
    c3 = (i >> 16u) & 255u;
    c4 = (i >> 24u) & 255u;
    return (static_cast<unsigned int>(c1) << 24u) + (static_cast<unsigned int>(c2) << 16u) +
           (static_cast<unsigned int>(c3) << 8u) + c4;
};
Scott Thornton's avatar
Scott Thornton committed
25

Paul's avatar
Paul committed
26
27
std::vector<float>
read_mnist_images(const std::string& full_path, int& number_of_images, int& image_size)
28
{
29
    using uchar = unsigned char;
Scott Thornton's avatar
Scott Thornton committed
30
31
32
33
34
35
36

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

    if(file.is_open())
    {
        int magic_number = 0, n_rows = 0, n_cols = 0;

37
38
        file.read(reinterpret_cast<char*>(&magic_number), sizeof(magic_number));
        magic_number = reverse_int(magic_number);
Scott Thornton's avatar
Scott Thornton committed
39
40
41
42

        if(magic_number != 2051)
            throw std::runtime_error("Invalid MNIST image file!");

43
44
        file.read(reinterpret_cast<char*>(&number_of_images), sizeof(number_of_images));
        number_of_images = reverse_int(number_of_images);
Scott Thornton's avatar
Scott Thornton committed
45
        file.read(reinterpret_cast<char*>(&n_rows), sizeof(n_rows));
46
47
48
        n_rows = reverse_int(n_rows);
        file.read(reinterpret_cast<char*>(&n_cols), sizeof(n_cols));
        n_cols = reverse_int(n_cols);
Scott Thornton's avatar
Scott Thornton committed
49
50
51
52
53
54
55
56
57

        image_size = n_rows * n_cols;

        std::vector<float> result(number_of_images * image_size);
        for(int i = 0; i < number_of_images; i++)
        {
            for(int j = 0; j < image_size; j++)
            {
                uchar tmp;
58
                file.read(reinterpret_cast<char*>(&tmp), 1);
Scott Thornton's avatar
Scott Thornton committed
59
60
61
62
63
64
65
66
67
68
69
                result[i * image_size + j] = tmp / 255.0;
            }
        }
        return result;
    }
    else
    {
        throw std::runtime_error("Cannot open file `" + full_path + "`!");
    }
}

Paul's avatar
Paul committed
70
std::vector<int32_t> read_mnist_labels(const std::string& full_path, int& number_of_labels)
Scott Thornton's avatar
Scott Thornton committed
71
{
72
    using uchar = unsigned char;
Scott Thornton's avatar
Scott Thornton committed
73
74
75
76
77
78

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

    if(file.is_open())
    {
        int magic_number = 0;
79
80
        file.read(reinterpret_cast<char*>(&magic_number), sizeof(magic_number));
        magic_number = reverse_int(magic_number);
Scott Thornton's avatar
Scott Thornton committed
81
82
83
84

        if(magic_number != 2049)
            throw std::runtime_error("Invalid MNIST label file!");

85
86
        file.read(reinterpret_cast<char*>(&number_of_labels), sizeof(number_of_labels));
        number_of_labels = reverse_int(number_of_labels);
Scott Thornton's avatar
Scott Thornton committed
87
88
89
90
91

        std::vector<int32_t> result(number_of_labels);
        for(int i = 0; i < number_of_labels; i++)
        {
            uchar tmp;
92
            file.read(reinterpret_cast<char*>(&tmp), 1);
Scott Thornton's avatar
Scott Thornton committed
93
94
95
96
97
98
99
100
101
102
103
104
            result[i] = tmp;
        }
        return result;
    }
    else
    {
        throw std::runtime_error("Unable to open file `" + full_path + "`!");
    }
}

int main(int argc, char const* argv[])
{
105
    if(argc > 3)
Scott Thornton's avatar
Scott Thornton committed
106
107
108
109
110
111
112
113
114
115
    {
        std::string datafile        = argv[2];
        std::string labelfile       = argv[3];
        int nimages                 = -1;
        int image_size              = -1;
        int nlabels                 = -1;
        std::vector<float> input    = read_mnist_images(datafile, nimages, image_size);
        std::vector<int32_t> labels = read_mnist_labels(labelfile, nlabels);

        std::string file = argv[1];
116
        auto prog        = migraph::parse_onnx(file);
117
118
        std::cout << prog << std::endl << std::endl;
        prog.compile(migraph::gpu::target{});
119
        auto s = migraph::shape{migraph::shape::float_type, {1, 1, 28, 28}};
Scott Thornton's avatar
Scott Thornton committed
120
        std::cout << s << std::endl;
Scott Thornton's avatar
Scott Thornton committed
121
        auto ptr = input.data();
122
123
124
        migraph::program::parameter_map m;
        m["output"] =
            migraph::gpu::to_gpu(migraph::generate_argument(prog.get_parameter_shape("output")));
Scott Thornton's avatar
Scott Thornton committed
125
        for(int i = 0; i < 20; i++)
Scott Thornton's avatar
Scott Thornton committed
126
        {
127
            std::cout << "label: " << labels[i] << "  ---->  ";
128
129
            m["0"]      = migraph::gpu::to_gpu(migraph::argument{s, &ptr[784 * i]});
            auto result = migraph::gpu::from_gpu(prog.eval(m));
Scott Thornton's avatar
Scott Thornton committed
130
131
132
            std::vector<float> logits;
            result.visit([&](auto output) { logits.assign(output.begin(), output.end()); });
            std::vector<float> probs = softmax(logits);
Scott Thornton's avatar
Scott Thornton committed
133
134
            for(auto x : probs)
                std::cout << x << "  ";
135
            std::cout << std::endl;
Scott Thornton's avatar
Scott Thornton committed
136
        }
137
        std::cout << std::endl;
Scott Thornton's avatar
Scott Thornton committed
138
139
    }
}