Merge pull request #11 from wsttiger/mnist2

Added MNIST test / demo

src/include/migraph/operators.hpp

@@ -193,6 +193,7 @@ struct pooling
     std::array<std::size_t, 2> stride  = {{1, 1}};
     std::array<std::size_t, 2> lengths = {{1, 1}};
     std::string name() const { return "pooling"; }
+
     shape compute_shape(std::vector<shape> inputs) const
     {
         check_shapes{inputs, *this}.has(1).only_dims(4);
@@ -474,6 +475,7 @@ struct broadcast
         auto input  = inputs.at(1);
 
         std::vector<size_t> bcast_strides(result.lens().size(), 0);
+
         if(std::all_of(
                result.lens().cbegin(), result.lens().cend(), [&](auto x) { return x == 1; }))
         {

src/include/migraph/tensor_view.hpp

@@ -29,12 +29,14 @@ struct tensor_view
     template <class... Ts, MIGRAPH_REQUIRES(std::is_integral<Ts>{}...)>
     const T& operator()(Ts... xs) const
     {
+        assert(m_shape.index({static_cast<std::size_t>(xs)...}) < m_shape.bytes() / sizeof(T));
         return m_data[m_shape.index({static_cast<std::size_t>(xs)...})];
     }
 
     template <class... Ts, MIGRAPH_REQUIRES(std::is_integral<Ts>{}...)>
     T& operator()(Ts... xs)
     {
+        assert(m_shape.index({static_cast<std::size_t>(xs)...}) < m_shape.bytes() / sizeof(T));
         return m_data[m_shape.index({static_cast<std::size_t>(xs)...})];
     }
 

src/onnx/CMakeLists.txt

@@ -16,6 +16,10 @@ add_executable(read_onnx read_onnx.cpp)
 rocm_clang_tidy_check(read_onnx)
 target_link_libraries(read_onnx migraph_onnx)
 
+add_executable(mnist mnist.cpp)
+rocm_clang_tidy_check(mnist)
+target_link_libraries(mnist migraph_cpu migraph_onnx)
+
 if(MIGRAPH_ENABLE_MIOPEN)
 add_executable(verify_onnx verify_onnx.cpp)
 rocm_clang_tidy_check(verify_onnx)

src/onnx/mnist.cpp

+#include <cstdio>
+#include <string>
+#include <fstream>
+#include <numeric>
+#include <stdexcept>
+
+#include <migraph/onnx.hpp>
+
+#include <migraph/cpu/cpu_target.hpp>
+#include <migraph/generate.hpp>
+
+auto reverse_int(unsigned int i)
+{
+    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;
+};
+
+std::vector<float> read_mnist_images(std::string full_path, int& number_of_images, int& image_size)
+{
+    using uchar = unsigned char;
+
+    std::ifstream file(full_path, std::ios::binary);
+
+    if(file.is_open())
+    {
+        int magic_number = 0, n_rows = 0, n_cols = 0;
+
+        file.read(reinterpret_cast<char*>(&magic_number), sizeof(magic_number));
+        magic_number = reverse_int(magic_number);
+
+        if(magic_number != 2051)
+            throw std::runtime_error("Invalid MNIST image file!");
+
+        file.read(reinterpret_cast<char*>(&number_of_images), sizeof(number_of_images));
+        number_of_images = reverse_int(number_of_images);
+        file.read(reinterpret_cast<char*>(&n_rows), sizeof(n_rows));
+        n_rows = reverse_int(n_rows);
+        file.read(reinterpret_cast<char*>(&n_cols), sizeof(n_cols));
+        n_cols = reverse_int(n_cols);
+
+        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;
+                file.read(reinterpret_cast<char*>(&tmp), 1);
+                result[i * image_size + j] = tmp / 255.0;
+            }
+        }
+        return result;
+    }
+    else
+    {
+        throw std::runtime_error("Cannot open file `" + full_path + "`!");
+    }
+}
+
+std::vector<int32_t> read_mnist_labels(std::string full_path, int& number_of_labels)
+{
+    using uchar = unsigned char;
+
+    std::ifstream file(full_path, std::ios::binary);
+
+    if(file.is_open())
+    {
+        int magic_number = 0;
+        file.read(reinterpret_cast<char*>(&magic_number), sizeof(magic_number));
+        magic_number = reverse_int(magic_number);
+
+        if(magic_number != 2049)
+            throw std::runtime_error("Invalid MNIST label file!");
+
+        file.read(reinterpret_cast<char*>(&number_of_labels), sizeof(number_of_labels));
+        number_of_labels = reverse_int(number_of_labels);
+
+        std::vector<int32_t> result(number_of_labels);
+        for(int i = 0; i < number_of_labels; i++)
+        {
+            uchar tmp;
+            file.read(reinterpret_cast<char*>(&tmp), 1);
+            result[i] = tmp;
+        }
+        return result;
+    }
+    else
+    {
+        throw std::runtime_error("Unable to 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;
+}
+
+int main(int argc, char const* argv[])
+{
+    if(argc > 3)
+    {
+        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];
+        auto prog        = migraph::parse_onnx(file);
+        prog.compile(migraph::cpu::cpu_target{});
+        auto s = migraph::shape{migraph::shape::float_type, {1, 1, 28, 28}};
+        std::cout << s << std::endl;
+        auto ptr = input.data();
+        for(int i = 0; i < 20; i++)
+        {
+            std::cout << "label: " << labels[i] << "  ---->  ";
+            auto input3 = migraph::argument{s, &ptr[784 * i]};
+            auto result = prog.eval({{"Input3", input3}});
+            std::vector<float> logits;
+            result.visit([&](auto output) { logits.assign(output.begin(), output.end()); });
+            std::vector<float> probs = softmax(logits);
+            for(auto x : probs)
+                std::cout << x << "  ";
+            std::cout << std::endl;
+        }
+        std::cout << std::endl;
+    }
+}

src/onnx/onnx.cpp

@@ -113,13 +113,49 @@ struct onnx_parser
             }
             return prog.add_instruction(add{}, args);
         });
-        add_op("Sub", [this](attribute_map, std::vector<instruction_ref> args) {
+        add_op("Sub", [this](attribute_map attributes, std::vector<instruction_ref> args) {
+            if(contains(attributes, "broadcast"))
+            {
+                uint64_t broadcasted = parse_value(attributes.at("broadcast")).at<uint64_t>();
+                if(broadcasted != 0)
+                {
+                    uint64_t axis = (contains(attributes, "axis"))
+                                        ? parse_value(attributes.at("axis")).at<uint64_t>()
+                                        : 0;
+                    auto l = prog.add_instruction(broadcast{axis}, args);
+                    return prog.add_instruction(sub{}, args[0], l);
+                }
+            }
             return prog.add_instruction(sub{}, args);
         });
-        add_op("Mul", [this](attribute_map, std::vector<instruction_ref> args) {
+        add_op("Mul", [this](attribute_map attributes, std::vector<instruction_ref> args) {
+            if(contains(attributes, "broadcast"))
+            {
+                uint64_t broadcasted = parse_value(attributes.at("broadcast")).at<uint64_t>();
+                if(broadcasted != 0)
+                {
+                    uint64_t axis = (contains(attributes, "axis"))
+                                        ? parse_value(attributes.at("axis")).at<uint64_t>()
+                                        : 0;
+                    auto l = prog.add_instruction(broadcast{axis}, args);
+                    return prog.add_instruction(mul{}, args[0], l);
+                }
+            }
             return prog.add_instruction(mul{}, args);
         });
-        add_op("Div", [this](attribute_map, std::vector<instruction_ref> args) {
+        add_op("Div", [this](attribute_map attributes, std::vector<instruction_ref> args) {
+            if(contains(attributes, "broadcast"))
+            {
+                uint64_t broadcasted = parse_value(attributes.at("broadcast")).at<uint64_t>();
+                if(broadcasted != 0)
+                {
+                    uint64_t axis = (contains(attributes, "axis"))
+                                        ? parse_value(attributes.at("axis")).at<uint64_t>()
+                                        : 0;
+                    auto l = prog.add_instruction(broadcast{axis}, args);
+                    return prog.add_instruction(div{}, args[0], l);
+                }
+            }
             return prog.add_instruction(div{}, args);
         });
     }

test/cpu_ops_test.cpp

@@ -252,6 +252,63 @@ void gemm_test()
     }
 }
 
+void maxpool_test()
+{
+    migraph::program p;
+    std::vector<float> a = {
+        -2.1314404,  -1.63041711, 1.54562736,  1.04625261,  -1.42931843, -0.48703974, 0.4065806,
+        -0.1524526,  1.30775225,  0.45538983,  -0.06631992, -1.75332725, 1.33493888,  0.47327688,
+        0.36873096,  1.18358743,  -0.34640595, 1.22098756,  0.01946825,  -0.20238149, 0.43348005,
+        -0.67991608, -0.83041084, 0.93537551,  0.70241445,  -0.5654031,  -1.30899191, -0.26735824,
+        -0.52444768, 1.99097753,  1.86504853,  -0.26506025, 0.26236168,  0.43763575,  0.95300823,
+        -1.02733946, -0.74655169, -0.5374338,  -0.28901565, -0.59789604, 0.5310151,   0.99125904,
+        0.40609556,  -1.57175648, 0.22031412,  1.45862222,  0.53217483,  1.39087725,  1.00170159,
+        -0.87175864, -1.7204628,  -1.72008383, -0.38656762, -0.01443311, 1.46645272,  -1.39995027,
+        0.22505587,  -0.43461126, -0.05511411, -0.79950953, -0.01439556, 0.08795211,  1.18943918,
+        -0.84079367, -1.73383629, -0.55662078, -0.30626822, -0.67339015, 0.44179603,  0.54316711,
+        0.40899998,  -0.27831686, -1.11900508, -0.0881724,  0.35483059,  2.36277103,  -0.04765317,
+        -0.36865309, 0.73814237,  1.47151589,  1.36546791,  -0.32649881, -1.0517807,  2.24768877,
+        0.68883753,  0.58646208,  -0.91017133, -0.50462508, -0.4013325,  -0.72348958, -0.47368807,
+        0.35285577,  -1.01817429, -0.5152272,  0.60321307,  0.43521205,  -0.23733577, 0.66427642,
+        0.82949388,  0.82443929,  0.71550399,  0.34561086,  0.68570769,  -0.40718508, -1.20350206,
+        0.15793853,  -2.31013632, -0.07934658, -0.09348056, 0.36576006,  2.46601582,  0.11090943,
+        0.9144392,   0.56759721,  -0.22112127, -0.21955389, 0.72474903,  -1.28448462, 1.53285873,
+        0.37437943,  0.31409341,  1.95433736,  0.91620457,  0.86205518,  1.24365854,  0.19248386,
+        0.22526583,  0.13462132,  -0.27561715, -2.06446075, -0.02306402, -1.38278747, 1.1411345,
+        1.31293464,  -1.86041689, 1.06763375,  -0.26541466, 1.4545635,   1.11430049,  -0.66491818,
+        0.87101674,  0.67768967,  -1.02062869, -1.05031872, -2.2764678,  -2.0200038,  0.37592548,
+        -0.26701379, -0.83388507, 0.19403623,  1.00968623,  0.11020003,  1.16736257,  -1.1160326,
+        0.47346735,  0.6126079,   -0.19135755, 1.33624589,  -0.29802522, -0.57873946, -1.06555879,
+        -0.20686582, 1.36892557,  -0.19937795, 0.8649236,   -1.40126073, 1.53441942,  0.34682792,
+        -1.31724346, -1.32898355, 2.40126371,  0.07845283,  1.35732043,  -0.63678312, 0.39429256,
+        -1.36487007, -0.31026676, -0.44981545, -0.28994772, -0.14657612, -1.75206447, -0.70612341,
+        1.20071781,  -1.64647579, -0.7133292,  0.88494766,  0.52119428,  -2.77387547, 2.07681108,
+        -0.90133125, 0.2847338,   0.6174528,   -0.20616426, -0.64263535, -1.08496261, 0.54275119,
+        -0.88503587, 0.6629802,   1.47319221,  -1.05829155, -0.97027361, -0.93187737, -1.39954746,
+        -0.52359426, -0.14743951, 1.51522756,  0.2078452,   -1.28156149, -1.19363916, -0.78680223,
+        -0.89094824, 1.30212069,  -0.77974445, -0.58411664, 0.48764706,  -0.67132682};
+    std::vector<float> c = {1.33493888, 1.54562736, 1.22098756, 1.33493888, 1.18358743, 1.99097753,
+                            1.00170159, 1.45862222, 1.39087725, 1.46645272, 1.18943918, -0.01443311,
+                            1.47151589, 2.36277103, 2.24768877, 0.68883753, 0.82949388, 0.71550399,
+                            1.95433736, 2.46601582, 1.53285873, 1.95433736, 1.06763375, 1.4545635,
+                            1.33624589, 1.16736257, 0.6126079,  1.36892557, 2.40126371, 1.53441942,
+                            0.52119428, 2.07681108, 0.88494766, 1.51522756, 0.54275119, 0.6629802};
+    migraph::shape a_shape{migraph::shape::float_type, {2, 3, 6, 6}};
+    auto al = p.add_literal(migraph::literal{a_shape, a});
+    p.add_instruction(migraph::pooling{"max", {{0, 0}}, {{2, 2}}, {{3, 2}}}, al);
+    p.compile(migraph::cpu::cpu_target{});
+    auto result = p.eval({});
+    std::cout << result.get_shape() << std::endl;
+    std::vector<float> results_vector(36);
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    float tol = 1e-6;
+    for(int i = 0; i < results_vector.size(); i++)
+    {
+        // std::cout << results_vector[i] << "          " << c[i] << std::endl;
+        EXPECT(std::abs(results_vector[i] - c[i]) < tol);
+    }
+}
+
 void softmax_test()
 {
     migraph::program p;
@@ -564,6 +621,7 @@ int main()
     transpose_test();
     contiguous_test();
     softmax_test();
+    // maxpool_test();
     conv2d_test();
     conv2d_padding_test();
     conv2d_padding_stride_test();