WIP: add a partially working CoreML conversion.

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

@@ -4,7 +4,7 @@
     ConvBNA,
     MaxPool,
     RouteSlice,
-    RouteConcat,
+    RouteConcat2,
     Shortcut,
     Upsample,
     Yolo

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

@@ -1,5 +1,6 @@
 import logging
 from pathlib import Path
+from typing import Optional
 
 import numpy as np
 
@@ -15,7 +16,8 @@
 
 class Network(nn.Module):
 
-    def __init__(self, cfg: darknet.Config, inference=True):
+    def __init__(self, cfg: darknet.Config, inference=True,
+                 up_to_layer: Optional[int]=None):
         super(Network, self).__init__()
 
         input_shape = (
@@ -26,6 +28,7 @@ def __init__(self, cfg: darknet.Config, inference=True):
         )
         self.in_shape_at_block = [input_shape]
         self.out_shape_at_block = [input_shape]
+        self.up_to_layer = up_to_layer
         self.model = self.create_network(cfg)
 
     def input_shape(self):
@@ -45,7 +48,10 @@ def create_network(self, cfg: darknet.Config):
         upsample_id = 0
         yolo_id = 0
 
-        for block in cfg._model:
+        for (i, block) in enumerate(cfg._model):
+            if self.up_to_layer is not None and i >= self.up_to_layer:
+                break
+
             layer_name = list(block.keys())[0]
             layer_params = block[layer_name]
 
@@ -71,13 +77,17 @@ def create_network(self, cfg: darknet.Config):
 
         return model
 
-    def load_convolutional_weights(self, weights_file: Path, version='v4'):
+    def load_convolutional_weights(self,
+                                   weights_file: Path,
+                                   version: str='v4'):
         if version != 'v4':
             raise NotImplementedError
 
         weight_loader = v4.NetworkWeightLoader(weights_file)
 
         for block_idx, block in enumerate(self.model):
+            if self.up_to_layer is not None and block_idx >= self.up_to_layer:
+                break
             if type(block) is not darknet.ConvBNA:
                 continue
 
@@ -125,7 +135,8 @@ def load_convolutional_weights(self, weights_file: Path, version='v4'):
 
         logging.debug(f'weight loader cursor = {weight_loader._cursor}')
         logging.debug(f'weights num elements = {weight_loader._weights.size}')
-        assert weight_loader._cursor == weight_loader._weights.size
+        if self.up_to_layer is None:
+            assert weight_loader._cursor == weight_loader._weights.size
 
     def _read_weights(self, shape, weight_loader):
         return weight_loader.read(shape[0]).reshape(shape)
@@ -200,7 +211,7 @@ def _append_route(self, model, layer_params, route_id):
             self.out_shape_at_block.append(shape_out)
 
             # Append the route-concat block.
-            model.append(darknet.RouteConcat(layers, route_id))
+            model.append(darknet.RouteConcat2(layers, route_id))
             logging.debug(
                 f'[Route {route_id}] (Concat): '
                 f'{shape_ins} -> {shape_out}\n'
@@ -276,7 +287,7 @@ def _forward(self, x):
                 x = ys[slice.layer]
                 y = slice(x)
                 ys.append(y)
-            elif type(block) is darknet.RouteConcat:
+            elif type(block) is darknet.RouteConcat2:
                 concat = block
 
                 ids = [l if l < 0 else l + 1 for l in concat.layers]
@@ -300,5 +311,8 @@ def _forward(self, x):
         return ys, boxes
 
     def forward(self, x):
-        _, boxes = self._forward(x)
-        return boxes
+        ys, boxes = self._forward(x)
+        if self.up_to_layer is None:
+            return boxes[0]
+        else:
+            return ys[-1]

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

@@ -71,42 +71,48 @@ class MaxPool(nn.Module):
 
     def __init__(self, kernel_size, stride):
         super(MaxPool, self).__init__()
-        self.kernel_size = kernel_size
+        self.kernel_size = [kernel_size, kernel_size]
         self.stride = stride
 
+        from functools import reduce
+        from operator import __add__
+        self.zero_pad_2d = nn.ZeroPad2d(reduce(__add__,
+            [(k // 2 + (k - 2 * (k // 2)) - 1, k // 2) for k in self.kernel_size[::-1]]))
+
     def forward(self, x):
-        # Let's use shortcut variables.
-        s = self.stride
-        # Get the height and width of the input signal.
-        h, w = x.shape[2:]
-
-        # The kernel radius is calculated as
-        r = self.kernel_size // 2
-
-        # We calculate the easy part of the padding.
-        p_left = r - 1 if self.kernel_size % 2 == 0 else r
-        p_top = r - 1 if self.kernel_size % 2 == 0 else r
-
-        # Now moving on the trickiest part of the padding.
-        #
-        # The max pool layers collects (w // s) x (h // s) samples from the
-        # input signal.
-        #
-        # If we reason in 1D, the samples are located at:
-        #   0, s, 2s, 3s, ... , (w // s) * s
-        #
-        # The input signal is extended spatially so that it contains the
-        # following sample points.
-        x_last = ((w - 1) // s) * s + r
-        y_last = ((h - 1) // s) * s + r
-        # Therefore the last two padding are
-        p_right = 0 if x_last == w - 1 else x_last - w + 1
-        p_bottom = 0 if y_last == h - 1 else y_last - h + 1
-
-        # Apply the padding with negative infinity value.
-        pad_size = (p_left, p_right, p_top, p_bottom)
-        x_padded = F.pad(x, pad_size, mode='constant', value=-float('inf'))
-        # print(f'x_padded = \n{x_padded}')
+        x_padded = self.zero_pad_2d(x)
+        # # Let's use shortcut variables.
+        # s = self.stride
+        # # Get the height and width of the input signal.
+        # h, w = x.shape[2:]
+
+        # # The kernel radius is calculated as
+        # r = self.kernel_size // 2
+
+        # # We calculate the easy part of the padding.
+        # p_left = r - 1 if self.kernel_size % 2 == 0 else r
+        # p_top = r - 1 if self.kernel_size % 2 == 0 else r
+
+        # # Now moving on the trickiest part of the padding.
+        # #
+        # # The max pool layers collects (w // s) x (h // s) samples from the
+        # # input signal.
+        # #
+        # # If we reason in 1D, the samples are located at:
+        # #   0, s, 2s, 3s, ... , (w // s) * s
+        # #
+        # # The input signal is extended spatially so that it contains the
+        # # following sample points.
+        # x_last = ((w - 1) // s) * s + r
+        # y_last = ((h - 1) // s) * s + r
+        # # Therefore the last two padding are
+        # p_right = 0 if x_last == w - 1 else x_last - w + 1
+        # p_bottom = 0 if y_last == h - 1 else y_last - h + 1
+
+        # # Apply the padding with negative infinity value.
+        # pad_size = (p_left, p_right, p_top, p_bottom)
+        # x_padded = F.pad(x, pad_size, mode='constant', value=-float('inf'))
+        # # print(f'x_padded = \n{x_padded}')
 
         # Apply the spatial max-pool function.
         y = F.max_pool2d(x_padded, self.kernel_size, self.stride)
@@ -140,17 +146,17 @@ def forward(self, x):
             return x[:, c1:c2, :, :]
 
 
-class RouteConcat(nn.Module):
+class RouteConcat2(nn.Module):
 
     def __init__(self, layers: [int], id: Optional[int] = None):
-        super(RouteConcat, self).__init__()
+        super(RouteConcat2, self).__init__()
         self.layers = layers
         self.id = id
 
-    def forward(self, *xs):
-        if len(self.layers) != len(xs):
-            raise RuntimeError(f"This route-concat layer requires {self.layers} inputs")
-        return torch.cat(xs, 1)
+    def forward(self, x1, x2):
+        if len(self.layers) != 2:
+            raise RuntimeError(f"This route-concat layer requires 2 inputs")
+        return torch.cat((x1, x2), 1)
 
 
 class Shortcut(nn.Module):

python/oddkiva/shakti/inference/darknet/v4/test/test_darknet_layers.py

@@ -23,16 +23,16 @@ def test_mish():
 
 def test_maxpool():
     for sz in range(1, 10):
-        # print(f'\nsz = {sz}')
+        print(f'\nsz = {sz}')
         w, h = sz, sz
         x_np = np.arange(sz ** 2).reshape(1, 1, h, w).astype(np.float32)
         x = torch.tensor(x_np)
-        # print(x)
+        print(x)
 
         max_pool = MaxPool(2, 2)
 
         y = max_pool(x)
-        # print(y)
+        print(y)
 
         # Just check the dimensions for now.
         hy, wy = y.shape[2:]

python/oddkiva/shakti/inference/darknet/v4/test/test_yolov4_config_parsing.py

@@ -7,6 +7,8 @@
 
 import torch
 
+import coremltools as ct
+
 import oddkiva.sara as sara
 import oddkiva.shakti.inference.darknet as darknet
 
@@ -58,40 +60,73 @@ def read_image(path: Path, yolo_net: darknet.Network):
     return image_tensor
 
 
-def test_yolo_v4_tiny_cfg():
+# def test_yolo_v4_tiny_cfg():
+#     yolo_cfg = darknet.Config()
+#     yolo_cfg.read(YOLO_V4_TINY_CFG_PATH)
+#     assert yolo_cfg._model is not None
+# 
+#     yolo_net = darknet.Network(yolo_cfg)
+#     yolo_net.load_convolutional_weights(YOLO_V4_TINY_WEIGHT_PATH);
+#     yolo_net.eval()
+# 
+#     in_tensor = read_image(DOG_IMAGE_PATH, yolo_net)
+#     in_tensor_saved = yolo_out_tensor(0)
+#     err = torch.norm(in_tensor - in_tensor_saved).item()
+#     logging.info(f'input err = {err}')
+#     assert err < 1e-12
+# 
+#     ys, boxes = yolo_net._forward(in_tensor)
+#     assert len(ys) == len(yolo_net.model) + 1
+#     assert torch.equal(ys[0], in_tensor)
+# 
+#     for i in range(1, len(yolo_net.model)):
+#         block = yolo_net.model[i]
+#         out_tensor_saved = yolo_out_tensor(i + 1)
+#         out_tensor_computed = ys[i + 1]
+# 
+#         assert out_tensor_saved.shape == out_tensor_computed.shape
+# 
+#         err = torch.norm(out_tensor_computed - out_tensor_saved).item()
+#         logging.info(f'[{i}] err = {err} for {block}')
+#         assert err < 3e-3
+# 
+# 
+#     ids = [31, 38]
+#     boxes_true = [yolo_out_tensor(id) for id in ids]
+# 
+#     for i, b, b_true in zip(ids, boxes, boxes_true):
+#         err = torch.norm(b - b_true)
+#         logging.info(f'[{i}] err = {err} for {yolo_net.model[i-1]}')
+#         assert err < 1e-4
+
+
+def test_yolo_v4_tiny_coreml_conversion():
     yolo_cfg = darknet.Config()
     yolo_cfg.read(YOLO_V4_TINY_CFG_PATH)
     assert yolo_cfg._model is not None
 
-    yolo_net = darknet.Network(yolo_cfg)
+    layer_idx = 37
+    yolo_net = darknet.Network(yolo_cfg, up_to_layer=layer_idx)
     yolo_net.load_convolutional_weights(YOLO_V4_TINY_WEIGHT_PATH);
+    yolo_net.eval()
+
+    print(yolo_net.model)
 
     in_tensor = read_image(DOG_IMAGE_PATH, yolo_net)
     in_tensor_saved = yolo_out_tensor(0)
-    err = torch.norm(in_tensor - in_tensor_saved).item()
-    logging.info(f'input err = {err}')
-    assert err < 1e-12
-
-    ys, boxes = yolo_net._forward(in_tensor)
-    assert len(ys) == len(yolo_net.model) + 1
-    assert torch.equal(ys[0], in_tensor)
-
-    for i in range(1, len(yolo_net.model)):
-        block = yolo_net.model[i]
-        out_tensor_saved = yolo_out_tensor(i + 1)
-        out_tensor_computed = ys[i + 1]
 
-        assert out_tensor_saved.shape == out_tensor_computed.shape
+    with torch.inference_mode():
+        traced_model = torch.jit.trace(yolo_net, in_tensor)
+        outs = traced_model(in_tensor)
 
-        err = torch.norm(out_tensor_computed - out_tensor_saved).item()
-        logging.info(f'[{i}] err = {err} for {block}')
-        assert err < 3e-3
+        ct_outs = [ct.TensorType(name=f'yolo_{i}') for i, _ in enumerate(outs)]
 
+        model = ct.convert(
+            traced_model,
+            inputs=[ct.TensorType(shape=in_tensor.shape)],
+            debug=True
+        )
 
-    ids = [31, 38]
-    boxes_true = [yolo_out_tensor(id) for id in ids]
+        model.save('/Users/oddkiva/Desktop/yolo.mlpackage')
 
-    for i, b, b_true in zip(ids, boxes, boxes_true):
-        err = torch.norm(b - b_true)
-        logging.info(f'[{i}] err = {err} for {yolo_net.model[i-1]}')
-        assert err < 1e-4
+    import IPython; IPython.embed()

requirements.txt

@@ -6,16 +6,13 @@ ipython
 ipdb
 ipdbplugin
 
-numpy
-scipy
-sympy
-
 Pillow
 
+torch==2.1.1
+torchvision==0.16.1
+torchaudio==2.1.1
+
 coremltools
-torch
-torchvision
-torchaudio
 
 PyOpenGL
-PySide2
+# PySide2