MultiScaleDeformableAttention (#143)

RoyYang0714 · web-flow · commit 83ba216b58ef · 2025-08-05T23:26:28.000+02:00
* feat: Add MultiScaleDeformableAttention.

* fix: Refactor common api.

* feat: Change version.
diff --git a/pyproject.toml b/pyproject.toml
@@ -113,7 +113,7 @@ plugins = ["numpy.typing.mypy_plugin"]
 
 [project]
 name = "vis4d"
-version = "0.1.3"
+version = "0.1.4"
 authors = [{name = "Vis4D Team"}]
 description = "Vis4D Python package for Visual 4D scene understanding"
 readme = "README.md"
diff --git a/vis4d/__init__.py b/vis4d/__init__.py
@@ -7,7 +7,7 @@
 
 import logging
 
-__version__ = "0.1.3"
+__version__ = "0.1.4"
 
 _root_logger = logging.getLogger()
 _logger = logging.getLogger(__name__)
diff --git a/vis4d/common/__init__.py b/vis4d/common/__init__.py
@@ -1,37 +1 @@
 """Contains common functions and types that are used across modules."""
-
-from .typing import (
-    ArgsType,
-    DictStrAny,
-    DictStrArrNested,
-    GenericFunc,
-    ListAny,
-    LossesType,
-    MetricLogs,
-    ModelOutput,
-    NDArrayF32,
-    NDArrayF64,
-    NDArrayI64,
-    NDArrayNumber,
-    NDArrayUI8,
-    TorchCheckpoint,
-    TorchLossFunc,
-)
-
-__all__ = [
-    "DictStrAny",
-    "DictStrArrNested",
-    "ModelOutput",
-    "ArgsType",
-    "NDArrayF32",
-    "NDArrayF64",
-    "NDArrayI64",
-    "NDArrayUI8",
-    "NDArrayNumber",
-    "MetricLogs",
-    "TorchCheckpoint",
-    "LossesType",
-    "TorchLossFunc",
-    "GenericFunc",
-    "ListAny",
-]
diff --git a/vis4d/common/ckpt.py b/vis4d/common/ckpt.py
@@ -15,14 +15,14 @@
 from torch import nn
 from torch.hub import load_state_dict_from_url as load_url
 
-from vis4d.common import TorchCheckpoint
 from vis4d.common.distributed import (
     get_rank,
     get_world_size,
     is_module_wrapper,
     synchronize,
 )
 from vis4d.common.logging import rank_zero_info, rank_zero_warn
+from vis4d.common.typing import TorchCheckpoint
 
 CheckpointLoadFunc = Callable[
     [str, Union[str, torch.device, None]], TorchCheckpoint
diff --git a/vis4d/common/dict.py b/vis4d/common/dict.py
@@ -4,7 +4,7 @@
 
 from typing import Any
 
-from vis4d.common import DictStrAny
+from vis4d.common.typing import DictStrAny
 
 
 def flatten_dict(dictionary: DictStrAny, seperator: str) -> list[str]:
diff --git a/vis4d/common/distributed.py b/vis4d/common/distributed.py
@@ -17,7 +17,7 @@
 from torch.distributed import broadcast_object_list
 from torch.nn.parallel import DataParallel, DistributedDataParallel
 
-from vis4d.common import ArgsType, DictStrAny, GenericFunc
+from vis4d.common.typing import ArgsType, DictStrAny, GenericFunc
 
 
 # no coverage for these functions, since we don't unittest distributed setting
diff --git a/vis4d/data/datasets/base.py b/vis4d/data/datasets/base.py
@@ -12,7 +12,7 @@
 
 from torch.utils.data import Dataset as TorchDataset
 
-from vis4d.common import ArgsType
+from vis4d.common.typing import ArgsType
 from vis4d.data.io.base import DataBackend
 from vis4d.data.io.file import FileBackend
 from vis4d.data.typing import DictData
diff --git a/vis4d/data/datasets/coco.py b/vis4d/data/datasets/coco.py
@@ -11,7 +11,7 @@
 import pycocotools.mask as maskUtils
 from pycocotools.coco import COCO as COCOAPI
 
-from vis4d.common import ArgsType, DictStrAny
+from vis4d.common.typing import ArgsType, DictStrAny
 from vis4d.data.const import CommonKeys as K
 from vis4d.data.typing import DictData
 
diff --git a/vis4d/engine/callbacks/evaluator.py b/vis4d/engine/callbacks/evaluator.py
@@ -7,14 +7,14 @@
 
 import lightning.pytorch as pl
 
-from vis4d.common import ArgsType, MetricLogs
 from vis4d.common.distributed import (
     all_gather_object_cpu,
     broadcast,
     rank_zero_only,
     synchronize,
 )
 from vis4d.common.logging import rank_zero_info
+from vis4d.common.typing import ArgsType, MetricLogs
 from vis4d.data.typing import DictData
 from vis4d.eval.base import Evaluator
 
diff --git a/vis4d/engine/callbacks/logging.py b/vis4d/engine/callbacks/logging.py
@@ -7,10 +7,10 @@
 
 import lightning.pytorch as pl
 
-from vis4d.common import ArgsType, MetricLogs
 from vis4d.common.logging import rank_zero_info
 from vis4d.common.progress import compose_log_str
 from vis4d.common.time import Timer
+from vis4d.common.typing import ArgsType, MetricLogs
 
 from .base import Callback
 
diff --git a/vis4d/engine/callbacks/visualizer.py b/vis4d/engine/callbacks/visualizer.py
@@ -7,8 +7,8 @@
 
 import lightning.pytorch as pl
 
-from vis4d.common import ArgsType
 from vis4d.common.distributed import broadcast, synchronize
+from vis4d.common.typing import ArgsType
 from vis4d.vis.base import Visualizer
 
 from .base import Callback
diff --git a/vis4d/engine/callbacks/yolox_callbacks.py b/vis4d/engine/callbacks/yolox_callbacks.py
@@ -13,7 +13,6 @@
 from torch.nn.modules.batchnorm import _NormBase
 from torch.utils.data import DataLoader
 
-from vis4d.common import ArgsType, DictStrAny
 from vis4d.common.distributed import (
     all_reduce_dict,
     broadcast,
@@ -22,6 +21,7 @@
     synchronize,
 )
 from vis4d.common.logging import rank_zero_info, rank_zero_warn
+from vis4d.common.typing import ArgsType, DictStrAny
 from vis4d.data.const import CommonKeys as K
 from vis4d.data.data_pipe import DataPipe
 from vis4d.op.detect.yolox import YOLOXHeadLoss
diff --git a/vis4d/engine/run.py b/vis4d/engine/run.py
@@ -9,8 +9,8 @@
 from absl import app  # pylint: disable=no-name-in-module
 from torch.utils.collect_env import get_pretty_env_info
 
-from vis4d.common import ArgsType
 from vis4d.common.logging import dump_config, rank_zero_info, setup_logger
+from vis4d.common.typing import ArgsType
 from vis4d.common.util import set_tf32
 from vis4d.config import instantiate_classes
 from vis4d.config.typing import ExperimentConfig
diff --git a/vis4d/engine/trainer.py b/vis4d/engine/trainer.py
@@ -11,9 +11,9 @@
 from lightning.pytorch.loggers.wandb import WandbLogger
 from lightning.pytorch.strategies.ddp import DDPStrategy
 
-from vis4d.common import ArgsType
 from vis4d.common.imports import TENSORBOARD_AVAILABLE
 from vis4d.common.logging import rank_zero_info
+from vis4d.common.typing import ArgsType
 
 
 class PLTrainer(Trainer):
diff --git a/vis4d/eval/common/seg.py b/vis4d/eval/common/seg.py
@@ -5,9 +5,13 @@
 import numpy as np
 from terminaltables import AsciiTable
 
-from vis4d.common import MetricLogs
 from vis4d.common.array import array_to_numpy
-from vis4d.common.typing import ArrayLike, NDArrayI64, NDArrayNumber
+from vis4d.common.typing import (
+    ArrayLike,
+    MetricLogs,
+    NDArrayI64,
+    NDArrayNumber,
+)
 from vis4d.eval.base import Evaluator
 
 
diff --git a/vis4d/eval/shift/multitask_writer.py b/vis4d/eval/shift/multitask_writer.py
@@ -11,10 +11,14 @@
 import numpy as np
 from PIL import Image
 
-from vis4d.common import MetricLogs
 from vis4d.common.array import array_to_numpy
 from vis4d.common.imports import SCALABEL_AVAILABLE
-from vis4d.common.typing import ArrayLike, GenericFunc, NDArrayNumber
+from vis4d.common.typing import (
+    ArrayLike,
+    GenericFunc,
+    MetricLogs,
+    NDArrayNumber,
+)
 from vis4d.data.datasets.shift import shift_det_map
 from vis4d.data.io import DataBackend, ZipBackend
 from vis4d.eval.base import Evaluator
diff --git a/vis4d/model/cls/vit.py b/vis4d/model/cls/vit.py
@@ -6,8 +6,8 @@
 import torch
 from torch import nn
 
-from vis4d.common import ArgsType
 from vis4d.common.ckpt import load_model_checkpoint
+from vis4d.common.typing import ArgsType
 from vis4d.op.base.vit import VisionTransformer, ViT_PRESET
 
 from .common import ClsOut
diff --git a/vis4d/model/detect/retinanet.py b/vis4d/model/detect/retinanet.py
@@ -4,8 +4,8 @@
 
 from torch import Tensor, nn
 
-from vis4d.common import LossesType
 from vis4d.common.ckpt import load_model_checkpoint
+from vis4d.common.typing import LossesType
 from vis4d.op.base.resnet import ResNet
 from vis4d.op.box.anchor import AnchorGenerator
 from vis4d.op.box.box2d import scale_and_clip_boxes
diff --git a/vis4d/op/base/csp_darknet.py b/vis4d/op/base/csp_darknet.py
@@ -12,7 +12,8 @@
 from torch import nn
 from torch.nn.modules.batchnorm import _BatchNorm
 
-from vis4d.op.layer import Conv2d, CSPLayer
+from vis4d.op.layer.conv2d import Conv2d
+from vis4d.op.layer.csp_layer import CSPLayer
 
 
 class Focus(nn.Module):
diff --git a/vis4d/op/base/vit.py b/vis4d/op/base/vit.py
@@ -6,7 +6,9 @@
 from timm.models import named_apply
 from torch import nn
 
-from ..layer import PatchEmbed, TransformerBlock
+from vis4d.op.layer.patch_embed import PatchEmbed
+from vis4d.op.layer.transformer import TransformerBlock
+
 from .base import BaseModel
 
 
diff --git a/vis4d/op/box/poolers/roi_pooler.py b/vis4d/op/box/poolers/roi_pooler.py
@@ -8,7 +8,7 @@
 import torch
 from torchvision.ops import roi_align, roi_pool
 
-from vis4d.common import ArgsType
+from vis4d.common.typing import ArgsType
 
 from .base import RoIPooler
 from .utils import assign_boxes_to_levels, boxes_to_tensor
diff --git a/vis4d/op/box/samplers/combined.py b/vis4d/op/box/samplers/combined.py
@@ -5,7 +5,7 @@
 import torch
 from torch import Tensor
 
-from vis4d.common import ArgsType
+from vis4d.common.typing import ArgsType
 
 from ..box2d import non_intersection, random_choice
 from ..matchers.base import MatchResult
diff --git a/vis4d/op/box/samplers/random.py b/vis4d/op/box/samplers/random.py
@@ -4,7 +4,7 @@
 
 import torch
 
-from vis4d.common import ArgsType
+from vis4d.common.typing import ArgsType
 
 from ..matchers.base import MatchResult
 from .base import Sampler, SamplingResult
diff --git a/vis4d/op/detect/dense_anchor.py b/vis4d/op/detect/dense_anchor.py
@@ -8,7 +8,7 @@
 import torch.nn.functional as F
 from torch import Tensor, nn
 
-from vis4d.common import TorchLossFunc
+from vis4d.common.typing import TorchLossFunc
 from vis4d.op.box.anchor import AnchorGenerator, anchor_inside_image
 from vis4d.op.box.encoder import DeltaXYWHBBoxEncoder
 from vis4d.op.box.matchers import Matcher
diff --git a/vis4d/op/detect/faster_rcnn.py b/vis4d/op/detect/faster_rcnn.py
@@ -17,9 +17,9 @@
     match_and_sample_proposals,
 )
 
-from ..typing import Proposals, Targets
 from .rcnn import RCNNHead, RCNNOut
 from .rpn import RPN2RoI, RPNHead, RPNOut
+from .typing import Proposals, Targets
 
 
 class FRCNNOut(NamedTuple):
diff --git a/vis4d/op/detect/mask_rcnn.py b/vis4d/op/detect/mask_rcnn.py
@@ -13,7 +13,7 @@
 from vis4d.op.box.poolers import MultiScaleRoIAlign
 from vis4d.op.mask.util import paste_masks_in_image, remove_overlap
 
-from ..typing import Proposals, Targets
+from .typing import Proposals, Targets
 
 
 class MaskRCNNHeadOut(NamedTuple):
diff --git a/vis4d/op/detect/rcnn.py b/vis4d/op/detect/rcnn.py
@@ -15,7 +15,7 @@
 from vis4d.op.box.encoder import DeltaXYWHBBoxDecoder, DeltaXYWHBBoxEncoder
 from vis4d.op.box.poolers import MultiScaleRoIAlign
 from vis4d.op.detect.common import DetOut
-from vis4d.op.layer import add_conv_branch
+from vis4d.op.layer.conv2d import add_conv_branch
 from vis4d.op.layer.weight_init import kaiming_init, normal_init, xavier_init
 from vis4d.op.loss.common import l1_loss
 from vis4d.op.loss.reducer import SumWeightedLoss
diff --git a/vis4d/op/detect/rpn.py b/vis4d/op/detect/rpn.py
@@ -16,11 +16,11 @@
 from vis4d.op.box.encoder import DeltaXYWHBBoxDecoder, DeltaXYWHBBoxEncoder
 from vis4d.op.box.matchers import Matcher, MaxIoUMatcher
 from vis4d.op.box.samplers import RandomSampler, Sampler
+from vis4d.op.layer.conv2d import Conv2d
 from vis4d.op.loss.common import l1_loss
 
-from ..layer import Conv2d
-from ..typing import Proposals
 from .dense_anchor import DenseAnchorHeadLoss, DenseAnchorHeadLosses
+from .typing import Proposals
 
 
 class RPNOut(NamedTuple):
diff --git a/vis4d/op/detect/typing.py b/vis4d/op/detect/typing.py
@@ -1,4 +1,4 @@
-"""Vis4D op typing."""
+"""Detect op typing."""
 
 from __future__ import annotations
 
diff --git a/vis4d/op/detect/yolox.py b/vis4d/op/detect/yolox.py
@@ -14,13 +14,13 @@
 from torch import Tensor, nn
 from torchvision.ops import batched_nms
 
-from vis4d.common import TorchLossFunc
 from vis4d.common.distributed import reduce_mean
+from vis4d.common.typing import TorchLossFunc
 from vis4d.op.box.anchor import MlvlPointGenerator
 from vis4d.op.box.encoder import YOLOXBBoxDecoder
 from vis4d.op.box.matchers import SimOTAMatcher
 from vis4d.op.box.samplers import PseudoSampler
-from vis4d.op.layer import Conv2d
+from vis4d.op.layer.conv2d import Conv2d
 from vis4d.op.layer.weight_init import bias_init_with_prob
 from vis4d.op.loss import IoULoss
 from vis4d.op.loss.reducer import SumWeightedLoss
diff --git a/vis4d/op/detect3d/qd_3dt.py b/vis4d/op/detect3d/qd_3dt.py
@@ -18,7 +18,7 @@
     match_and_sample_proposals,
 )
 from vis4d.op.geometry.rotation import generate_rotation_output
-from vis4d.op.layer import Conv2d, add_conv_branch
+from vis4d.op.layer.conv2d import Conv2d, add_conv_branch
 from vis4d.op.layer.weight_init import kaiming_init, xavier_init
 from vis4d.op.loss.base import Loss
 from vis4d.op.loss.common import rotation_loss, smooth_l1_loss
diff --git a/vis4d/op/fpp/dla_up.py b/vis4d/op/fpp/dla_up.py
@@ -11,9 +11,10 @@
 import torch
 from torch import nn
 
-from vis4d.common import NDArrayI64
+from vis4d.common.typing import NDArrayI64
+from vis4d.op.layer.conv2d import Conv2d
+from vis4d.op.layer.deform_conv import DeformConv
 
-from ..layer import Conv2d, DeformConv
 from .base import FeaturePyramidProcessing
 
 
diff --git a/vis4d/op/fpp/yolox_pafpn.py b/vis4d/op/fpp/yolox_pafpn.py
@@ -10,7 +10,8 @@
 import torch
 from torch import nn
 
-from vis4d.op.layer import Conv2d, CSPLayer
+from vis4d.op.layer.conv2d import Conv2d
+from vis4d.op.layer.csp_layer import CSPLayer
 
 from .base import FeaturePyramidProcessing
 
diff --git a/vis4d/op/layer/__init__.py b/vis4d/op/layer/__init__.py
diff --git a/vis4d/op/layer/attention.py b/vis4d/op/layer/attention.py
diff --git a/vis4d/op/layer/mlp.py b/vis4d/op/layer/mlp.py
diff --git a/vis4d/op/layer/ms_deform_attn.py b/vis4d/op/layer/ms_deform_attn.py
diff --git a/vis4d/op/track/qdtrack.py b/vis4d/op/track/qdtrack.py
diff --git a/vis4d/state/track3d/motion/lstm_3d.py b/vis4d/state/track3d/motion/lstm_3d.py
diff --git a/vis4d/zoo/run.py b/vis4d/zoo/run.py
