Merge pull request #18 from spencerwooo:vdc-attack

[New Attack] VDC Attack (Virtual Dense Connection)

.gitignore

@@ -1,9 +1,8 @@
 # Dataset files
 datasets/
 
-# PDM files
-.pdm-python
-pdm.lock
+# Lockfiles
+uv.lock
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

pyproject.toml

@@ -12,7 +12,7 @@ dependencies = [
     "scipy>=1.10.1",
 ]
 dynamic = ["version"]
-optional-dependencies = { dev = ["mypy", "rich"] }
+optional-dependencies = { dev = ["mypy", "rich", "timm"] }
 
 [build-system]
 requires = ["setuptools"]

src/torchattack/__init__.py

@@ -15,6 +15,7 @@
 from torchattack.ssp import SSP
 from torchattack.tgr import TGR
 from torchattack.tifgsm import TIFGSM
+from torchattack.vdc import VDC
 from torchattack.vmifgsm import VMIFGSM
 from torchattack.vnifgsm import VNIFGSM
 
@@ -38,6 +39,7 @@
     'SSP',
     'TGR',
     'TIFGSM',
+    'VDC',
     'VMIFGSM',
     'VNIFGSM',
 ]

src/torchattack/base.py

@@ -30,7 +30,7 @@ def __repr__(self) -> str:
         def repr_map(k, v):
             if isinstance(v, float):
                 return f'{k}={v:.3f}'
-            if k in ['model', 'normalize', 'feature_layer']:
+            if k in ['model', 'normalize', 'feature_layer', 'hooks']:
                 return f'{k}={v.__class__.__name__}'
             if isinstance(v, torch.Tensor):
                 return f'{k}={v.shape}'

src/torchattack/eval.py

@@ -159,7 +159,10 @@ def from_pretrained(
             return cls(model_name, device, model, transform, normalize)
 
         except ValueError:
-            print('Model not found in torchvision.models, falling back to timm.')
+            print(
+                f'Warning: Model `{model_name}` not found in torchvision.models, '
+                'falling back to loading weights from timm.'
+            )
             return cls.from_pretrained(model_name, device, from_timm=True)
 
     def forward(self, x):

src/torchattack/pna_patchout.py

@@ -95,7 +95,9 @@ def __init__(
         self.targeted = targeted
         self.lossfn = nn.CrossEntropyLoss()
 
+        # Register hooks
         if self.pna_skip:
+            self.hooks: list[torch.utils.hooks.RemovableHandle] = []
             self._register_vit_model_hook()
 
         # Set default image size and number of patches for PatchOut
@@ -143,6 +145,9 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
             delta.grad.detach_()
             delta.grad.zero_()
 
+            for hook in self.hooks:
+                hook.remove()
+
         return x + delta
 
     def _register_vit_model_hook(self):
@@ -162,7 +167,8 @@ def attn_drop_mask_grad(
         # Register backward hook for layers specified in _supported_vit_cfg
         for layer in self._supported_vit_cfg[self.model_name]:
             module = rgetattr(self.model, layer)
-            module.register_backward_hook(drop_hook_func)
+            hook = module.register_backward_hook(drop_hook_func)
+            self.hooks.append(hook)
 
     def _apply_patch_out(self, delta: torch.Tensor, seed: int) -> torch.Tensor:
         delta_mask = torch.zeros_like(delta)

src/torchattack/tgr.py

@@ -61,6 +61,8 @@ def __init__(
         self.targeted = targeted
         self.lossfn = nn.CrossEntropyLoss()
 
+        # Register hooks
+        self.hooks: list[torch.utils.hooks.RemovableHandle] = []
         self._register_tgr_model_hooks()
 
     def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
@@ -110,6 +112,9 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
             delta.grad.detach_()
             delta.grad.zero_()
 
+            for hook in self.hooks:
+                hook.remove()
+
         return x + delta
 
     def _register_tgr_model_hooks(self):
@@ -348,10 +353,11 @@ def mlp_tgr(
 
         assert self.model_name in _supported_vit_cfg
 
-        for hook, layers in _supported_vit_cfg[self.model_name]:
+        for hook_func, layers in _supported_vit_cfg[self.model_name]:
             for layer in layers:
                 module = rgetattr(self.model, layer)
-                module.register_backward_hook(hook)
+                hook = module.register_backward_hook(hook_func)
+                self.hooks.append(hook)
 
 
 if __name__ == '__main__':