Speedup adversarial training with a simple perturber

mart/attack/adversary.py

@@ -153,11 +153,6 @@ def on_run_start(
         # Set up the optimizer.
         self.cur_iter = 0
 
-        # We could be at the inference/no-grad mode here.
-        # Initialize lazy module
-        # FIXME: Perturbers can just use on_run_start/on_run_end to initialize
-        self.perturber(input, target)
-
         # param_groups with learning rate and other optim params.
         param_groups = self.perturber.parameter_groups()
 

mart/attack/perturber/batch.py

@@ -60,6 +60,14 @@ def on_run_start(self, adversary, input, target, model, **kwargs):
             perturber = self.perturber_factory(*self.perturber_args, **self.perturber_kwargs)
             self.perturbers[f"input_{i}_perturber"] = perturber
 
+        # Trigger callback
+        for i, (input_i, target_i) in enumerate(zip(input, target)):
+            perturber = self.perturbers[f"input_{i}_perturber"]
+            if isinstance(perturber, Callback):
+                perturber.on_run_start(
+                    adversary=adversary, input=input_i, target=target_i, model=model, **kwargs
+                )
+
     def forward(self, input: torch.Tensor, target: Union[torch.Tensor, Dict[str, Any]]) -> None:
         output = []
         for i, (input_i, target_i) in enumerate(zip(input, target)):

mart/attack/perturber/perturber.py

@@ -9,14 +9,16 @@
 
 import torch
 
+from mart.attack.callbacks import Callback
+
 from ..gradient_modifier import GradientModifier
 from ..initializer import Initializer
 from ..projector import Projector
 
 __all__ = ["Perturber"]
 
 
-class Perturber(torch.nn.modules.lazy.LazyModuleMixin, torch.nn.Module):
+class Perturber(Callback, torch.nn.Module):
     """The base class of perturbers.
 
     A perturber wraps a nn.Parameter and returns this parameter when called. It also enables one to
@@ -46,10 +48,11 @@ def __init__(
         self.projector = projector
         self.optim_params = optim_params
 
-        self.perturbation = torch.nn.UninitializedParameter()
+        # Pre-occupy the name of the buffer, so that extra_repr() always gets perturbation.
+        self.register_buffer("perturbation", torch.nn.UninitializedBuffer(), persistent=False)
 
         def projector_wrapper(perturber_module, args):
-            if isinstance(perturber_module.perturbation, torch.nn.UninitializedParameter):
+            if isinstance(perturber_module.perturbation, torch.nn.UninitializedBuffer):
                 raise ValueError("Perturbation must be initialized")
 
             input, target = args
@@ -59,6 +62,20 @@ def projector_wrapper(perturber_module, args):
         if projector is not None:
             self.register_forward_pre_hook(projector_wrapper)
 
+    def on_run_start(self, *, adversary, input, target, model, **kwargs):
+        # Initialize perturbation.
+        perturbation = torch.zeros_like(input, requires_grad=True)
+
+        # Register perturbation as a non-persistent buffer even though we will optimize it. This is because it is not
+        # a parameter of the underlying model but a parameter of the adversary.
+        self.register_buffer("perturbation", perturbation, persistent=False)
+
+        # A backward hook that will be called when a gradient w.r.t the Tensor is computed.
+        if self.gradient_modifier is not None:
+            self.perturbation.register_hook(self.gradient_modifier)
+
+        self.initializer(self.perturbation)
+
     def parameter_groups(self):
         """Return parameters along with the pre-defined optimization parameters.
 
@@ -71,28 +88,13 @@ def parameter_groups(self):
 
         return [{"params": self.perturbation} | self.optim_params]
 
-    def initialize_parameters(
-        self, input: torch.Tensor, target: Union[torch.Tensor, Dict[str, Any]]
-    ) -> None:
-        assert isinstance(self.perturbation, torch.nn.UninitializedParameter)
-
-        self.perturbation.materialize(input.shape, device=input.device)
-
-        # A backward hook that will be called when a gradient w.r.t the Tensor is computed.
-        if self.gradient_modifier is not None:
-            self.perturbation.register_hook(self.gradient_modifier)
-
-        self.initializer(self.perturbation)
-
     def forward(
         self, input: torch.Tensor, target: Union[torch.Tensor, Dict[str, Any]]
     ) -> torch.Tensor:
         return self.perturbation
 
     def extra_repr(self):
         perturbation = self.perturbation
-        if not self.has_uninitialized_params():
-            perturbation = (perturbation.shape, perturbation.min(), perturbation.max())
 
         return (
             f"{repr(perturbation)}, initializer={self.initializer},"

mart/configs/attack/classification_eps1.75_fgsm.yaml

@@ -1,12 +1,12 @@
 defaults:
   - iterative_sgd
-  - perturber: batch
+  - perturber: default
   - perturber/initializer: constant
   - perturber/gradient_modifier: sign
   - perturber/projector: linf_additive_range
   - objective: misclassification
   - gain: cross_entropy
-  - threat_model: batch_additive
+  - threat_model: additive
 
 optimizer:
   lr: 1.75

mart/configs/attack/classification_eps2_pgd10_step1.yaml

@@ -1,12 +1,12 @@
 defaults:
   - iterative_sgd
-  - perturber: batch
+  - perturber: default
   - perturber/initializer: uniform_lp
   - perturber/gradient_modifier: sign
   - perturber/projector: linf_additive_range
   - objective: misclassification
   - gain: cross_entropy
-  - threat_model: batch_additive
+  - threat_model: additive
 
 optimizer:
   lr: 1

mart/configs/attack/classification_eps8_pgd10_step1.yaml

@@ -1,12 +1,12 @@
 defaults:
   - iterative_sgd
-  - perturber: batch
+  - perturber: default
   - perturber/initializer: uniform_lp
   - perturber/gradient_modifier: sign
   - perturber/projector: linf_additive_range
   - objective: misclassification
   - gain: cross_entropy
-  - threat_model: batch_additive
+  - threat_model: additive
 
 optimizer:
   lr: 1

mart/configs/attack/perturber/default.yaml

@@ -0,0 +1,4 @@
+_target_: mart.attack.Perturber
+initializer: ???
+gradient_modifier: ???
+projector: ???

mart/utils/utils.py

@@ -12,6 +12,7 @@
 from pytorch_lightning import Callback
 from pytorch_lightning.loggers import LightningLoggerBase
 from pytorch_lightning.utilities import rank_zero_only
+from pytorch_lightning.utilities.model_summary import summarize
 
 from mart.utils import pylogger, rich_utils
 
@@ -164,13 +165,11 @@ def log_hyperparameters(object_dict: Dict[str, Any]) -> None:
     hparams["model"] = cfg["model"]
 
     # save number of model parameters
-    hparams["model/params/total"] = sum(p.numel() for p in model.parameters())
-    hparams["model/params/trainable"] = sum(
-        p.numel() for p in model.parameters() if p.requires_grad
-    )
-    hparams["model/params/non_trainable"] = sum(
-        p.numel() for p in model.parameters() if not p.requires_grad
-    )
+    summary = summarize(model)
+
+    hparams["model/params/total"] = summary.total_parameters
+    hparams["model/params/trainable"] = summary.trainable_parameters
+    hparams["model/params/non_trainable"] = summary.total_parameters - summary.trainable_parameters
 
     hparams["datamodule"] = cfg["datamodule"]
     hparams["trainer"] = cfg["trainer"]

tests/test_adversary.py

@@ -5,13 +5,16 @@
 # agreement between Intel Corporation and you.
 #
 
+from functools import partial
 from unittest.mock import Mock
 
 import pytest
 import torch
+from torch.optim import SGD
 
 import mart
 from mart.attack import Adversary, NoAdversary
+from mart.attack.perturber import Perturber
 
 
 def test_no_adversary(input_data, target_data):
@@ -75,8 +78,64 @@ def test_adversary_with_model(input_data, target_data, perturbation):
     assert gain.call_count == max_iters + 1
     assert model.call_count == max_iters + 1
 
-    # Twice because once with model=None to get perturbation and once with
-    # model=model as above.
-    assert perturber.call_count == 2
+    # Once with model=None to get perturbation.
+    # When model=model, perturber.initialize_parameters() is called.
+    assert perturber.call_count == 1
 
     torch.testing.assert_close(output_data, input_data + perturbation)
+
+
+def test_adversary_perturber_hidden_params(input_data, target_data):
+    initializer = Mock()
+    perturber = Perturber(initializer)
+
+    threat_model = mart.attack.threat_model.Additive()
+    optimizer = Mock()
+    gain = Mock(return_value=torch.tensor(0.0, requires_grad=True))
+    model = Mock(return_value={})
+
+    adversary = Adversary(
+        threat_model=threat_model, perturber=perturber, optimizer=optimizer, max_iters=1, gain=gain
+    )
+    output_data = adversary(input_data, target_data, model=model)
+
+    # Adversarial perturbation should not be updated by a regular training optimizer.
+    params = [p for p in adversary.parameters()]
+    assert len(params) == 0
+
+    # Adversarial perturbation should not be saved to the model checkpoint.
+    state_dict = adversary.state_dict()
+    assert "perturber.perturbation" not in state_dict
+
+
+def test_adversary_perturbation(input_data, target_data):
+    threat_model = mart.attack.threat_model.Additive()
+    optimizer = partial(SGD, lr=1.0, maximize=True)
+
+    def gain(logits):
+        return logits.mean()
+
+    # Perturbation initialized as zero.
+    def initializer(x):
+        torch.nn.init.constant_(x, 0)
+
+    perturber = Perturber(initializer)
+
+    adversary = Adversary(
+        threat_model=threat_model, perturber=perturber, optimizer=optimizer, max_iters=1, gain=gain
+    )
+
+    def model(input, target, model=None, **kwargs):
+        return {"logits": adversary(input, target)}
+
+    output1 = adversary(input_data.requires_grad_(), target_data, model=model)
+    pert1 = perturber.perturbation.clone()
+    output2 = adversary(input_data.requires_grad_(), target_data, model=model)
+    pert2 = perturber.perturbation.clone()
+
+    # The perturbation from multiple runs should be the same.
+    torch.testing.assert_close(pert1, pert2)
+
+    # Simulate a new batch of data of different size.
+    new_input_data = torch.cat([input_data, input_data])
+    output3 = adversary(new_input_data, target_data, model=model)

tests/test_perturber.py

@@ -21,7 +21,7 @@ def test_perturber_repr(input_data, target_data):
     perturber = Perturber(initializer, gradient_modifier, projector)
 
     # get additive perturber representation
-    perturbation = torch.nn.UninitializedParameter()
+    perturbation = torch.nn.UninitializedBuffer()
     expected_repr = (
         f"{repr(perturbation)}, initializer={initializer},"
         f"gradient_modifier={gradient_modifier}, projector={projector}"
@@ -31,7 +31,7 @@ def test_perturber_repr(input_data, target_data):
 
     # generate again the perturber with an initialized
     # perturbation
-    perturber.initialize_parameters(input_data, target_data)
+    perturber.on_run_start(adversary=None, input=input_data, target=target_data, model=None)
     representation = perturber.extra_repr()
     assert expected_repr != representation
 
@@ -40,6 +40,7 @@ def test_perturber_forward(input_data, target_data):
     initializer = Mock()
     perturber = Perturber(initializer)
 
+    perturber.on_run_start(adversary=None, input=input_data, target=target_data, model=None)
     output = perturber(input_data, target_data)
     expected_output = perturber.perturbation
     torch.testing.assert_close(output, expected_output, equal_nan=True)