Misc fixes tests

pytest.ini

@@ -0,0 +1,8 @@
+[pytest]
+filterwarnings =
+    ignore:pkg_resources is deprecated as an API.*:DeprecationWarning:.*lightning_utilities.*
+    ignore:.*Deprecated call to `pkg_resources.declare_namespace[(]'lightning'[)].*:DeprecationWarning:.*pkg_resources.*
+    ignore:.*Deprecated call to `pkg_resources.declare_namespace[(]'lightning.fabric'[)].*:DeprecationWarning:.*lightning.fabric.*
+    ignore:.*Deprecated call to `pkg_resources.declare_namespace[(]'lightning.pytorch'[)].*:DeprecationWarning:.*lightning.pytorch.*
+    ignore:distutils Version classes are deprecated.*:DeprecationWarning:.*torchmetrics.*
+    ignore:distutils Version classes are deprecated.*:DeprecationWarning:.*torch.utils.tensorboard.*

src/cupbearer/data/backdoors.py

@@ -149,13 +149,13 @@ def init_warping_field(self, px: int, py: int):
         # upsample expects a batch dimesion, so we add a singleton. We permute after
         # upsampling, since grid_sample expects the length-2 axis to be the last one.
         field = F.interpolate(
-            self.control_grid[None], size=(px, py), mode="bicubic", align_corners=True
+            self.control_grid[None], size=(py, px), mode="bicubic", align_corners=True
         )[0].permute(1, 2, 0)
 
         # Create coordinates by adding to identity field
         xs = torch.linspace(-1, 1, steps=px)
         ys = torch.linspace(-1, 1, steps=py)
-        xx, yy = torch.meshgrid(xs, ys)
+        yy, xx = torch.meshgrid(ys, xs, indexing="ij")
         identity_grid = torch.stack((yy, xx), 2)
         field = field + identity_grid
 

src/cupbearer/detectors/abstraction/__init__.py

@@ -26,6 +26,7 @@ class AbstractionTrainConfig(TrainConfig):
     check_val_every_n_epoch: int = 1
     enable_progress_bar: bool = False
     max_steps: Optional[int] = None
+    log_every_n_steps: Optional[int] = None
     # TODO: should be possible to configure loggers (e.g. wandb)
 
     def setup_and_validate(self):
@@ -34,6 +35,7 @@ def setup_and_validate(self):
             self.batch_size = 2
             self.num_epochs = 1
             self.max_steps = 1
+            self.log_every_n_steps = self.max_steps
 
 
 # This is all unnessarily verbose right now, it's a remnant from when we had

src/cupbearer/detectors/abstraction/abstraction_detector.py

@@ -148,6 +148,7 @@ def train(
         num_epochs: int = 10,
         validation_datasets: Optional[dict[str, Dataset]] = None,
         max_steps: Optional[int] = None,
+        log_every_n_steps: Optional[int] = None,
         **kwargs,
     ):
         # Possibly we should store this as a submodule to save optimizers and continue
@@ -169,6 +170,7 @@ def train(
             enable_checkpointing=False,
             logger=None,
             default_root_dir=self.save_path,
+            log_every_n_steps=log_every_n_steps,
         )
         self.model.eval()
         # We don't need gradients for base model parameters:

src/cupbearer/detectors/finetuning.py

@@ -1,4 +1,5 @@
 import copy
+import warnings
 from dataclasses import dataclass, field
 from typing import Optional
 
@@ -30,6 +31,7 @@ def train(
         num_epochs: int = 10,
         batch_size: int = 128,
         max_steps: Optional[int] = None,
+        log_every_n_steps: Optional[int] = None,
         **kwargs,
     ):
         classifier = Classifier(
@@ -50,11 +52,20 @@ def train(
             max_epochs=num_epochs,
             max_steps=max_steps or -1,
             default_root_dir=self.save_path,
+            log_every_n_steps=log_every_n_steps,
         )
-        trainer.fit(
-            model=classifier,
-            train_dataloaders=clean_loader,
-        )
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                "ignore",
+                message=(
+                    "You defined a `validation_step` but have no `val_dataloader`."
+                    " Skipping val loop."
+                ),
+            )
+            trainer.fit(
+                model=classifier,
+                train_dataloaders=clean_loader,
+            )
 
     def layerwise_scores(self, batch):
         raise NotImplementedError(
@@ -102,13 +113,15 @@ class FinetuningTrainConfig(TrainConfig):
     num_epochs: int = 10
     batch_size: int = 128
     max_steps: Optional[int] = None
+    log_every_n_steps: Optional[int] = (None,)
 
     def setup_and_validate(self):
         super().setup_and_validate()
         if self.debug:
             self.num_epochs = 1
             self.max_steps = 1
             self.batch_size = 2
+            self.log_every_n_steps = self.max_steps
 
 
 @dataclass

src/cupbearer/scripts/conf/train_classifier_conf.py

@@ -25,6 +25,7 @@ class Config(ScriptConfig):
     dir: DirConfig = mutable_field(
         DirConfig, base=os.path.join("logs", "train_classifier")
     )
+    log_every_n_steps: Optional[int] = (None,)
 
     @property
     def num_classes(self):
@@ -61,4 +62,4 @@ def setup_and_validate(self):
             self.max_batch_size = 2
             self.wandb = False
             self.batch_size = 2
-            self.num_workers = 0
+            self.log_every_n_steps = self.max_steps

src/cupbearer/scripts/make_adversarial_examples.py

@@ -35,7 +35,8 @@ def main(cfg: Config):
     )
     rob_acc, l2, elapsed_time = atk.save(dataloader, save_path, return_verbose=True)
 
-    if rob_acc > cfg.success_threshold:
+    # N.B. rob_acc is in percent while success_threshold is not
+    if rob_acc > 100 * cfg.success_threshold:
         raise RuntimeError(
             f"Attack failed, new accuracy is {rob_acc} > {cfg.success_threshold}."
         )

src/cupbearer/scripts/train_classifier.py

@@ -1,8 +1,6 @@
 import warnings
 
 import lightning as L
-from cupbearer.data.backdoor_data import BackdoorData
-from cupbearer.data.backdoors import WanetBackdoor
 from cupbearer.scripts._shared import Classifier
 from cupbearer.utils.scripts import run
 from lightning.pytorch.callbacks import ModelCheckpoint
@@ -13,16 +11,6 @@
 
 
 def main(cfg: Config):
-    if (
-        cfg.num_workers > 0
-        and isinstance(cfg.train_data, BackdoorData)
-        and isinstance(cfg.train_data.backdoor, WanetBackdoor)
-    ):
-        # TODO: actually fix this bug (warping field not being shared among workers)
-        raise NotImplementedError(
-            "WanetBackdoor is not compatible with num_workers > 0 right now."
-        )
-
     dataset = cfg.train_data.build()
 
     train_loader = DataLoader(
@@ -78,6 +66,7 @@ def main(cfg: Config):
         callbacks=callbacks,
         logger=metrics_logger,
         default_root_dir=cfg.dir.path,
+        log_every_n_steps=cfg.log_every_n_steps,
     )
     if not val_loaders:
         warnings.filterwarnings(

tests/test_data.py

@@ -7,7 +7,7 @@
 # We shouldn't import TestDataMix directly because that will make pytest think
 # it's a test.
 from cupbearer import data
-from torch.utils.data import Dataset
+from torch.utils.data import DataLoader, Dataset
 
 
 class DummyDataset(Dataset):
@@ -41,7 +41,7 @@ def __init__(self, length: int, num_classes: int, shape: tuple[int, int]):
         self.num_classes = num_classes
         self.img = torch.tensor(
             [
-                [[i_y % 2, i_x % 2, (i_x + i_y) % 2] for i_x in range(shape[1])]
+                [[i_y % 2, i_x % 2, (i_x + i_y + 1) % 2] for i_x in range(shape[1])]
                 for i_y in range(shape[0])
             ],
             dtype=torch.float32,
@@ -61,7 +61,7 @@ def __getitem__(self, index) -> tuple[torch.Tensor, int]:
 class DummyImageConfig(data.DatasetConfig):
     length: int
     num_classes: int = 10
-    shape: tuple[int, int] = (8, 8)
+    shape: tuple[int, int] = (8, 12)
 
     def _build(self) -> Dataset:
         return DummyImageData(self.length, self.num_classes, self.shape)
@@ -282,3 +282,28 @@ def test_wanet_backdoor(clean_image_config):
         assert torch.max(clean_img) <= 1
         assert torch.max(anoma_img) <= 1
         assert torch.max(noise_img) <= 1
+
+
+def test_wanet_backdoor_on_multiple_workers(
+    clean_image_config,
+):
+    clean_image_config.num_classes = 1
+    target_class = 1
+    anomalous_config = data.BackdoorData(
+        original=clean_image_config,
+        backdoor=data.backdoors.WanetBackdoor(
+            p_backdoor=1.0,
+            p_noise=0.0,
+            target_class=target_class,
+        ),
+    )
+    data_loader = DataLoader(
+        dataset=anomalous_config.build(),
+        num_workers=2,
+        batch_size=1,
+    )
+    imgs = [img for img_batch, label_batch in data_loader for img in img_batch]
+    assert all(torch.allclose(imgs[0], img) for img in imgs)
+
+    clean_image = clean_image_config.build().dataset.img
+    assert not any(torch.allclose(clean_image, img) for img in imgs)

tests/test_pipeline.py

@@ -13,6 +13,15 @@
 from cupbearer.utils.scripts import run
 from simple_parsing import ArgumentGenerationMode, parse
 
+# Ignore warnings about num_workers
+pytestmark = pytest.mark.filterwarnings(
+    "ignore"
+    ":The '[a-z]*_dataloader' does not have many workers which may be a bottleneck. "
+    "Consider increasing the value of the `num_workers` argument` to "
+    "`num_workers=[0-9]*` in the `DataLoader` to improve performance."
+    ":UserWarning"
+)
+
 
 @pytest.fixture(scope="module")
 def backdoor_classifier_path(module_tmp_path):
@@ -81,7 +90,6 @@ def test_train_autoencoder_corner_backdoor(backdoor_classifier_path, tmp_path):
     assert (tmp_path / "eval.json").is_file()
 
 
-# N.B. this test is flaky, sometimes no adversarial examples are found
 @pytest.mark.slow
 def test_train_mahalanobis_advex(backdoor_classifier_path, tmp_path):
     # This test doesn't need a backdoored classifier, but we already have one
@@ -145,7 +153,8 @@ def test_wanet(tmp_path):
         "--train_data backdoor --train_data.original gtsrb "
         "--train_data.backdoor wanet --model mlp "
         "--val_data.backdoor backdoor --val_data.backdoor.original gtsrb "
-        "--val_data.backdoor.backdoor wanet",
+        "--val_data.backdoor.backdoor wanet "
+        "--num_workers=1",
         argument_generation_mode=ArgumentGenerationMode.NESTED,
     )
     run(train_classifier.main, cfg)
@@ -157,8 +166,8 @@ def test_wanet(tmp_path):
     for name, data_cfg in cfg.val_data.items():
         if name == "backdoor":
             assert torch.allclose(
-                data_cfg.backdoor.warping_field,
-                cfg.train_data.backdoor.warping_field,
+                data_cfg.backdoor.control_grid,
+                cfg.train_data.backdoor.control_grid,
             )
         else:
             with pytest.raises(NotImplementedError):
@@ -174,6 +183,6 @@ def test_wanet(tmp_path):
     )
     run(train_detector.main, train_detector_cfg)
     assert torch.allclose(
-        train_detector_cfg.task.backdoor.warping_field,
-        cfg.train_data.backdoor.warping_field,
+        train_detector_cfg.task.backdoor.control_grid,
+        cfg.train_data.backdoor.control_grid,
     )