Speedup adversarial training with a simple perturber

mart/attack/perturber/perturber.py

@@ -46,7 +46,9 @@ def __init__(
         self.projector = projector
         self.optim_params = optim_params
 
-        self.perturbation = torch.nn.UninitializedParameter()
+        # 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", torch.nn.UninitializedParameter(), persistent=False)
 
         def projector_wrapper(perturber_module, args):
             if isinstance(perturber_module.perturbation, torch.nn.UninitializedParameter):

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/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

@@ -12,6 +12,7 @@
 
 import mart
 from mart.attack import Adversary, NoAdversary
+from mart.attack.perturber import Perturber
 
 
 def test_no_adversary(input_data, target_data):
@@ -80,3 +81,27 @@ def test_adversary_with_model(input_data, target_data, perturbation):
     assert perturber.call_count == 2
 
     torch.testing.assert_close(output_data, input_data + perturbation)
+
+
+def test_adversary_perturber_hidden_params(input_data, target_data):
+    initializer = Mock()
+    perturber = Perturber(initializer)
+    perturber(input_data, target_data)
+
+    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