Revision models.detection.yolo

pl_bolts/models/detection/yolo/yolo_config.py

@@ -170,7 +170,7 @@ def _create_layer(config: dict, num_inputs: List[int]) -> Tuple[nn.Module, int]:
     return create_func[config["type"]](config, num_inputs)
 
 
-def _create_convolutional(config, num_inputs):
+def _create_convolutional(config: dict, num_inputs: int) -> Tuple[nn.Module, int]:
     module = nn.Sequential()
 
     batch_normalize = config.get("batch_normalize", False)
@@ -206,13 +206,13 @@ def _create_convolutional(config, num_inputs):
     return module, config["filters"]
 
 
-def _create_maxpool(config, num_inputs):
+def _create_maxpool(config: dict, num_inputs: int) -> Tuple[nn.Module, int]:
     padding = (config["size"] - 1) // 2
     module = nn.MaxPool2d(config["size"], config["stride"], padding)
     return module, num_inputs[-1]
 
 
-def _create_route(config, num_inputs):
+def _create_route(config: dict, num_inputs: int) -> Tuple[nn.Module, int]:
     num_chunks = config.get("groups", 1)
     chunk_idx = config.get("group_id", 0)
 
@@ -228,17 +228,17 @@ def _create_route(config, num_inputs):
     return module, num_outputs
 
 
-def _create_shortcut(config, num_inputs):
+def _create_shortcut(config: dict, num_inputs: int) -> Tuple[nn.Module, int]:
     module = yolo_layers.ShortcutLayer(config["from"])
     return module, num_inputs[-1]
 
 
-def _create_upsample(config, num_inputs):
+def _create_upsample(config: dict, num_inputs: int) -> Tuple[nn.Module, int]:
     module = nn.Upsample(scale_factor=config["stride"], mode="nearest")
     return module, num_inputs[-1]
 
 
-def _create_yolo(config, num_inputs):
+def _create_yolo(config: dict, num_inputs: int) -> Tuple[nn.Module, int]:
     # The "anchors" list alternates width and height.
     anchor_dims = config["anchors"]
     anchor_dims = [(anchor_dims[i], anchor_dims[i + 1]) for i in range(0, len(anchor_dims), 2)]

pl_bolts/models/detection/yolo/yolo_layers.py

@@ -1,4 +1,4 @@
-from typing import Callable, Dict, List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
@@ -464,7 +464,7 @@ def _calculate_losses(
 class Mish(nn.Module):
     """Mish activation."""
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         return x * torch.tanh(nn.functional.softplus(x))
 
 
@@ -483,7 +483,7 @@ def __init__(self, source_layers: List[int], num_chunks: int, chunk_idx: int) ->
         self.num_chunks = num_chunks
         self.chunk_idx = chunk_idx
 
-    def forward(self, x, outputs):
+    def forward(self, x, outputs: List[Union[Tensor, None]]) -> Tensor:
         chunks = [torch.chunk(outputs[layer], self.num_chunks, dim=1)[self.chunk_idx] for layer in self.source_layers]
         return torch.cat(chunks, dim=1)
 
@@ -500,5 +500,5 @@ def __init__(self, source_layer: int) -> None:
         super().__init__()
         self.source_layer = source_layer
 
-    def forward(self, x, outputs):
+    def forward(self, x, outputs: List[Union[Tensor, None]]) -> Tensor:
         return outputs[-1] + outputs[self.source_layer]

tests/models/test_detection.py

@@ -79,8 +79,9 @@ def test_fasterrcnn_pyt_module_bbone_train(tmpdir):
     trainer.fit(model, train_dl, valid_dl)
 
 
-def test_yolo(tmpdir):
-    config_path = Path(TEST_ROOT) / "data" / "yolo.cfg"
+@pytest.mark.parametrize("config", [("yolo"), ("yolo_giou")])
+def test_yolo(config):
+    config_path = Path(TEST_ROOT) / "data" / f"{config}.cfg"
     config = YOLOConfiguration(config_path)
     model = YOLO(config.get_network())
 

tests/models/yolo/unit/test_yolo_config.py

@@ -1,15 +1,11 @@
-from pathlib import Path
-
 import pytest
 
 from pl_bolts.models.detection.yolo.yolo_config import (
-    YOLOConfiguration,
     _create_convolutional,
     _create_maxpool,
     _create_shortcut,
     _create_upsample,
 )
-from tests import TEST_ROOT
 
 
 @pytest.mark.parametrize(
@@ -22,7 +18,7 @@
     ],
 )
 def test_create_convolutional(config):
-    conv = _create_convolutional(config, [3])[0]
+    conv, _ = _create_convolutional(config, [3])
 
     assert conv.conv.out_channels == config["filters"]
     assert conv.conv.kernel_size == (config["size"], config["size"])
@@ -63,7 +59,7 @@ def test_create_convolutional(config):
 )
 def test_create_maxpool(config):
     pad_size = (config["size"] - 1) // 2
-    maxpool = _create_maxpool(config, [3])[0]
+    maxpool, _ = _create_maxpool(config, [3])
 
     assert maxpool.kernel_size == config["size"]
     assert maxpool.stride == config["stride"]
@@ -78,7 +74,7 @@ def test_create_maxpool(config):
     ],
 )
 def test_create_shortcut(config):
-    shortcut = _create_shortcut(config, [3])[0]
+    shortcut, _ = _create_shortcut(config, [3])
 
     assert shortcut.source_layer == config["from"]
 
@@ -91,13 +87,6 @@ def test_create_shortcut(config):
     ],
 )
 def test_create_upsample(config):
-    upsample = _create_upsample(config, [3])[0]
+    upsample, _ = _create_upsample(config, [3])
 
     assert upsample.scale_factor == float(config["stride"])
-
-
-@pytest.mark.parametrize("config", [("yolo"), ("yolo_giou")])
-def test_yolo_config(config):
-    config_path = Path(TEST_ROOT) / "data" / f"{config}.cfg"
-    config = YOLOConfiguration(config_path)
-    config.get_network()