Cannot compute the gradient of non-differentiable PreprocessorPyTorch

[Lodour]

Describe the bug
PyTorchClassifier cannot compute the gradient of a non-differentiable PreprocessorPyTorch.

While the document of PreprocessorPyTorch says that we need to define estimate_forward for BPDA, I find that estimate_forward is never invoked for PyTorch-based preprocessor. See more details at the end.

To Reproduce

import torch
import torch.nn as nn
from art.defences.preprocessor.preprocessor import PreprocessorPyTorch
from art.estimators.classification import PyTorchClassifier
from torchvision.models import resnet18


class Defense(PreprocessorPyTorch):

    def forward(self, x, y=None):
        return x.detach(), y

    def estimate_forward(self, x, y=None):
        return x, y


if __name__ == '__main__':
    classifier = PyTorchClassifier(
        model=resnet18(),
        loss=nn.CrossEntropyLoss(),
        input_shape=(3, 224, 224),
        nb_classes=1000,
        preprocessing_defences=Defense(),
        device_type='cpu'
    )
    x = torch.rand(1, 3, 224, 224)
    y = torch.eye(1000)[[0]]
    grad = classifier.loss_gradient(x, y)

The above code will raise the following error:

Traceback (most recent call last):
  File "/xxx/test.py", line 20, in <module>
    classifier.loss_gradient(x, y)
  File "/xxx/envs/armory/lib/python3.7/site-packages/art/estimators/classification/pytorch.py", line 810, in loss_gradient
    assert grads.shape == x.shape
AttributeError: 'NoneType' object has no attribute 'shape'

Expected behavior
classifier.loss_gradient(x, y) should return a gradient of the same shape as x.

Screenshots
N/A

System information (please complete the following information):

• OS: Linux
• Python version: 3.7.9
• ART version or commit number: 1.9.1
• PyTorch version: 1.10.1

More information:
The document below says that we need to define estimate_forward for BPDA, which is a neat solution.

adversarial-robustness-toolbox/art/defences/preprocessor/preprocessor.py

Lines 168 to 179 in bc49671

	def estimate_forward(self, x: "torch.Tensor", y: Optional["torch.Tensor"] = None) -> "torch.Tensor":
	"""
	Provide a differentiable estimate of the forward function, so that autograd can calculate gradients
	of the defence for the backward pass. If the defence is differentiable, just call `self.forward()`.
	If the defence is not differentiable and a differentiable estimate is not available, replace with
	an identity function.
	:param x: Dataset to be preprocessed.
	:param y: Labels to be preprocessed.
	:return: Preprocessed data.
	"""
	return self.forward(x, y=y)[0]

After tracking down the function estimate_forward, I find that it is only invoked through _apply_preprocessing_gradient, which is only enabled when all_framework_preprocessing = False.

adversarial-robustness-toolbox/art/estimators/classification/pytorch.py

Lines 807 to 808 in bc49671

	if not self.all_framework_preprocessing:
	grads = self._apply_preprocessing_gradient(x, grads)

I am not sure how to make all_framework_preprocessing = True. I then tried chaining a non-PyTorch Preprocessor as below

from art.defences.preprocessor.preprocessor import Preprocessor

class DummyDefense(Preprocessor):

    def __call__(self, x, y=None):
        return x, y

    def forward(self, x, y=None):
        return x, y

if __name__ == '__main__':
    classifier = PyTorchClassifier(
        model=resnet18(),
        loss=nn.CrossEntropyLoss(),
        input_shape=(3, 224, 224),
        nb_classes=1000,
        preprocessing_defences=[DummyDefense(), Defense()],
        device_type='cpu'
    )
    x = torch.rand(1, 3, 224, 224)
    y = torch.eye(1000)[[0]]
    grad = classifier.loss_gradient(x, y)

But got the following error:

Traceback (most recent call last):
  File "/xxx/envs/armory/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/xxx/envs/armory/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/xxx/dummy.py", line 36, in <module>
    grad = classifier.loss_gradient(x, y)
  File "/xxx/envs/armory/lib/python3.7/site-packages/art/estimators/classification/pytorch.py", line 775, in loss_gradient
    raise NotImplementedError("Combination of inputs and preprocessing not supported.")
NotImplementedError: Combination of inputs and preprocessing not supported.

[Lodour]

Hi @beat-buesser, sorry to bother but I was wondering why this is not a bug.
The estimate_forward method in PreprocessorPyTorch is documented but not invoked anywhere.

You can verify so by searching "estimate_forward" in the entire codebase:
https://github.com/Trusted-AI/adversarial-robustness-toolbox/search?q=estimate_forward

If I understand it correctly, the current PreprocessorPyTorch does not support BPDA. It seems that PreprocessorPyTorch is assumed to have a differentiable forward, and the estimate_forward is only put there for documentation.

Thanks!

[beat-buesser]

Hi @Lodour No problem. You are asking an important question and I wanted to use your great examples to show the solution. For PreprocessorPyTorch we should only modify the abstract method forward. To apply BPDA or a non-differentiable defense in PreprocessorPyTorch we need to include it in the method forward and the following pattern applies:

import torch
import torch.nn as nn
from torch.autograd import Function

from art.defences.preprocessor.preprocessor import PreprocessorPyTorch
from art.estimators.classification import PyTorchClassifier
from torchvision.models import resnet18


class Defense(PreprocessorPyTorch):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        class NonDifferentialBPDA(Function):  # pylint: disable=W0223
            """
            Function running Preprocessor.
            """

            @staticmethod
            def forward(ctx, input):  # pylint: disable=W0622,W0221
                # Defense can be run on CPU after transferring data (shown here), but can also stay on device
                numpy_input = input.detach().cpu().numpy()
                # Apply defense here, no actual defense applied in this example, input = output
                result = numpy_input
                return input.new(result)

            @staticmethod
            def backward(ctx, grad_output):  # pylint: disable=W0221
                numpy_go = grad_output.cpu().numpy()
                # Apply BPDA or other gradient estimation
                result = numpy_go
                return grad_output.new(result)

        self._non_differential_bpda = NonDifferentialBPDA

    def forward(self, x, y=None):

        x.detach()

        x_bpda = self._non_differential_bpda.apply(x)

        return x_bpda, y


if __name__ == '__main__':
    classifier = PyTorchClassifier(
        model=resnet18(),
        loss=nn.CrossEntropyLoss(),
        input_shape=(3, 224, 224),
        nb_classes=1000,
        preprocessing_defences=Defense(),
        device_type='cpu'
    )
    x = torch.rand(1, 3, 224, 224)
    y = torch.eye(1000)[[0]]
    grad = classifier.loss_gradient(x, y)

We have currently two preprocessing defenses applying this pattern: VideoCompressionPyTorch and Mp3CompressionPyTorch. Differently to the short example above, these example move the non-differentiable part of the defense into separate classes VideoCompression and Mp3Compression of type Preprocessor. We think this pattern can be further simplified and automated if the attack exists in type Preprocessor and plan to make it a feature of future releases.

[Lodour]

Thanks for the clarification, this is a good solution to reuse Preprocessor for PyTorch.

I thought ART was going for the unified estimate_gradient approach because it seemed more straightforward (at least to me) to define estimate_forward rather than backward. For example, if the defense is complicated and there is a better estimate than the identity function, writing down the backward pass might be counterintuitive.

A more concrete example is Section 5.2.1. THERMOMETER ENCODING in the BPDA paper http://proceedings.mlr.press/v80/athalye18a/athalye18a.pdf

Considering cases like this, I would vote for a solution that keeps estimate_forward (unless there were other issues I am not aware of). For example, the estimate_gradient in PreprocessorPyTorch could potentially be a better place to reuse the existing estimate_gradient in Preprocessor, if it detects that a PyTorch version of estimate_forward is not defined.