ENH: finish the conversion of YOLO v4 to CoreML package.

cpp/examples/Sara/NeuralNetworks/check_yolo_network.cpp

@@ -9,11 +9,11 @@
 // you can obtain one at http://mozilla.org/MPL/2.0/.
 // ========================================================================== //
 
-#include <DO/Sara/Core.hpp>
 #include <DO/Sara/Core/TicToc.hpp>
 #include <DO/Sara/Graphics.hpp>
 #include <DO/Sara/ImageIO.hpp>
 #include <DO/Sara/ImageProcessing.hpp>
+#include <DO/Sara/NeuralNetworks/Darknet/Debug.hpp>
 #include <DO/Sara/NeuralNetworks/Darknet/Network.hpp>
 #include <DO/Sara/NeuralNetworks/Darknet/Parser.hpp>
 #include <DO/Sara/NeuralNetworks/Darknet/YoloUtilities.hpp>
@@ -29,15 +29,10 @@
 #  include <omp.h>
 #endif
 
-#define COMPARE_WITH_DARKNET_OUTPUT
-#if defined(COMPARE_WITH_DARKNET_OUTPUT)
-#  include <DO/Sara/NeuralNetworks/Darknet/Debug.hpp>
-#endif
-
 
-namespace d = DO::Sara::Darknet;
 namespace fs = std::filesystem;
 namespace sara = DO::Sara;
+namespace d = DO::Sara::Darknet;
 
 
 auto check_yolo_implementation(d::Network& model, const fs::path& output_dir)
@@ -89,6 +84,35 @@ auto check_yolo_implementation(d::Network& model, const fs::path& output_dir)
 }
 
 
+auto save_network_fused_conv_weights(d::Network& model,
+                                     const fs::path& output_dir) -> void
+{
+  if (!fs::exists(output_dir))
+    throw std::runtime_error{fmt::format("Ouput directory '{}' does not exist!",
+                                         output_dir.string())};
+
+  model.debug = true;
+
+  const auto& net = model.net;
+  auto conv_idx = 0;
+
+  for (auto layer = 1u; layer < net.size(); ++layer)
+  {
+    auto conv = dynamic_cast<d::Convolution*>(net[layer].get());
+    if (!conv)
+      continue;
+
+    d::write_tensor(conv->weights.w,
+                    output_dir / fmt::format("conv_weight_{}.bin", conv_idx));
+
+    const auto b_tensor = sara::TensorView_<float, 1>{
+        conv->weights.b.data(), static_cast<int>(conv->weights.b.size())};
+    d::write_tensor(b_tensor,
+                    output_dir / fmt::format("conv_bias_{}.bin", conv_idx));
+    ++conv_idx;
+  }
+}
+
 auto save_network_intermediate_outputs(d::Network& model,
                                        const fs::path& image_path,
                                        const fs::path& output_dir) -> void
@@ -105,7 +129,7 @@ auto save_network_intermediate_outputs(d::Network& model,
   const auto input_chw = sara::tensor_view(input_hwc).transpose({2, 0, 1});
   static_assert(std::is_same_v<decltype(input_chw),  //
                                const sara::Tensor_<float, 3>>);
-  sara::toc("Image transpose");
+  sara::toc("Image transpose (HWC -> CHW)");
 
   sara::tic();
   const auto& input_layer = dynamic_cast<const d::Input&>(*model.net.front());
@@ -123,9 +147,7 @@ auto save_network_intermediate_outputs(d::Network& model,
   model.forward(input_nchw);
 
   // Save the input tensor.
-  static const auto data_dir_path =
-      fs::path{"/Users/oddkiva/Desktop/yolo-intermediate-out/"};
-  d::write_tensor(input_nchw, data_dir_path / "yolo_inter_0.bin");
+  d::write_tensor(input_nchw, output_dir / "yolo_inter_0.bin");
 
   // Save the intermediate output tensors.
   const auto& net = model.net;
@@ -136,39 +158,16 @@ auto save_network_intermediate_outputs(d::Network& model,
               << *net[layer] << std::endl;
 
     d::write_tensor(net[layer]->output,
-                    data_dir_path / fmt::format("yolo_inter_{}.bin", layer));
+                    output_dir / fmt::format("yolo_inter_{}.bin", layer));
   }
 }
 
 auto graphics_main(int, char**) -> int
 {
-#if defined(DEBUG_TENSOR_IO)
-  const auto shape = Eigen::Vector4i{1, 1, 3, 3};
-  auto x = sara::Tensor_<float, 4>{shape};
-  // clang-format off
-  x.flat_array() <<
-    0, 1, 2,
-    3, 4, 5,
-    6, 7, 8;
-  // clang-format on
-
-  const auto data_dir_path =
-      fs::path{"/Users/oddkiva/Desktop/yolo-intermediate-out"};
-  if (!fs::exists(data_dir_path))
-    fs::create_directory(data_dir_path);
-
-  const auto x_path = data_dir_path / "x.bin";
-  d::write_tensor(x, x_path.string());
-
-  const auto x2 = d::read_tensor(x_path.string());
-
-  if (x != x2)
-    throw std::runtime_error{"Implementation error for 4D-tensor IO"};
-  fmt::print("TENSOR IO: ALL GOOD!\n");
-#else
   const auto model_dir_path = fs::canonical(  //
       fs::path{src_path("trained_models")}    //
   );
+
   fmt::print("YOLO: {}\n", model_dir_path.string());
   if (!fs::exists(model_dir_path))
     throw std::runtime_error{"trained_models directory does not exist"};
@@ -182,18 +181,18 @@ auto graphics_main(int, char**) -> int
   auto model = d::load_yolo_model(yolo_dir_path, yolo_version, is_tiny);
   fmt::print("Load model OK!");
 
-  const auto yolo_out_dir_path =
-      fs::path{"/Users/oddkiva/Desktop/yolo-intermediate-out"};
-  if (!fs::exists(yolo_out_dir_path))
-    fs::create_directory(yolo_out_dir_path);
+  const auto yolo_data_check_dir_path = yolo_dir_path / "data_check";
+  if (!fs::exists(yolo_data_check_dir_path))
+    fs::create_directory(yolo_data_check_dir_path);
 
   const auto data_dir_path = fs::canonical(fs::path{src_path("data")});
   const auto image_path = data_dir_path / "dog.jpg";
   if (!fs::exists(image_path))
     throw std::runtime_error{"image does not exist!"};
 
-  save_network_intermediate_outputs(model, image_path, yolo_out_dir_path);
-#endif
+  save_network_fused_conv_weights(model, yolo_data_check_dir_path);
+  save_network_intermediate_outputs(model, image_path,
+                                    yolo_data_check_dir_path);
 
   return 0;
 }

cpp/src/DO/Sara/NeuralNetworks/Darknet/Debug.hpp

@@ -22,18 +22,20 @@
 namespace DO::Sara::Darknet {
 
   // CAVEAT: this is sensitive to the CPU architecture endianness.
-  inline auto write_tensor(const TensorView_<float, 4>& x,
+  template <typename T, int N>
+  inline auto write_tensor(const TensorView_<T, N>& x,
                            const std::string& filepath) -> void
   {
     auto file = std::ofstream{filepath, std::ios::binary};
     if (!file.is_open())
       throw std::runtime_error{"Error: could not open file: " + filepath + "!"};
 
+    // Write the tensor dimension.
     file.write(reinterpret_cast<const char*>(x.sizes().data()),
                x.sizes().size() * sizeof(int));
 
-    file.write(reinterpret_cast<const char*>(x.data()),
-               x.size() * sizeof(float));
+    // Copy the tensor content.
+    file.write(reinterpret_cast<const char*>(x.data()), x.size() * sizeof(T));
   }
 
   // CAVEAT: this is sensitive to the CPU architecture endianness.

cpp/src/DO/Sara/NeuralNetworks/Darknet/Layer.hpp

@@ -187,7 +187,8 @@ namespace DO::Sara::Darknet {
 
     auto to_output_stream(std::ostream&) const -> void override;
 
-    auto forward(const TensorView_<float, 4>& fx, const TensorView_<float, 4>& x)
+    auto forward(const TensorView_<float, 4>& fx,
+                 const TensorView_<float, 4>& x)
         -> const TensorView_<float, 4>&;
   };
 

cpp/test/Sara/NeuralNetworks/test_neuralnetworks_yolo_v4_config_parsing.cpp

@@ -13,6 +13,7 @@
 
 #include <DO/Sara/Defines.hpp>
 
+#include <DO/Sara/NeuralNetworks/Darknet/Debug.hpp>
 #include <DO/Sara/NeuralNetworks/Darknet/Layer.hpp>
 #include <DO/Sara/NeuralNetworks/Darknet/Parser.hpp>
 
@@ -27,6 +28,27 @@ namespace sara = DO::Sara;
 
 BOOST_AUTO_TEST_SUITE(TestLayers)
 
+BOOST_AUTO_TEST_CASE(test_tensor_io)
+{
+  namespace d = sara::Darknet;
+
+  const auto shape = Eigen::Vector4i{1, 1, 3, 3};
+  auto x = sara::Tensor_<float, 4>{shape};
+  // clang-format off
+  x.flat_array() <<
+    0, 1, 2,
+    3, 4, 5,
+    6, 7, 8;
+  // clang-format on
+
+  const auto x_path = fs::path{"x.bin"};
+  d::write_tensor(x, x_path.string());
+
+  const auto x2 = d::read_tensor(x_path.string());
+
+  BOOST_CHECK(x == x2);
+}
+
 BOOST_AUTO_TEST_CASE(test_yolov4_tiny_config_parsing)
 {
   const auto model_dir_path =

python/oddkiva/shakti/inference/coreml/yolo_v4.py

@@ -0,0 +1,30 @@
+from pathlib import Path
+
+import torch
+
+import coremltools as ct
+
+import oddkiva.shakti.inference.darknet as darknet
+
+
+def convert_yolo_v4_to_coreml(yolo_net: darknet.Network,
+                              in_tensor: torch.Tensor,
+                              path: Path):
+    with torch.inference_mode():
+        traced_model = torch.jit.trace(yolo_net, in_tensor)
+        outs = traced_model(in_tensor)
+
+        ct_ins = [ct.ImageType(name="image",
+                               shape=in_tensor.shape,
+                               scale=1 / 255)]
+        ct_outs = [ct.TensorType(name=f'yolo_{i}') for i, _ in enumerate(outs)]
+
+        model = ct.convert(traced_model, inputs=ct_ins, outputs=ct_outs,
+                           debug=True)
+
+        model.input_description["image"] = "Input RGB image"
+        for i in range(len(outs)):
+            model.output_description[f"yolo_{i}"] =\
+                f"Box predictions at scale {i}"
+
+        model.save(path)

python/oddkiva/shakti/inference/darknet/network.py

@@ -112,8 +112,6 @@ def load_convolutional_weights(self,
                         shape, weight_loader)
                     block.bn_weights['running_var'] = self._read_weights(
                         shape, weight_loader)
-                    if block_idx == 0:
-                        logging.info(block.bn_weights)
                 else:
                     bn = block.layers[1]
                     self._load_weights(bn.weight, weight_loader)
@@ -236,9 +234,8 @@ def _append_shortcut(self, model, layer_params, shortcut_id):
 
         # Calculate the output shape.
         assert shape_ins[0][2:] == shape_ins[1][2:]
-        n, _, h_in, w_in = shape_ins[0]
-        c_out = sum([shape_in[1] for shape_in in shape_ins])
-        shape_out = (n, c_out, h_in, w_in)
+        n, c_in, h_in, w_in = shape_ins[0]
+        shape_out = (n, c_in, h_in, w_in)
 
         # Store.
         self.in_shape_at_block.append(shape_ins)
@@ -316,7 +313,8 @@ def _forward(self, x):
                 ys.append(y)
             elif type(block) is darknet.RouteSlice:
                 slice = block
-                x = ys[slice.layer]
+                y_idx = slice.layer if slice.layer < 0 else slice.layer + 1
+                x = ys[y_idx]
                 y = slice(x)
                 ys.append(y)
             elif type(block) is darknet.RouteConcat2:
@@ -336,7 +334,7 @@ def _forward(self, x):
             elif type(block) is darknet.Shortcut:
                 shortcut = block
                 i1 = -1
-                if shortcut.self.from_layer < 0:
+                if shortcut.from_layer < 0:
                     i2 = shortcut.from_layer
                 else:
                     i2 = shortcut.from_layer + 1