ShapeShifter for YOLOv3

mart/callbacks/eval_mode.py

@@ -4,23 +4,47 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 
+from __future__ import annotations
+
 from pytorch_lightning.callbacks import Callback
 
+from mart import utils
+
+logger = utils.get_pylogger(__name__)
+
 __all__ = ["AttackInEvalMode"]
 
 
 class AttackInEvalMode(Callback):
     """Switch the model into eval mode during attack."""
 
-    def __init__(self):
-        self.training_mode_status = None
-
-    def on_train_start(self, trainer, model):
-        self.training_mode_status = model.training
-        model.train(False)
-
-    def on_train_end(self, trainer, model):
-        assert self.training_mode_status is not None
-
-        # Resume the previous training status of the model.
-        model.train(self.training_mode_status)
+    def __init__(self, module_classes: type | list[type]):
+        # FIXME: convert strings to classes using hydra.utils.get_class? This will clean up some verbosity in configuration but will require importing hydra in this callback.
+        if isinstance(module_classes, type):
+            module_classes = [module_classes]
+
+        self.module_classes = tuple(module_classes)
+
+    def setup(self, trainer, pl_module, stage):
+        if stage != "fit":
+            return
+
+        # Log to the console so the user can see visually see which modules will be in eval mode during training.
+        for name, module in pl_module.named_modules():
+            if isinstance(module, self.module_classes):
+                logger.info(
+                    f"Setting eval mode for {name} ({module.__class__.__module__}.{module.__class__.__name__})"
+                )
+
+    def on_train_epoch_start(self, trainer, pl_module):
+        # We must use on_train_epoch_start because PL will set pl_module to train mode right before this callback.
+        # See: https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#hooks
+        for name, module in pl_module.named_modules():
+            if isinstance(module, self.module_classes):
+                module.eval()
+
+    def on_train_epoch_end(self, trainer, pl_module):
+        # FIXME: Why is this necessary?
+        for name, module in pl_module.named_modules():
+            if isinstance(module, self.module_classes):
+                module.train()

mart/callbacks/no_grad_mode.py

@@ -4,8 +4,15 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 
+from __future__ import annotations
+
+import torch
 from pytorch_lightning.callbacks import Callback
 
+from mart import utils
+
+logger = utils.get_pylogger(__name__)
+
 __all__ = ["ModelParamsNoGrad"]
 
 
@@ -15,10 +22,25 @@ class ModelParamsNoGrad(Callback):
     This callback should not change the result. Don't use unless an attack runs faster.
     """
 
-    def on_train_start(self, trainer, model):
-        for param in model.parameters():
-            param.requires_grad_(False)
+    def __init__(self, module_names: str | list[str] = None):
+        if isinstance(module_names, str):
+            module_names = [module_names]
+
+        self.module_names = module_names
+
+    def setup(self, trainer, pl_module, stage):
+        if stage != "fit":
+            return
+
+        # We use setup, and not on_train_start, so that mart.optim.OptimizerFactory can ignore parameters with no gradients.
+        # See: https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#hooks
+        for name, param in pl_module.named_parameters():
+            if any(name.startswith(module_name) for module_name in self.module_names):
+                logger.info(f"Disabling gradient for {name}")
+                param.requires_grad_(False)
 
-    def on_train_end(self, trainer, model):
-        for param in model.parameters():
-            param.requires_grad_(True)
+    def teardown(self, trainer, pl_module, stage):
+        for name, param in pl_module.named_parameters():
+            if any(name.startswith(module_name) for module_name in self.module_names):
+                # FIXME: Why is this necessary?
+                param.requires_grad_(True)

mart/callbacks/visualizer.py

@@ -4,38 +4,37 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 
-import os
+from operator import attrgetter
 
 from pytorch_lightning.callbacks import Callback
-from torchvision.transforms import ToPILImage
 
-__all__ = ["PerturbedImageVisualizer"]
+__all__ = ["ImageVisualizer"]
 
 
-class PerturbedImageVisualizer(Callback):
-    """Save adversarial images as files."""
+class ImageVisualizer(Callback):
+    def __init__(self, frequency: int = 100, **tag_paths):
+        self.frequency = frequency
+        self.tag_paths = tag_paths
 
-    def __init__(self, folder):
-        super().__init__()
+    def log_image(self, trainer, tag, image):
+        # Add image to each logger
+        for logger in trainer.loggers:
+            # FIXME: Should we just use isinstance(logger.experiment, SummaryWriter)?
+            if not hasattr(logger.experiment, "add_image"):
+                continue
 
-        # FIXME: This should use the Trainer's logging directory.
-        self.folder = folder
-        self.convert = ToPILImage()
+            logger.experiment.add_image(tag, image, global_step=trainer.global_step)
 
-        if not os.path.isdir(self.folder):
-            os.makedirs(self.folder)
+    def log_images(self, trainer, pl_module):
+        for tag, path in self.tag_paths.items():
+            image = attrgetter(path)(pl_module)
+            self.log_image(trainer, tag, image)
 
-    def on_train_batch_end(self, trainer, model, outputs, batch, batch_idx):
-        # Save input and target for on_train_end
-        self.input = batch["input"]
-        self.target = batch["target"]
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+        if batch_idx % self.frequency != 0:
+            return
 
-    def on_train_end(self, trainer, model):
-        # FIXME: We should really just save this to outputs instead of recomputing adv_input
-        adv_input = model(input=self.input, target=self.target)
+        self.log_images(trainer, pl_module)
 
-        for img, tgt in zip(adv_input, self.target):
-            fname = tgt["file_name"]
-            fpath = os.path.join(self.folder, fname)
-            im = self.convert(img / 255)
-            im.save(fpath)
+    def on_train_end(self, trainer, pl_module):
+        self.log_images(trainer, pl_module)

mart/configs/callbacks/attack_in_eval_mode.yaml

@@ -1,2 +1,11 @@
 attack_in_eval_mode:
   _target_: mart.callbacks.AttackInEvalMode
+  module_classes: ???
+  #  - _target_: hydra.utils.get_class
+  #    path: mart.models.LitModular
+  #  - _target_: hydra.utils.get_class
+  #    path: torch.nn.BatchNorm2d
+  #  - _target_: hydra.utils.get_class
+  #    path: torch.nn.Dropout
+  #  - _target_: hydra.utils.get_class
+  #    path: torch.nn.SyncBatchNorm

mart/configs/callbacks/no_grad_mode.yaml

@@ -1,2 +1,3 @@
-attack_in_eval_mode:
+no_grad_mode:
   _target_: mart.callbacks.ModelParamsNoGrad
+  module_names: ???

mart/configs/callbacks/perturbation_visualizer.yaml

@@ -0,0 +1,4 @@
+perturbation_visualizer:
+  _target_: mart.callbacks.ImageVisualizer
+  frequency: 100
+  perturbation: ???

mart/configs/datamodule/coco.yaml

@@ -1,5 +1,5 @@
 defaults:
-  - default.yaml
+  - default
 
 train_dataset:
   _target_: mart.datamodules.coco.CocoDetection

mart/configs/datamodule/coco_yolov3.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - coco
+
+num_workers: 1
+
+train_dataset:
+  transforms:
+    transforms:
+      - _target_: torchvision.transforms.ToTensor
+      - _target_: mart.transforms.ConvertCocoPolysToMask
+      - _target_: mart.transforms.PadToSquare
+        fill: 0.5
+      - _target_: mart.transforms.Resize
+        size: [416, 416]
+      - _target_: mart.transforms.ConvertBoxesToCXCYHW
+      - _target_: mart.transforms.RemapLabels
+      - _target_: mart.transforms.PackBoxesAndLabels
+        num_classes: 80
+      - _target_: mart.transforms.ConvertInstanceSegmentationToPerturbable
+
+val_dataset:
+  transforms: ${..train_dataset.transforms}
+
+test_dataset:
+  transforms: ${..val_dataset.transforms}
+
+collate_fn:
+  _target_: hydra.utils.get_method
+  path: mart.datamodules.coco.yolo_collate_fn

mart/configs/experiment/COCO_YOLOv3.yaml

@@ -0,0 +1,34 @@
+# @package _global_
+
+defaults:
+  - override /datamodule: coco_yolov3
+  - override /model: yolov3
+  - override /metric: average_precision
+  - override /optimization: super_convergence
+
+task_name: "COCO_YOLOv3"
+tags: ["evaluation"]
+
+optimized_metric: "test_metrics/map"
+
+trainer:
+  # 117,266 training images, 6 epochs, batch_size=16, 43,974.75
+  max_steps: 43975
+  # FIXME: "nms_kernel" not implemented for 'BFloat16', torch.ops.torchvision.nms().
+  precision: 32
+
+datamodule:
+  num_workers: 32
+  ims_per_batch: 16
+
+model:
+  load_state_dict:
+    yolov3: ${paths.data_dir}/yolov3_original.pt
+
+  # yolov3 model does not produce preds/targets in training sequence
+  training_metrics: null
+
+  optimizer:
+    lr: 0.001
+    momentum: 0.9
+    weight_decay: 0.0005