ENH: better printing.

oddkiva · Dec 21, 2023 · a2cc65b · a2cc65b
1 parent f715dd4
commit a2cc65b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/python/oddkiva/shakti/inference/darknet/network.py b/python/oddkiva/shakti/inference/darknet/network.py
@@ -77,7 +77,7 @@ def _append_conv(self, model, layer_params, conv_id):
 
         # Append the convolutional block to the model.
         model.append(darknet.ConvBNA(c_in, layer_params, conv_id))
-        print(f'[Conv{conv_id}]: '
+        print(f'[Conv {conv_id}] '
               f'{self.in_shape_at_block[-1]} -> {self.out_shape_at_block[-1]}')
 
     def _append_route(self, model, layer_params, route_id):
@@ -100,7 +100,7 @@ def _append_route(self, model, layer_params, route_id):
             # Append the route-slice block.
             model.append(darknet.RouteSlice(
                 layers[0], groups, group_id, route_id))
-            print(f'[Route{route_id}] (Slide): '
+            print(f'[Route {route_id}] (Slide): '
                   f'{self.in_shape_at_block[-1]} -> {self.out_shape_at_block[-1]}')
         else:
             # Fetch all the input shapes.
@@ -117,8 +117,9 @@ def _append_route(self, model, layer_params, route_id):
 
             # Append the route-concat block.
             model.append(darknet.RouteConcat(layers, route_id))
-            print(f'[Route{route_id}] (Concat): '
+            print(f'[Route {route_id}] (Concat): '
                   f'{self.in_shape_at_block[-1]} -> {self.out_shape_at_block[-1]}')
+            print(f'                    layers = {layers}')
 
     def _append_max_pool(self, model, layer_params, max_pool_id):
         # Extract the input shape
@@ -135,7 +136,7 @@ def _append_max_pool(self, model, layer_params, max_pool_id):
         self.out_shape_at_block.append(shape_out)
 
         model.append(darknet.MaxPool(size, stride))
-        print(f'[MaxPool{max_pool_id}] '
+        print(f'[MaxPool {max_pool_id}] '
               f'{self.in_shape_at_block[-1]} -> {self.out_shape_at_block[-1]}')
 
     def _append_upsample(self, model, layer_params, upsample_id):
@@ -152,7 +153,7 @@ def _append_upsample(self, model, layer_params, upsample_id):
         self.out_shape_at_block.append(shape_out)
 
         model.append(darknet.Upsample(stride))
-        print(f'[Upsample{upsample_id}] '
+        print(f'[Upsample {upsample_id}] '
               f'{self.in_shape_at_block[-1]} -> {self.out_shape_at_block[-1]}')
 
     def _append_yolo(self, model, layer_params, yolo_id):
@@ -167,7 +168,7 @@ def _append_yolo(self, model, layer_params, yolo_id):
         self.out_shape_at_block.append(shape_out)
 
         model.append(darknet.Yolo(layer_params))
-        print(f'[YOLO{yolo_id}] '
+        print(f'[YOLO {yolo_id}] '
               f'{self.in_shape_at_block[-1]} -> {self.out_shape_at_block[-1]}')
 
     def load_convolutional_weights(self, conv, weights_file: Path):

diff --git a/python/oddkiva/shakti/inference/darknet/torch_layers.py b/python/oddkiva/shakti/inference/darknet/torch_layers.py
@@ -206,10 +206,10 @@ def forward(self, x):
         y[:, xs, :, :] = self.alpha * nn.Sigmoid(x[:, xs, :, :]) + self.beta
         y[:, ys, :, :] = self.alpha * nn.Sigmoid(x[:, ys, :, :]) + self.beta
 
-        # Class probabilities.
+        # P[object] and P[class|object] probabilities.
         for box in range(0, 3):
             c_begin = box * num_box_features + 4
             c_end = (box + 1) * num_box_features 
             y[:, c_begin:c_end, :, :] = nn.Sigmoid(x[:, c_begin:c_end, :, :])
 
-        return y;
+        return y