Refactor VDC attack with rgetattr

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/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__':