Modality dispatch

mart/attack/adversary.py

@@ -149,7 +149,7 @@ def configure_gradient_clipping(
 
         if self.gradient_modifier:
             for group in optimizer.param_groups:
-                self.gradient_modifier(group["params"])
+                self.gradient_modifier(group)
 
     @silent()
     def fit(self, input, target, *, model: Callable):

mart/attack/composer.py

@@ -11,6 +11,8 @@
 
 import torch
 
+from ..utils.modality_dispatch import DEFAULT_MODALITY, modality_dispatch
+
 
 class Composer(abc.ABC):
     def __call__(
@@ -78,3 +80,34 @@ def compose(self, perturbation, *, input, target):
         masked_perturbation = perturbation * mask
 
         return input + masked_perturbation
+
+
+class Modality(Composer):
+    def __init__(self, **modality_method):
+        self.modality_method = modality_method
+
+    def __call__(
+        self,
+        perturbation: torch.Tensor | Iterable[torch.Tensor] | Iterable[dict[str, torch.Tensor]],
+        *,
+        input: torch.Tensor | Iterable[torch.Tensor] | Iterable[dict[str, torch.Tensor]],
+        target: torch.Tensor | Iterable[torch.Tensor] | Iterable[dict[str, Any]],
+        **kwargs,
+    ) -> torch.Tensor | Iterable[torch.Tensor]:
+        return modality_dispatch(
+            input,
+            data=perturbation,
+            target=target,
+            modality_func=self.compose,
+            modality=DEFAULT_MODALITY,
+        )
+
+    def compose(
+        self,
+        perturbation: torch.Tensor,
+        *,
+        input: torch.Tensor,
+        target: torch.Tensor | dict[str, Any],
+        modality: str,
+    ) -> torch.Tensor:
+        return self.modality_method[modality](perturbation, input=input, target=target)

mart/attack/enforcer.py

@@ -11,6 +11,8 @@
 
 import torch
 
+from ..utils.modality_dispatch import modality_dispatch
+
 __all__ = ["Enforcer"]
 
 
@@ -96,36 +98,41 @@ def verify(self, input_adv, *, input, target):
 
 
 class Enforcer:
-    def __init__(self, constraints: dict[str, Constraint]) -> None:
-        self.constraints = list(constraints.values())  # intentionally ignore keys
+    def __init__(self, **modality_constraints: dict[str, dict[str, Constraint]]) -> None:
+        self.modality_constraints = {}
+
+        for modality, constraints in modality_constraints.items():
+            # Intentionally ignore keys after modality.
+            # The keys are there for combining constraints easily in Hydra.
+            self.modality_constraints[modality] = constraints.values()
 
     @torch.no_grad()
     def __call__(
         self,
-        input_adv: torch.Tensor | Iterable[torch.Tensor],
+        input_adv: torch.Tensor | Iterable[torch.Tensor] | Iterable[dict[str, torch.Tensor]],
         *,
-        input: torch.Tensor | Iterable[torch.Tensor],
+        input: torch.Tensor | Iterable[torch.Tensor] | Iterable[dict[str, torch.Tensor]],
         target: torch.Tensor | Iterable[torch.Tensor] | Iterable[dict[str, Any]],
         **kwargs,
     ):
-        if isinstance(input_adv, torch.Tensor) and isinstance(input, torch.Tensor):
-            self.enforce(input_adv, input=input, target=target)
-
-        elif (
-            isinstance(input_adv, Iterable)
-            and isinstance(input, Iterable)  # noqa: W503
-            and isinstance(target, Iterable)  # noqa: W503
-        ):
-            for input_adv_i, input_i, target_i in zip(input_adv, input, target):
-                self.enforce(input_adv_i, input=input_i, target=target_i)
+        # The default modality is set to "constraints", so that it is backward compatible with existing configs.
+        modality_dispatch(
+            input,
+            data=input_adv,
+            target=target,
+            modality_func=self.enforce,
+            modality="constraints",
+        )
 
     @torch.no_grad()
     def enforce(
         self,
-        input_adv: torch.Tensor,
+        input_adv: torch.Tensor | Iterable[torch.Tensor],
         *,
-        input: torch.Tensor,
-        target: torch.Tensor | dict[str, Any],
+        input: torch.Tensor | Iterable[torch.Tensor],
+        target: torch.Tensor | Iterable[torch.Tensor] | Iterable[dict[str, Any]],
+        modality: str,
     ):
-        for constraint in self.constraints:
+
+        for constraint in self.modality_constraints[modality]:
             constraint(input_adv, input=input, target=target)

mart/attack/gradient_modifier.py

@@ -6,17 +6,20 @@
 
 from __future__ import annotations
 
-from typing import Iterable
+from typing import Any
 
 import torch
 
+from ..utils.modality_dispatch import DEFAULT_MODALITY
+
 __all__ = ["GradientModifier"]
 
 
 class GradientModifier:
     """Gradient modifier base class."""
 
-    def __call__(self, parameters: torch.Tensor | Iterable[torch.Tensor]) -> None:
+    def __call__(self, param_group: dict[str, Any]) -> None:
+        parameters = param_group["params"]
         if isinstance(parameters, torch.Tensor):
             parameters = [parameters]
 
@@ -43,3 +46,15 @@ def __init__(self, p: int | float):
     def modify_(self, parameter: torch.Tensor) -> None:
         p_norm = torch.norm(parameter.grad.detach(), p=self.p)
         parameter.grad.detach().div_(p_norm)
+
+
+class Modality(GradientModifier):
+    def __init__(self, **modality_method):
+        if len(modality_method) == 0:
+            modality_method = {DEFAULT_MODALITY: self.modify_}
+
+        self.modality_method_ = modality_method
+
+    def __call__(self, param_group: dict[str, Any]) -> None:
+        modality = param_group["modality"] if "modality" in param_group else DEFAULT_MODALITY
+        self.modality_method_[modality](param_group)

mart/attack/initializer.py

@@ -6,7 +6,7 @@
 
 from __future__ import annotations
 
-from typing import Iterable
+from typing import Any
 
 import torch
 import torchvision
@@ -21,11 +21,15 @@ class Initializer:
     """Initializer base class."""
 
     @torch.no_grad()
-    def __call__(self, parameters: torch.Tensor | Iterable[torch.Tensor]) -> None:
-        if isinstance(parameters, torch.Tensor):
-            parameters = [parameters]
-
-        [self.initialize_(parameter) for parameter in parameters]
+    def __call__(
+        self,
+        parameter: torch.Tensor,
+        *,
+        input: torch.Tensor | None = None,
+        target: torch.Tensor | dict[str, Any] | None = None,
+    ) -> None:
+        # Accept input and target from modality_dispatch().
+        self.initialize_(parameter)
 
     @torch.no_grad()
     def initialize_(self, parameter: torch.Tensor) -> None:

mart/attack/perturber.py

@@ -11,6 +11,11 @@
 import torch
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
 
+from ..utils.modality_dispatch import (
+    DEFAULT_MODALITY,
+    ModalityParameterDict,
+    modality_dispatch,
+)
 from .projector import Projector
 
 if TYPE_CHECKING:
@@ -23,8 +28,8 @@ class Perturber(torch.nn.Module):
     def __init__(
         self,
         *,
-        initializer: Initializer,
-        projector: Projector | None = None,
+        initializer: Initializer | dict[str, Initializer],
+        projector: Projector | dict[str, Projector] | None = None,
     ):
         """_summary_
 
@@ -34,8 +39,17 @@ def __init__(
         """
         super().__init__()
 
+        projector = projector or Projector()
+
+        # Modality-specific objects.
+        # Backward compatibility, in case modality is unknown, and not given in input.
+        if not isinstance(initializer, dict):
+            initializer = {DEFAULT_MODALITY: initializer}
+        if not isinstance(projector, dict):
+            projector = {DEFAULT_MODALITY: projector}
+
         self.initializer_ = initializer
-        self.projector_ = projector or Projector()
+        self.projector_ = projector
 
         self.perturbation = None
 
@@ -65,6 +79,10 @@ def create_from_tensor(tensor):
                 return torch.nn.Parameter(
                     torch.empty_like(tensor, dtype=torch.float, requires_grad=True)
                 )
+            elif isinstance(tensor, dict):
+                return ModalityParameterDict(
+                    {modality: create_from_tensor(t) for modality, t in tensor.items()}
+                )
             elif isinstance(tensor, Iterable):
                 return torch.nn.ParameterList([create_from_tensor(t) for t in tensor])
             else:
@@ -76,7 +94,13 @@ def create_from_tensor(tensor):
             self.perturbation = create_from_tensor(input)
 
         # Always (re)initialize perturbation.
-        self.initializer_(self.perturbation)
+        modality_dispatch(
+            input,
+            data=self.perturbation,
+            target=None,
+            modality_func=self.initializer_,
+            modality=DEFAULT_MODALITY,
+        )
 
     def named_parameters(self, *args, **kwargs):
         if self.perturbation is None:
@@ -90,12 +114,18 @@ def parameters(self, *args, **kwargs):
 
         return super().parameters(*args, **kwargs)
 
-    def forward(self, **batch):
+    def forward(self, *, input, target, **batch):
         if self.perturbation is None:
             raise MisconfigurationException(
                 "You need to call the configure_perturbation before forward."
             )
 
-        self.projector_(self.perturbation, **batch)
+        modality_dispatch(
+            input,
+            data=self.perturbation,
+            target=target,
+            modality_func=self.projector_,
+            modality=DEFAULT_MODALITY,
+        )
 
         return self.perturbation

mart/attack/projector.py

@@ -6,7 +6,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Iterable
+from typing import Any
 
 import torch
 
@@ -17,33 +17,21 @@ class Projector:
     @torch.no_grad()
     def __call__(
         self,
-        perturbation: torch.Tensor | Iterable[torch.Tensor],
+        perturbation: torch.Tensor,
         *,
-        input: torch.Tensor | Iterable[torch.Tensor],
-        target: torch.Tensor | Iterable[torch.Tensor] | Iterable[dict[str, Any]],
+        input: torch.Tensor,
+        target: torch.Tensor | dict[str, Any],
         **kwargs,
     ) -> None:
-        if isinstance(perturbation, torch.Tensor) and isinstance(input, torch.Tensor):
-            self.project_(perturbation, input=input, target=target)
-
-        elif (
-            isinstance(perturbation, Iterable)
-            and isinstance(input, Iterable)  # noqa: W503
-            and isinstance(target, Iterable)  # noqa: W503
-        ):
-            for perturbation_i, input_i, target_i in zip(perturbation, input, target):
-                self.project_(perturbation_i, input=input_i, target=target_i)
-
-        else:
-            raise NotImplementedError
+        self.project_(perturbation, input=input, target=target)
 
     @torch.no_grad()
     def project_(
         self,
-        perturbation: torch.Tensor | Iterable[torch.Tensor],
+        perturbation: torch.Tensor,
         *,
-        input: torch.Tensor | Iterable[torch.Tensor],
-        target: torch.Tensor | Iterable[torch.Tensor] | Iterable[dict[str, Any]],
+        input: torch.Tensor,
+        target: torch.Tensor | dict[str, Any],
     ) -> None:
         pass
 
@@ -57,10 +45,10 @@ def __init__(self, projectors: list[Projector]):
     @torch.no_grad()
     def __call__(
         self,
-        perturbation: torch.Tensor | Iterable[torch.Tensor],
+        perturbation: torch.Tensor,
         *,
-        input: torch.Tensor | Iterable[torch.Tensor],
-        target: torch.Tensor | Iterable[torch.Tensor] | Iterable[dict[str, Any]],
+        input: torch.Tensor,
+        target: torch.Tensor | dict[str, Any],
         **kwargs,
     ) -> None:
         for projector in self.projectors:

mart/configs/attack/composer/modality.yaml

@@ -0,0 +1 @@
+_target_: mart.attack.composer.Modality

mart/configs/attack/enforcer/default.yaml

@@ -1,2 +1,3 @@
 _target_: mart.attack.Enforcer
-constraints: ???
+# FIXME: Hydra does not detect modality-aware constraints defined as sub-components.
+# constraints: ???

mart/configs/attack/gradient_modifier/modality.yaml

@@ -0,0 +1 @@
+_target_: mart.attack.gradient_modifier.Modality

mart/configs/attack/object_detection_rgb_mask_adversary.yaml

@@ -0,0 +1,28 @@
+defaults:
+  - adversary
+  - perturber: default
+  - perturber/initializer@perturber.initializer.rgb: constant
+  - perturber/projector@perturber.projector.rgb: mask_range
+  - /optimizer@optimizer: sgd
+  - gain: rcnn_training_loss
+  - objective: zero_ap
+  - gradient_modifier: modality
+  - gradient_modifier@gradient_modifier.rgb: sign
+  - composer: modality
+  - composer@composer.rgb: overlay
+  - enforcer: default
+  - enforcer/constraints@enforcer.rgb: [mask, pixel_range]
+
+# Make a 5-step attack for the demonstration purpose.
+optimizer:
+  # Though we only use modality-aware hyper-params, the config requires a value for optimizer.lr.
+  lr: 0
+  rgb:
+    lr: 55
+
+max_iters: 5
+
+perturber:
+  initializer:
+    rgb:
+      constant: 127