Repository: bethgelab/foolbox (language: python)

Issue:
DifferentiableModelWrapper equivalent for Expectation Over Transformation?
Following what was discussed in #191, I'm trying to implement Expectation Over Transformation in a similar fashion. The original poster mentioned using DifferentiableModelWrapper to do so, but that is obviously no longer available. The only wrapper class that seems to be available currently is a ThresholdingWrapper. I was wondering if anyone could help point me in the right direction for accomplishing this. 

Ideally, I would think implementing EOT would involve having a wrapper around a PyTorchModel object that basically averages the gradients from that model then returns the final, averaged gradients as per EOT.  


Discussion / hints:
This is certainly something we should add back to Foolbox 3 in one way or the other. Unfortunately, I probably won't have to time to work on this in the near future but a PR would be very welcome.
Ideally, it should be a solution that works for all frameworks, but I would need to think more about it to know what the right approach is.
Hi, 

are there any updates on this? 
By checking the documentation I did not find any implementation of Expectation over Transformation. 

Do you have any suggestions on how to implement it?  
Right now there is no implementation for this in foolbox. However, this should be pretty easy to implement. You could start by creating a class that inherits from the framework-specific model classes (e.g., [PyTorchModel](https://github.com/bethgelab/foolbox/blob/master/foolbox/models/pytorch.py#L20)) and then modify the internal `def _model` definition to call the inner model multiple times and the output over those repetitions.

An even better approach would be to create a new class that takes a PyTorch/JaxModel as an argument and averages over multiple inner calls when computing the forward pass.

If you want to try this out, it would be great if you could create a PR for this. 

Fix patch:
diff --git a/docs/modules/models.rst b/docs/modules/models.rst
index 4f743512f..f12ba7c94 100644
--- a/docs/modules/models.rst
+++ b/docs/modules/models.rst
@@ -22,6 +22,7 @@ Wrappers
    :nosignatures:
 
    TransformBoundsWrapper
+   ExpectationOverTransformationWrapper
 
 Detailed description
 --------------------
@@ -43,3 +44,6 @@ Detailed description
 
 .. autoclass:: TransformBoundsWrapper
    :members:
+
+.. autoclass:: ExpectationOverTransformationWrapper
+   :members:
diff --git a/examples/eot_attack_pytorch_resnet18.py b/examples/eot_attack_pytorch_resnet18.py
new file mode 100644
index 000000000..ce68061d6
--- /dev/null
+++ b/examples/eot_attack_pytorch_resnet18.py
@@ -0,0 +1,129 @@
+#!/usr/bin/env python3
+"""
+A simple example that demonstrates how to run Expectation over Transformation
+coupled with any attack, on a Resnet-18 PyTorch model.
+"""
+from typing import Any
+
+import torch
+from torch import Tensor
+import torchvision.models as models
+import torchvision.transforms as transforms
+import eagerpy as ep
+from foolbox import PyTorchModel, accuracy, samples
+from foolbox.attacks import LinfPGD
+from foolbox.models import ExpectationOverTransformationWrapper
+
+
+class RandomizedResNet18(torch.nn.Module):
+    def __init__(self) -> None:
+
+        super().__init__()
+
+        # base model
+        self.model = models.resnet18(pretrained=True)
+
+        # random apply rotation
+        self.transforms = transforms.RandomRotation(degrees=25)
+
+    def forward(self, x: Tensor) -> Any:
+
+        # random transform
+        x = self.transforms(x)
+
+        return self.model(x)
+
+
+def main() -> None:
+    # instantiate a model (could also be a TensorFlow or JAX model)
+    model = models.resnet18(pretrained=True).eval()
+    preprocessing = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], axis=-3)
+    fmodel = PyTorchModel(model, bounds=(0, 1), preprocessing=preprocessing)
+
+    # get data and test the model
+    # wrapping the tensors with ep.astensors is optional, but it allows
+    # us to work with EagerPy tensors in the following
+    images, labels = ep.astensors(*samples(fmodel, dataset="imagenet", batchsize=16))
+
+    print("Testing attack on the base model (no transformations applied)")
+    clean_acc = accuracy(fmodel, images, labels)
+    print(f"clean accuracy:  {clean_acc * 100:.1f} %")
+
+    # apply an attack with different eps
+    attack = LinfPGD()
+    epsilons = [
+        0.0,
+        0.0002,
+        0.0005,
+        0.0008,
+        0.001,
+        0.0015,
+        0.002,
+        0.003,
+        0.01,
+        0.02,
+        0.03,
+        0.1,
+        0.3,
+        0.5,
+        1.0,
+    ]
+
+    raw_advs, clipped_advs, success = attack(fmodel, images, labels, epsilons=epsilons)
+
+    # calculate and report the robust accuracy (the accuracy of the model when
+    # it is attacked)
+    robust_accuracy = 1 - success.float32().mean(axis=-1)
+    print("robust accuracy for perturbations with")
+    for eps, acc in zip(epsilons, robust_accuracy):
+        print(f"  Linf norm ≤ {eps:<6}: {acc.item() * 100:4.1f} %")
+
+    # Let's apply the same LinfPGD attack, but on a model with random transformations
+    rand_model = RandomizedResNet18().eval()
+    fmodel = PyTorchModel(rand_model, bounds=(0, 1), preprocessing=preprocessing)
+    seed = 1111
+
+    print("#" * 40)
+    print("Testing attack on the randomized model (random rotation applied)")
+
+    # Note: accuracy may slightly decrease, depending on seed
+    torch.manual_seed(seed)
+    clean_acc = accuracy(fmodel, images, labels)
+    print(f"clean accuracy:  {clean_acc * 100:.1f} %")
+
+    # test the base attack on the randomized model
+    print("robust accuracy for perturbations with")
+    for eps in epsilons:
+
+        # reset seed to have the same perturbations in each attack
+        torch.manual_seed(seed)
+        _, _, success = attack(fmodel, images, labels, epsilons=eps)
+
+        # calculate and report the robust accuracy
+        # the attack is performing worse on the randomized models, since gradient computation is affected!
+        robust_accuracy = 1 - success.float32().mean(axis=-1)
+        print(f"  Linf norm ≤ {eps:<6}: {robust_accuracy.item() * 100:4.1f} %")
+
+    # Now, Let's use Expectation Over Transformation to counter the randomization
+    eot_model = ExpectationOverTransformationWrapper(fmodel, n_steps=16)
+
+    print("#" * 40)
+    print("Testing EoT attack on the randomized model (random crop applied)")
+    torch.manual_seed(seed)
+    clean_acc = accuracy(eot_model, images, labels)
+    print(f"clean accuracy:  {clean_acc * 100:.1f} %")
+
+    print("robust accuracy for perturbations with")
+    for eps in epsilons:
+        # reset seed to have the same perturbations in each attack
+        torch.manual_seed(seed)
+        _, _, success = attack(eot_model, images, labels, epsilons=eps)
+
+        # calculate and report the robust accuracy
+        # with EoT, the base attack is working again!
+        robust_accuracy = 1 - success.float32().mean(axis=-1)
+        print(f"  Linf norm ≤ {eps:<6}: {robust_accuracy.item() * 100:4.1f} %")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/foolbox/models/__init__.py b/foolbox/models/__init__.py
index bb6eef236..1947ae5a3 100644
--- a/foolbox/models/__init__.py
+++ b/foolbox/models/__init__.py
@@ -7,3 +7,4 @@
 from .numpy import NumPyModel  # noqa: F401
 
 from .wrappers import ThresholdingWrapper  # noqa: F401
+from .wrappers import ExpectationOverTransformationWrapper  # noqa: F401
diff --git a/foolbox/models/wrappers.py b/foolbox/models/wrappers.py
index 390be6b4d..7201b9509 100644
--- a/foolbox/models/wrappers.py
+++ b/foolbox/models/wrappers.py
@@ -21,3 +21,29 @@ def __call__(self, inputs: T) -> T:
         y = ep.where(x < self._threshold, min_, max_).astype(x.dtype)
         z = self._model(y)
         return restore_type(z)
+
+
+class ExpectationOverTransformationWrapper(Model):
+    def __init__(self, model: Model, n_steps: int = 16):
+        self._model = model
+        self._n_steps = n_steps
+
+    @property
+    def bounds(self) -> Bounds:
+        return self._model.bounds
+
+    def __call__(self, inputs: T) -> T:
+
+        x, restore_type = ep.astensor_(inputs)
+
+        for i in range(self._n_steps):
+            z_t = self._model(x)
+
+            if i == 0:
+                z = z_t.expand_dims(0)
+            else:
+                z = ep.concatenate([z, z_t.expand_dims(0)], axis=0)
+
+        z = z.mean(0)
+
+        return restore_type(z)


Test patch:
diff --git a/tests/test_eot_wrapper.py b/tests/test_eot_wrapper.py
new file mode 100644
index 000000000..dd22605ae
--- /dev/null
+++ b/tests/test_eot_wrapper.py
@@ -0,0 +1,53 @@
+import pytest
+
+import eagerpy as ep
+
+from foolbox import accuracy
+from foolbox.attacks import (
+    LinfBasicIterativeAttack,
+    L1BasicIterativeAttack,
+    L2BasicIterativeAttack,
+)
+from foolbox.models import ExpectationOverTransformationWrapper
+from foolbox.types import L2, Linf
+
+from conftest import ModeAndDataAndDescription
+
+
+def test_eot_wrapper(
+    fmodel_and_data_ext_for_attacks: ModeAndDataAndDescription,
+) -> None:
+
+    (fmodel, x, y), real, low_dimensional_input = fmodel_and_data_ext_for_attacks
+
+    if isinstance(x, ep.NumPyTensor):
+        pytest.skip()
+
+    # test clean accuracy when wrapping EoT
+    x = (x - fmodel.bounds.lower) / (fmodel.bounds.upper - fmodel.bounds.lower)
+    fmodel = fmodel.transform_bounds((0, 1))
+    acc = accuracy(fmodel, x, y)
+
+    rand_model = ExpectationOverTransformationWrapper(fmodel, n_steps=4)
+    rand_acc = accuracy(rand_model, x, y)
+    assert acc - rand_acc == 0
+
+    # test with base attacks
+    # (accuracy should not change, since fmodel is not random)
+    attacks = (
+        L1BasicIterativeAttack(),
+        L2BasicIterativeAttack(),
+        LinfBasicIterativeAttack(),
+    )
+    epsilons = (5000.0, L2(50.0), Linf(1.0))
+
+    for attack, eps in zip(attacks, epsilons):
+
+        # acc on standard model
+        advs, _, _ = attack(fmodel, x, y, epsilons=eps)
+        adv_acc = accuracy(fmodel, advs, y)
+
+        # acc on eot model
+        advs, _, _ = attack(rand_model, x, y, epsilons=eps)
+        r_adv_acc = accuracy(rand_model, advs, y)
+        assert adv_acc - r_adv_acc == 0