diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 00000000..99952996 --- /dev/null +++ b/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.* diff --git a/src/cupbearer/data/backdoors.py b/src/cupbearer/data/backdoors.py index 10c6bdf0..0868c466 100644 --- a/src/cupbearer/data/backdoors.py +++ b/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 diff --git a/src/cupbearer/detectors/abstraction/__init__.py b/src/cupbearer/detectors/abstraction/__init__.py index afbcad97..5e11198a 100644 --- a/src/cupbearer/detectors/abstraction/__init__.py +++ b/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 diff --git a/src/cupbearer/detectors/abstraction/abstraction_detector.py b/src/cupbearer/detectors/abstraction/abstraction_detector.py index 487e9e70..cbe23e8f 100644 --- a/src/cupbearer/detectors/abstraction/abstraction_detector.py +++ b/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: diff --git a/src/cupbearer/detectors/finetuning.py b/src/cupbearer/detectors/finetuning.py index efaef734..e9b7f713 100644 --- a/src/cupbearer/detectors/finetuning.py +++ b/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,6 +113,7 @@ 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() @@ -109,6 +121,7 @@ def setup_and_validate(self): self.num_epochs = 1 self.max_steps = 1 self.batch_size = 2 + self.log_every_n_steps = self.max_steps @dataclass diff --git a/src/cupbearer/scripts/conf/train_classifier_conf.py b/src/cupbearer/scripts/conf/train_classifier_conf.py index e8323296..dc071c66 100644 --- a/src/cupbearer/scripts/conf/train_classifier_conf.py +++ b/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 diff --git a/src/cupbearer/scripts/make_adversarial_examples.py b/src/cupbearer/scripts/make_adversarial_examples.py index 5b8c704d..4b94805a 100644 --- a/src/cupbearer/scripts/make_adversarial_examples.py +++ b/src/cupbearer/scripts/make_adversarial_examples.py @@ -35,9 +35,11 @@ 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}." + "Attack failed, new accuracy is" + f" {rob_acc}% > {100 * cfg.success_threshold}%." ) # Plot a few adversarial examples in a grid and save the plot as a pdf diff --git a/src/cupbearer/scripts/train_classifier.py b/src/cupbearer/scripts/train_classifier.py index a1a9410d..e7d81d2a 100644 --- a/src/cupbearer/scripts/train_classifier.py +++ b/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( diff --git a/tests/test_data.py b/tests/test_data.py index 2a3b37e2..d2b15459 100644 --- a/tests/test_data.py +++ b/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) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 369bf97b..2377a3e4 100644 --- a/tests/test_pipeline.py +++ b/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, )