Merge pull request #109 from ain-soph/unstable

fix #105

docs/source/conf.py

@@ -83,6 +83,7 @@ def linkcode_resolve(domain, info):
     'numpy': ('https://numpy.org/doc/stable', None),
     'pillow': ('https://pillow.readthedocs.io/en/stable/', None),
     'python': ('https://docs.python.org/3', None),
+    # 'skimage': ('https://scikit-image.org/docs/dev/', None),
     'sklearn': ('https://scikit-learn.org/stable/', None),
     'torch': ('https://pytorch.org/docs/stable/', None),
     'torchvision': ('https://pytorch.org/vision/stable/', None),

projects/dataset_condensation/main.py

@@ -290,7 +290,7 @@ def get_real_data(c: int, batch_size: int, **kwargs) -> torch.Tensor:
                   f'median: {float(image_syn.median()):.2f} '
                   f'({float(image_syn.min()):.2f}, {float(image_syn.max()):.2f})')
             cur_result = robusts.global_avg + accs.global_avg if eval_model.adv_train else accs.global_avg
-            if cur_result > best_result:
+            if cur_result > best_result + 1e-3:
                 best_result = cur_result
                 print(' ' * 12, '{purple}best result update!{reset}'.format(**ansi))
                 if kwargs['file_path']:

requirements.txt

@@ -5,6 +5,7 @@ tqdm
 torch
 torchvision
 numpy
+# scikit-image
 scikit-learn
 scipy
 matplotlib

setup.cfg

@@ -33,6 +33,7 @@ install_requires =
   torchvision>=0.11.1
   numpy>=1.20.3
   matplotlib>=3.4.2
+  # scikit-image>=0.19.2
   scikit-learn>=0.24.0
   scipy>=1.5.4
   pyyaml>=5.3.1

trojanvision/attacks/backdoor/imc.py

@@ -80,6 +80,6 @@ def optimize_mark(self, loss_fn: Callable[..., torch.Tensor] = None, **kwargs):
                 optimizer.step()
                 optimizer.zero_grad()
                 self.mark.mark.detach_()
-        self.mark.mark[:-1] = tanh_func(atanh_mark)
         atanh_mark.requires_grad_(False)
+        self.mark.mark[:-1] = tanh_func(atanh_mark)
         self.mark.mark.detach_()

trojanvision/attacks/backdoor/latent_backdoor.py

@@ -217,8 +217,8 @@ def preprocess_mark(self, other_input: torch.Tensor, other_label: torch.Tensor):
                 optimizer.step()
                 optimizer.zero_grad()
                 self.mark.mark.detach_()
-        self.mark.mark[:-1] = tanh_func(atanh_mark)
         atanh_mark.requires_grad_(False)
+        self.mark.mark[:-1] = tanh_func(atanh_mark)
         self.mark.mark.detach_()
 
     # -------------------------------- Loss Utils ------------------------------ #

trojanvision/attacks/backdoor/trojannn.py

@@ -1,7 +1,9 @@
 #!/usr/bin/env python3
 
 r"""
-CUDA_VISIBLE_DEVICES=0 python examples/backdoor_attack.py --color --verbose 1 --pretrained --validate_interval 1 --epochs 10 --lr 0.01 --mark_random_init --attack trojannn
+CUDA_VISIBLE_DEVICES=0 python examples/backdoor_attack.py --color --verbose 1 --tqdm --pretrained --validate_interval 1 --epochs 10 --lr 0.01 --mark_random_init --attack trojannn
+
+CUDA_VISIBLE_DEVICES=0 python examples/backdoor_attack.py --color --verbose 1 --tqdm --pretrained --validate_interval 1 --epochs 10 --lr 0.01 --mark_random_init --attack trojannn --model vgg13_comp --preprocess_layer classifier.fc1 --preprocess_next_layer classifier.fc2
 """  # noqa: E501
 
 from .badnet import BadNet
@@ -11,6 +13,9 @@
 
 import torch
 import torch.optim as optim
+# import numpy as np
+# import skimage.restoration
+
 import argparse
 
 
@@ -76,8 +81,6 @@ def __init__(self, preprocess_layer: str = 'flatten', preprocess_next_layer: str
                  neuron_lr: float = 0.1, neuron_epoch: int = 1000,
                  **kwargs):
         super().__init__(**kwargs)
-        if not self.mark.mark_random_init:
-            raise Exception('TrojanNN requires "mark_random_init" to be True to initialize watermark.')
         if self.mark.mark_random_pos:
             raise Exception('TrojanNN requires "mark_random_pos" to be False to max activate neurons.')
 
@@ -93,9 +96,15 @@ def __init__(self, preprocess_layer: str = 'flatten', preprocess_next_layer: str
         self.neuron_num = neuron_num
 
         self.neuron_idx: torch.Tensor = None
+        self.background = torch.zeros(self.dataset.data_shape, device=env['device']).unsqueeze(0)
+        # Original code: doesn't work on resnet18_comp
+        # self.background = torch.normal(mean=175.0 / 255, std=8.0 / 255,
+        #                                size=self.dataset.data_shape,
+        #                                device=env['device']).clamp(0, 1).unsqueeze(0)
 
     def attack(self, *args, **kwargs):
         self.neuron_idx = self.get_neuron_idx()
+        print('Neuron Idx: ', self.neuron_idx.cpu().tolist())
         self.preprocess_mark(neuron_idx=self.neuron_idx)
         super().attack(*args, **kwargs)
 
@@ -111,14 +120,13 @@ def get_neuron_idx(self) -> torch.Tensor:
         """
         weight = self.model.state_dict()[self.preprocess_next_layer + '.weight'].abs()
         if weight.dim() > 2:
-            weight = weight.flatten(2).mean(2)
-        weight = weight.mean(0)
-        return weight.argsort(descending=True)[:self.neuron_num]
+            weight = weight.flatten(2).sum(2)
+        return weight.sum(0).argsort(descending=True)[:self.neuron_num]
 
     def get_neuron_value(self, trigger_input: torch.Tensor, neuron_idx: torch.Tensor) -> float:
         r"""Get average neuron activation value of :attr:`trigger_input` for :attr:`neuron_idx`.
 
-        The feature map is obtained by calling :meth:`ImageModel.get_layer()`.
+        The feature map is obtained by calling :meth:`trojanvision.models.ImageModel.get_layer()`.
 
         Args:
             trigger_input (torch.Tensor): Triggered input tensor with shape ``(N, C, H, W)``.
@@ -130,53 +138,100 @@ def get_neuron_value(self, trigger_input: torch.Tensor, neuron_idx: torch.Tensor
         trigger_feats = self.model.get_layer(
             trigger_input, layer_output=self.preprocess_layer)[:, neuron_idx].abs()
         if trigger_feats.dim() > 2:
-            trigger_feats = trigger_feats.flatten(2).mean(2)
-        return trigger_feats.mean().item()
+            trigger_feats = trigger_feats.flatten(2).sum(2)
+        return trigger_feats.sum().item()
 
-    # train the mark to activate the least-used neurons.
     def preprocess_mark(self, neuron_idx: torch.Tensor):
         r"""Optimize mark to maxmize activation on :attr:`neuron_idx`.
         It uses :any:`torch.optim.Adam` and
         :any:`torch.optim.lr_scheduler.CosineAnnealingLR`
         with tanh objective funcion.
 
-        The feature map is obtained by calling :meth:`ImageModel.get_layer()`.
+        The feature map is obtained by calling :meth:`trojanvision.models.ImageModel.get_layer()`.
 
         Args:
             neuron_idx (torch.Tensor): Neuron index list tensor with shape ``(self.neuron_num)``.
         """
-        zeros = torch.zeros(self.dataset.data_shape, device=env['device']).unsqueeze(0)
-        with torch.no_grad():
-            trigger_input = self.add_mark(zeros, mark_alpha=1.0)
-            print('Neuron Value Before Preprocessing:',
-                  f'{self.get_neuron_value(trigger_input, neuron_idx):.5f}')
-
         atanh_mark = torch.randn_like(self.mark.mark[:-1], requires_grad=True)
+        # Original code: no difference
+        # start_h, start_w = self.mark.mark_height_offset, self.mark.mark_width_offset
+        # end_h, end_w = start_h + self.mark.mark_height, start_w + self.mark.mark_width
+        # self.mark.mark[:-1] = self.background[0, :, start_h:end_h, start_w:end_w]
+        # atanh_mark = (self.mark.mark[:-1] * (2 - 1e-5) - 1).atanh()
+        # atanh_mark.requires_grad_()
+        self.mark.mark[:-1] = tanh_func(atanh_mark.detach())
+        self.mark.mark.detach_()
+
         optimizer = optim.Adam([atanh_mark], lr=self.neuron_lr)
+        # No difference for SGD
+        # optimizer = optim.SGD([atanh_mark], lr=self.neuron_lr)
         optimizer.zero_grad()
         lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(
             optimizer, T_max=self.neuron_epoch)
 
+        with torch.no_grad():
+            trigger_input = self.add_mark(self.background, mark_alpha=1.0)
+            print('Neuron Value Before Preprocessing:',
+                  f'{self.get_neuron_value(trigger_input, neuron_idx):.5f}')
+
         for _ in range(self.neuron_epoch):
             self.mark.mark[:-1] = tanh_func(atanh_mark)
-            trigger_input = self.add_mark(zeros, mark_alpha=1.0)
-            trigger_feats = self.model.get_layer(trigger_input, layer_output=self.preprocess_layer).abs()
+            trigger_input = self.add_mark(self.background, mark_alpha=1.0)
+            trigger_feats = self.model.get_layer(trigger_input, layer_output=self.preprocess_layer)
+            trigger_feats = trigger_feats[:, neuron_idx].abs()
             if trigger_feats.dim() > 2:
-                trigger_feats = trigger_feats.flatten(2).mean(2)
-            loss = (trigger_feats[0] - self.target_value).square().sum()
+                trigger_feats = trigger_feats.flatten(2).sum(2)  # .amax(2)
+            loss = (trigger_feats - self.target_value).square().sum()
+            # Original code: no difference
+            # loss = -self.target_value * trigger_feats.sum()
             loss.backward(inputs=[atanh_mark])
             optimizer.step()
             lr_scheduler.step()
             optimizer.zero_grad()
             self.mark.mark.detach_()
-        self.mark.mark[:-1] = tanh_func(atanh_mark)
+
+            # Original Code: no difference
+            # self.mark.mark[:-1] = tanh_func(atanh_mark.detach())
+            # trigger = self.denoise(self.add_mark(torch.zeros_like(self.background), mark_alpha=1.0)[0])
+            # mark = trigger[:, start_h:end_h, start_w:end_w].clamp(0, 1)
+            # atanh_mark.data = (mark * (2 - 1e-5) - 1).atanh()
+
         atanh_mark.requires_grad_(False)
+        self.mark.mark[:-1] = tanh_func(atanh_mark)
         self.mark.mark.detach_()
 
     def validate_fn(self, **kwargs) -> tuple[float, float]:
         if self.neuron_idx is not None:
             with torch.no_grad():
-                zeros = torch.zeros(self.dataset.data_shape, device=env['device']).unsqueeze(0)
-                trigger_input = self.add_mark(zeros, mark_alpha=1.0)
+                trigger_input = self.add_mark(self.background, mark_alpha=1.0)
                 print(f'Neuron Value: {self.get_neuron_value(trigger_input, self.neuron_idx):.5f}')
         return super().validate_fn(**kwargs)
+
+    # @staticmethod
+    # def denoise(img: torch.Tensor, weight: float = 1.0, max_num_iter: int = 100, eps: float = 1e-3) -> torch.Tensor:
+    #     r"""Denoise image by calling :any:`skimage.restoration.denoise_tv_bregman`.
+
+    #     Warning:
+    #         This method is currently unused in :meth:`preprocess_mark()`
+    #         because no performance difference is observed.
+
+    #     Args:
+    #         img (torch.Tensor): Noisy image tensor with shape ``(C, H, W)``.
+
+    #     Returns:
+    #         torch.Tensor: Denoised image tensor with shape ``(C, H, W)``.
+    #     """
+    #     if img.size(0) == 1:
+    #         img_np: np.ndarray = img[0].detach().cpu().numpy()
+    #     else:
+    #         img_np = img.detach().cpu().permute(1, 2, 0).contiguous().numpy()
+
+    #     denoised_img_np = skimage.restoration.denoise_tv_bregman(
+    #         img_np, weight=weight, max_num_iter=max_num_iter, eps=eps)
+    #     denoised_img = torch.from_numpy(denoised_img_np)
+
+    #     if denoised_img.dim() == 2:
+    #         denoised_img.unsqueeze_(0)
+    #     else:
+    #         denoised_img = denoised_img.permute(2, 0, 1).contiguous()
+    #     return img.to(device=img.device)

trojanzoo/trainer.py

@@ -196,6 +196,8 @@ def summary(self, indent: int = None):
             if value:
                 prints('{green}{0:<10s}{reset}'.format(key, **ansi),
                        indent=indent + 10)
+                if isinstance(value, dict):
+                    value = {k: str(v).split('\n')[0] for k, v in value.items()}
                 prints(value, indent=indent + 10)
                 prints('-' * 20, indent=indent + 10)
 

trojanzoo/utils/logger.py

@@ -27,23 +27,23 @@ class SmoothedValue:
         https://github.com/pytorch/vision/blob/main/references/classification/utils.py
 
     Args:
+        name (str): Name string.
         window_size (int): The :attr:`maxlen` of :class:`~collections.deque`.
         fmt (str): The format pattern of ``str(self)``.
 
     Attributes:
+        name (str): Name string.
+        fmt (str): The string pattern.
         deque (~collections.deque): The unique data series.
         count (int): The amount of data.
         total (float): The sum of all data.
-        fmt (str): The string pattern.
 
-        last (float): The last value of :attr:`deque`.
         median (float): The median of :attr:`deque`.
         avg (float): The avg of :attr:`deque`.
-        global_avg (float):
-          :math:`\frac{\text{total}}{\text{count}}`
+        global_avg (float): :math:`\frac{\text{total}}{\text{count}}`
         max (float): The max of :attr:`deque`.
         min (float): The min of :attr:`deque`.
-        value (float): The last value of :attr:`deque`.
+        last_value (float): The last value of :attr:`deque`.
     """
 
     def __init__(self, name: str = '', window_size: int = None, fmt: str = '{global_avg:.3f}'):
@@ -53,13 +53,6 @@ def __init__(self, name: str = '', window_size: int = None, fmt: str = '{global_
         self.total: float = 0.0
         self.fmt = fmt
 
-    @property
-    def last(self) -> float:
-        try:
-            return self.deque[-1]
-        except IndexError:
-            raise IndexError(f'{self.name} is empty')
-
     def update(self, value: float, n: int = 1) -> 'SmoothedValue':
         r"""Update :attr:`n` pieces of data with same :attr:`value`.
 
@@ -169,7 +162,7 @@ def min(self) -> float:
             return 0.0
 
     @property
-    def value(self) -> float:
+    def last_value(self) -> float:
         try:
             return self.deque[-1]
         except Exception:
@@ -183,7 +176,7 @@ def __str__(self):
             global_avg=self.global_avg,
             min=self.min,
             max=self.max,
-            value=self.value)
+            last_value=self.last_value)
 
     def __format__(self, format_spec: str) -> str:
         return self.__str__()

trojanzoo/utils/module/process.py

@@ -49,7 +49,7 @@ def summary(self, indent: int = None):
             if value:
                 prints('{green}{0:<20s}{reset}'.format(
                     key, **ansi), indent=indent + 10)
-                prints({v: getattr(self, v)
+                prints({v: str(getattr(self, v)).split('\n')[0]
                        for v in value}, indent=indent + 10)
                 prints('-' * 20, indent=indent + 10)