Add runtime type annotations to rfdetr util helpers (#551)

* Add runtime type annotations (re-authored from Copilot)
* Standardize docstring argument and return type formatting across utility functions
* Apply suggestions from code review


---------

Co-authored-by: jirka <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 3 people

rfdetr/util/benchmark.py

@@ -32,19 +32,19 @@
 import time
 
 
-from typing import Any, Callable, List, Optional, Union
+from typing import Any, Callable, Dict, List, Sequence, Union
 from numbers import Number
 
 Handle = Callable[[List[Any], List[Any]], Union[typing.Counter[str], Number]]
 
 
-def get_shape(val: object) -> typing.List[int]:
+def get_shape(val: Any) -> typing.List[int]:
     """
     Get the shapes from a jit value object.
     Args:
-        val (torch._C.Value): jit value object.
+        val: jit value object.
     Returns:
-        list(int): return a list of ints.
+        return a list of ints.
     """
     if val.isCompleteTensor():  # pyre-ignore
         r = val.type().sizes()  # pyre-ignore
@@ -64,17 +64,17 @@ def get_shape(val: object) -> typing.List[int]:
 
 
 def addmm_flop_jit(
-    inputs: typing.List[object], outputs: typing.List[object]
+    inputs: typing.List[Any], outputs: typing.List[Any]
 ) -> typing.Counter[str]:
     """
     This method counts the flops for fully connected layers with torch script.
     Args:
-        inputs (list(torch._C.Value)): The input shape in the form of a list of
+        inputs: The input shape in the form of a list of
             jit object.
-        outputs (list(torch._C.Value)): The output shape in the form of a list
+        outputs: The output shape in the form of a list
             of jit object.
     Returns:
-        Counter: A Counter dictionary that records the number of flops for each
+        A Counter dictionary that records the number of flops for each
             operation.
     """
     # Count flop for nn.Linear
@@ -91,7 +91,7 @@ def addmm_flop_jit(
     return flop_counter
 
 
-def bmm_flop_jit(inputs, outputs):
+def bmm_flop_jit(inputs: typing.List[Any], outputs: typing.List[Any]) -> Counter[str]:
     # Count flop for nn.Linear
     # inputs is a list of length 3.
     input_shapes = [get_shape(v) for v in inputs]
@@ -106,7 +106,7 @@ def bmm_flop_jit(inputs, outputs):
     return flop_counter
 
 
-def basic_binary_op_flop_jit(inputs, outputs, name):
+def basic_binary_op_flop_jit(inputs: typing.List[Any], outputs: typing.List[Any], name: str) -> Counter[str]:
     input_shapes = [get_shape(v) for v in inputs]
     # for broadcasting
     input_shapes = [s[::-1] for s in input_shapes]
@@ -116,29 +116,34 @@ def basic_binary_op_flop_jit(inputs, outputs, name):
     return flop_counter
 
 
-def rsqrt_flop_jit(inputs, outputs):
+def rsqrt_flop_jit(inputs: typing.List[Any], outputs: typing.List[Any]) -> Counter[str]:
     input_shapes = [get_shape(v) for v in inputs]
     flop = prod(input_shapes[0]) * 2
     flop_counter = Counter({"rsqrt": flop})
     return flop_counter
 
 
-def dropout_flop_jit(inputs, outputs):
+def dropout_flop_jit(inputs: typing.List[Any], outputs: typing.List[Any]) -> Counter[str]:
     input_shapes = [get_shape(v) for v in inputs[:1]]
     flop = prod(input_shapes[0])
     flop_counter = Counter({"dropout": flop})
     return flop_counter
 
 
-def softmax_flop_jit(inputs, outputs):
+def softmax_flop_jit(inputs: typing.List[Any], outputs: typing.List[Any]) -> Counter[str]:
     # from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/profiler/internal/flops_registry.py
     input_shapes = [get_shape(v) for v in inputs[:1]]
     flop = prod(input_shapes[0]) * 5
     flop_counter = Counter({"softmax": flop})
     return flop_counter
 
 
-def _reduction_op_flop_jit(inputs, outputs, reduce_flops=1, finalize_flops=0):
+def _reduction_op_flop_jit(
+    inputs: typing.List[Any],
+    outputs: typing.List[Any],
+    reduce_flops: int = 1,
+    finalize_flops: int = 0,
+) -> int:
     input_shapes = [get_shape(v) for v in inputs]
     output_shapes = [get_shape(v) for v in outputs]
 
@@ -161,11 +166,11 @@ def conv_flop_count(
     This method counts the flops for convolution. Note only multiplication is
     counted. Computation for addition and bias is ignored.
     Args:
-        x_shape (list(int)): The input shape before convolution.
-        w_shape (list(int)): The filter shape.
-        out_shape (list(int)): The output shape after convolution.
+        x_shape: The input shape before convolution.
+        w_shape: The filter shape.
+        out_shape: The output shape after convolution.
     Returns:
-        Counter: A Counter dictionary that records the number of flops for each
+        A Counter dictionary that records the number of flops for each
             operation.
     """
     batch_size, Cin_dim, Cout_dim = x_shape[0], w_shape[1], out_shape[1]
@@ -177,17 +182,17 @@ def conv_flop_count(
 
 
 def conv_flop_jit(
-    inputs: typing.List[object], outputs: typing.List[object]
+    inputs: typing.List[Any], outputs: typing.List[Any]
 ) -> typing.Counter[str]:
     """
     This method counts the flops for convolution using torch script.
     Args:
-        inputs (list(torch._C.Value)): The input shape in the form of a list of
+        inputs: The input shape in the form of a list of
             jit object before convolution.
-        outputs (list(torch._C.Value)): The output shape in the form of a list
+        outputs: The output shape in the form of a list
             of jit object after convolution.
     Returns:
-        Counter: A Counter dictionary that records the number of flops for each
+        A Counter dictionary that records the number of flops for each
             operation.
     """
     # Inputs of Convolution should be a list of length 12. They represent:
@@ -206,18 +211,18 @@ def conv_flop_jit(
 
 
 def einsum_flop_jit(
-    inputs: typing.List[object], outputs: typing.List[object]
+    inputs: typing.List[Any], outputs: typing.List[Any]
 ) -> typing.Counter[str]:
     """
     This method counts the flops for the einsum operation. We currently support
     two einsum operations: "nct,ncp->ntp" and "ntg,ncg->nct".
     Args:
-        inputs (list(torch._C.Value)): The input shape in the form of a list of
+        inputs: The input shape in the form of a list of
             jit object before einsum.
-        outputs (list(torch._C.Value)): The output shape in the form of a list
+        outputs: The output shape in the form of a list
             of jit object after einsum.
     Returns:
-        Counter: A Counter dictionary that records the number of flops for each
+        A Counter dictionary that records the number of flops for each
             operation.
     """
     # Inputs of einsum should be a list of length 2.
@@ -254,17 +259,17 @@ def einsum_flop_jit(
 
 
 def matmul_flop_jit(
-    inputs: typing.List[object], outputs: typing.List[object]
+    inputs: typing.List[Any], outputs: typing.List[Any]
 ) -> typing.Counter[str]:
     """
     This method counts the flops for matmul.
     Args:
-        inputs (list(torch._C.Value)): The input shape in the form of a list of
+        inputs: The input shape in the form of a list of
             jit object before matmul.
-        outputs (list(torch._C.Value)): The output shape in the form of a list
+        outputs: The output shape in the form of a list
             of jit object after matmul.
     Returns:
-        Counter: A Counter dictionary that records the number of flops for each
+        A Counter dictionary that records the number of flops for each
             operation.
     """
 
@@ -287,17 +292,17 @@ def matmul_flop_jit(
 
 
 def batchnorm_flop_jit(
-    inputs: typing.List[object], outputs: typing.List[object]
+    inputs: typing.List[Any], outputs: typing.List[Any]
 ) -> typing.Counter[str]:
     """
     This method counts the flops for batch norm.
     Args:
-        inputs (list(torch._C.Value)): The input shape in the form of a list of
+        inputs: The input shape in the form of a list of
             jit object before batch norm.
-        outputs (list(torch._C.Value)): The output shape in the form of a list
+        outputs: The output shape in the form of a list
             of jit object after batch norm.
     Returns:
-        Counter: A Counter dictionary that records the number of flops for each
+        A Counter dictionary that records the number of flops for each
             operation.
     """
     # Inputs[0] contains the shape of the input.
@@ -470,26 +475,26 @@ def elementwise_flop(inputs: List[Any], outputs: List[Any]) -> Number:
 
 def flop_count(
     model: nn.Module,
-    inputs: typing.Tuple[object, ...],
-    whitelist: typing.Union[typing.List[str], None] = None,
-    customized_ops: typing.Union[typing.Dict[str, typing.Callable], None] = None,
+    inputs: typing.Tuple[Any, ...],
+    whitelist: typing.Optional[typing.List[str]] = None,
+    customized_ops: typing.Optional[typing.Dict[str, typing.Callable]] = None,
 ) -> typing.DefaultDict[str, float]:
     """
     Given a model and an input to the model, compute the Gflops of the given
     model. Note the input should have a batch size of 1.
     Args:
-        model (nn.Module): The model to compute flop counts.
-        inputs (tuple): Inputs that are passed to `model` to count flops.
+        model: The model to compute flop counts.
+        inputs: Inputs that are passed to `model` to count flops.
             Inputs need to be in a tuple.
-        whitelist (list(str)): Whitelist of operations that will be counted. It
+        whitelist: Whitelist of operations that will be counted. It
             needs to be a subset of _SUPPORTED_OPS. By default, the function
             computes flops for all supported operations.
-        customized_ops (dict(str,Callable)) : A dictionary contains customized
+        customized_ops: A dictionary contains customized
             operations and their flop handles. If customized_ops contains an
             operation in _SUPPORTED_OPS, then the default handle in
              _SUPPORTED_OPS will be overwritten.
     Returns:
-        defaultdict: A dictionary that records the number of gflops for each
+        A dictionary that records the number of gflops for each
             operation.
     """
     # Copy _SUPPORTED_OPS to flop_count_ops.
@@ -557,13 +562,13 @@ def flop_count(
     return final_count
 
 
-def warmup(model, inputs, N=10):
+def warmup(model: torch.nn.Module, inputs: Any, N: int = 10) -> None:
     for i in range(N):
         out = model(inputs)
     torch.cuda.synchronize()
 
 
-def measure_time(model, inputs, N=10):
+def measure_time(model: torch.nn.Module, inputs: Any, N: int = 10) -> float:
     warmup(model, inputs)
     s = time.time()
     for i in range(N):
@@ -573,7 +578,7 @@ def measure_time(model, inputs, N=10):
     return t
 
 
-def fmt_res(data):
+def fmt_res(data: np.ndarray) -> Dict[str, float]:
     # return data.mean(), data.std(), data.min(), data.max()
     return {
         "mean": data.mean(),
@@ -583,7 +588,7 @@ def fmt_res(data):
     }
 
 
-def benchmark(model, dataset, output_dir):
+def benchmark(model: torch.nn.Module, dataset: Sequence[Any], output_dir: Any) -> Dict[str, Any]:
     print("Get model size, FLOPs, and FPS")
     # import pdb; pdb.set_trace()
     _outputs = {}

rfdetr/util/box_ops.py

@@ -14,27 +14,36 @@
 """
 Utilities for bounding box manipulation and GIoU.
 """
+from typing import Tuple
+
 import torch
 import torch.nn.functional as F
 from torchvision.ops.boxes import box_area
 
 
-def box_cxcywh_to_xyxy(x):
+def box_cxcywh_to_xyxy(x: torch.Tensor) -> torch.Tensor:
     x_c, y_c, w, h = x.unbind(-1)
     b = [(x_c - 0.5 * w.clamp(min=0.0)), (y_c - 0.5 * h.clamp(min=0.0)),
          (x_c + 0.5 * w.clamp(min=0.0)), (y_c + 0.5 * h.clamp(min=0.0))]
     return torch.stack(b, dim=-1)
 
 
-def box_xyxy_to_cxcywh(x):
+def box_xyxy_to_cxcywh(x: torch.Tensor) -> torch.Tensor:
     x0, y0, x1, y1 = x.unbind(-1)
     b = [(x0 + x1) / 2, (y0 + y1) / 2,
          (x1 - x0), (y1 - y0)]
     return torch.stack(b, dim=-1)
 
 
 # modified from torchvision to also return the union
-def box_iou(boxes1, boxes2):
+def box_iou(boxes1: torch.Tensor, boxes2: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Returns:
+        iou: the NxM matrix containing the pairwise
+            IoU values for every element in boxes1 and boxes2
+        union: the NxM matrix containing the pairwise
+            union values for every element in boxes1 and boxes2
+    """
     area1 = box_area(boxes1)
     area2 = box_area(boxes2)
 
@@ -50,7 +59,7 @@ def box_iou(boxes1, boxes2):
     return iou, union
 
 
-def generalized_box_iou(boxes1, boxes2):
+def generalized_box_iou(boxes1: torch.Tensor, boxes2: torch.Tensor) -> torch.Tensor:
     """
     Generalized IoU from https://giou.stanford.edu/
 
@@ -72,7 +81,7 @@ def generalized_box_iou(boxes1, boxes2):
     return iou - (area - union) / area
 
 
-def masks_to_boxes(masks):
+def masks_to_boxes(masks: torch.Tensor) -> torch.Tensor:
     """Compute the bounding boxes around the provided masks
 
     The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions.
@@ -99,7 +108,7 @@ def masks_to_boxes(masks):
     return torch.stack([x_min, y_min, x_max, y_max], 1)
 
 
-def batch_dice_loss(inputs: torch.Tensor, targets: torch.Tensor):
+def batch_dice_loss(inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
     """
     Compute the DICE loss, similar to generalized IOU for masks
     Args:
@@ -122,7 +131,7 @@ def batch_dice_loss(inputs: torch.Tensor, targets: torch.Tensor):
 )  # type: torch.jit.ScriptModule
 
 
-def batch_sigmoid_ce_loss(inputs: torch.Tensor, targets: torch.Tensor):
+def batch_sigmoid_ce_loss(inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
     """
     Args:
         inputs: A float tensor of arbitrary shape.

rfdetr/util/drop_scheduler.py

@@ -5,9 +5,17 @@
 # ------------------------------------------------------------------------
 """util for drop scheduler."""
 import numpy as np
+from typing import Literal
 
 
-def drop_scheduler(drop_rate, epochs, niter_per_ep, cutoff_epoch=0, mode='standard', schedule='constant'):
+def drop_scheduler(
+    drop_rate: float,
+    epochs: int,
+    niter_per_ep: int,
+    cutoff_epoch: int = 0,
+    mode: Literal['standard', 'early', 'late'] = 'standard',
+    schedule: Literal['constant', 'linear'] = 'constant',
+) -> np.ndarray:
     """drop scheduler"""
     assert mode in ['standard', 'early', 'late']
     if mode == 'standard':