diff --git a/src/secmlt/adv/evasion/base_evasion_attack.py b/src/secmlt/adv/evasion/base_evasion_attack.py index 4e291503..1fe8a380 100644 --- a/src/secmlt/adv/evasion/base_evasion_attack.py +++ b/src/secmlt/adv/evasion/base_evasion_attack.py @@ -2,7 +2,7 @@ import importlib.util from abc import abstractmethod -from typing import Literal +from typing import Literal, Union import torch from secmlt.adv.backends import Backends @@ -144,14 +144,14 @@ def get_backends() -> set[str]: class BaseEvasionAttack: """Base class for evasion attacks.""" - def __call__(self, model: BaseModel, data_loader: DataLoader) -> DataLoader: + def __call__(self, model: Union[BaseModel, list[BaseModel]], data_loader: DataLoader) -> DataLoader: """ Compute the attack against the model, using the input data. Parameters ---------- - model : BaseModel - Model to test. + model : BaseModel | list[BaseModel] + Model or list of models to test. data_loader : DataLoader Test dataloader. diff --git a/src/secmlt/adv/evasion/modular_attack.py b/src/secmlt/adv/evasion/modular_attack.py index 93f2e0f5..963cbf07 100644 --- a/src/secmlt/adv/evasion/modular_attack.py +++ b/src/secmlt/adv/evasion/modular_attack.py @@ -158,8 +158,8 @@ def forward_loss( Parameters ---------- - model : BaseModel - Model used by the attack run. + model : BaseModel | list[BaseModel] + Model(s) used by the attack run. x : torch.Tensor Input sample. target : torch.Tensor @@ -170,9 +170,17 @@ def forward_loss( tuple[torch.Tensor, torch.Tensor] Output scores and loss. """ - scores = model.decision_function(x) - target = target.to(scores.device) - losses = self.loss_function(scores, target) + if isinstance(model, list): + scores_list = [m.decision_function(x) for m in model] + scores = torch.stack(scores_list).mean(dim=0) + losses = torch.stack( + [self.loss_function(scores, target.to(scores.device)) + for scores in scores_list] + ).mean(dim=0) + else: + scores = model.decision_function(x) + losses = self.loss_function(scores, + target.to(scores.device)) return scores, losses def _run(