feat: LOKT attack method (#58)

* feat: simple tcgan for lokt

* fix: trainer error

* feat: lokt trainer

* feat: raw lokt attacker

* feat: lokt

* lokt readme

* fix: name error in lokt

* fix: lokt error & lokt attack script

Author: Chrisqcwx

README.md

@@ -45,7 +45,7 @@ The reason why we developed this toolbox is that the research line of **MI** suf
 
 ## :memo: Model Inversion Attacks
 
-|Method|Paper|Publication|Whitebox or Blackbox|Key Characteristics|
+|Method|Paper|Publication|Scenario|Key Characteristics|
 |:-:|:-:|:-:|:-:|:-:|
 |[DeepInversion](./src/modelinversion/attack/DeepInversion/)|Dreaming to Distill: Data-Free Knowledge Transfer via DeepInversion|[CVPR'2020](https://openaccess.thecvf.com/content_CVPR_2020/html/Yin_Dreaming_to_Distill_Data-Free_Knowledge_Transfer_via_DeepInversion_CVPR_2020_paper.html)|whitebox|student-teacher, data-free|
 |[GMI](./src/modelinversion/attack/GMI/)|The Secret Revealer: Generative Model-Inversion Attacks Against Deep Neural Networks|[CVPR'2020](https://openaccess.thecvf.com/content_CVPR_2020/html/Zhang_The_Secret_Revealer_Generative_Model-Inversion_Attacks_Against_Deep_Neural_Networks_CVPR_2020_paper.html)|whitebox|the first GAN-based MIA, instance-level|
@@ -58,6 +58,7 @@ The reason why we developed this toolbox is that the research line of **MI** suf
 |[C2FMI](./src/modelinversion/attack/C2FMI/)|C2FMI: Corse-to-Fine Black-box Model Inversion Attack|[TDSC'2023](https://ieeexplore.ieee.org/abstract/document/10148574)|whitebox, blackbox|gradient-free, two-stage|
 |[Lomma](./src/modelinversion/attack/Lomma/)|Re-Thinking Model Inversion Attacks Against Deep Neural Networks|[CVPR'2023](https://openaccess.thecvf.com/content/CVPR2023/html/Nguyen_Re-Thinking_Model_Inversion_Attacks_Against_Deep_Neural_Networks_CVPR_2023_paper.html)|blackbox|special loss, model augmentation|
 |[RLBMI](./src/modelinversion/attack/RLBMI/)|Reinforcement Learning-Based Black-Box Model Inversion Attacks|[CVPR'2023](https://openaccess.thecvf.com/content/CVPR2023/html/Han_Reinforcement_Learning-Based_Black-Box_Model_Inversion_Attacks_CVPR_2023_paper.html)|blackbox|reinforcement learning|
+|[LOKT](./src/modelinversion/attack/Lokt/)|Label-Only Model Inversion Attacks via Knowledge Transfer|[NeurIPS'2023](https://openreview.net/forum?id=NuoIThPPag)|blackbox|surrogate models, label-only|
 
 
 

attack_scripts/lokt.py

@@ -0,0 +1,48 @@
+import sys
+sys.path.append('.')
+sys.path.append('./src')
+sys.path.append('./src/modelinversion')
+
+from modelinversion.attack.Lokt.attacker import LoktAttackConfig, LoktAttacker
+from development_config import get_dirs
+
+if __name__ == '__main__':
+    dirs = get_dirs('lokt')
+    cache_dir, result_dir, ckpt_dir, dataset_dir = dirs['work_dir'], dirs['result_dir'], dirs['ckpt_dir'], dirs['dataset_dir']
+    
+    # target name support: vgg16, ir152, facenet64, facenet
+    target_name = 'vgg16'
+    # eval name support: vgg16, ir152, facenet64, facenet
+    eval_name = 'facenet'
+    # gan target name support: vgg16
+    gan_target_name = 'vgg16'
+    # dataset name support: celeba
+    dataset_name = 'celeba'
+    # gan dataset name support: celeba, ffhq, facescrub
+    gan_dataset_name = 'celeba'
+    
+    surrogate_names = ['densenet121', 'densenet161', 'densenet169']
+    
+    batch_size = 64
+    target_labels = list(range(0, 10))
+    device = 'cuda:2'
+    
+    config = LoktAttackConfig(
+        target_name=target_name,
+        eval_name=eval_name,
+        ckpt_dir=ckpt_dir,
+        result_dir=result_dir,
+        dataset_dir=dataset_dir,
+        cache_dir=cache_dir,
+        dataset_name=dataset_name,
+        device=device,
+        gan_target_name=gan_target_name,
+        gan_dataset_name=gan_dataset_name,
+        surrogate_names = surrogate_names
+    )
+    
+    attacker = LoktAttacker(config)
+    
+    attacker.attack(batch_size, target_labels)
+    
+    attacker.evaluation(200, knn=True, feature_distance=True, fid=True)

attack_scripts/train_models/lokt.py

@@ -0,0 +1,51 @@
+import sys
+sys.path.append('.')
+sys.path.append('./src')
+sys.path.append('./src/modelinversion')
+
+import os
+
+import torch
+from torch.utils.data import DataLoader
+from torchvision.datasets import ImageFolder
+from torchvision.transforms import ToTensor
+from modelinversion.attack.Lokt.gan_trainer import LoktGANTrainArgs, LoktGANTrainer
+from development_config import get_dirs
+
+from modelinversion.foldermanager import FolderManager
+os.environ["CUDA_VISIBLE_DEVICES"] = '1'
+if __name__ == '__main__':
+    batch_size = 64
+    
+    target_name = 'vgg16'
+    dataset_name = 'celeba'
+    device = 'cuda'
+    
+    dirs = get_dirs('lokt_gan')
+    cache_dir, result_dir, ckpt_dir, dataset_dir = dirs['work_dir'], dirs['result_dir'], dirs['ckpt_dir'], dirs['dataset_dir']
+    dataset = ImageFolder(dataset_dir, transform=ToTensor())
+    epoch_num = 300
+    
+    # trainloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+    
+    train_args = LoktGANTrainArgs(
+        dataset_name=dataset_name,
+        batch_size=batch_size,
+        epoch_num=epoch_num,
+        device=device,
+        target_name=target_name,
+        class_loss_start_iter=15000
+    )
+    
+    folder_manager = FolderManager(ckpt_dir, dataset_dir, cache_dir, result_dir, None)
+    
+    trainer = LoktGANTrainer(train_args, folder_manager)
+    
+    trainer.train()
+    # trainer.prepare_training()
+    # trainer.G.load_state_dict(torch.load('checkpoints/lokt/lokt_celeba_vgg16_G.pt')['state_dict'])
+    z = torch.randn((8, 128), device=device)
+    c = torch.arange(0, 8, device=device, dtype=torch.long)
+    img = trainer.G(z, c)
+    import torchvision
+    torchvision.utils.save_image(img, 'aa.png', normalize=True, nrow = 4)

attack_scripts/train_models/lokt_surrogate.py

@@ -0,0 +1,56 @@
+import sys
+sys.path.append('.')
+sys.path.append('./src')
+sys.path.append('./src/modelinversion')
+
+import os
+
+import torch
+from torch.utils.data import DataLoader
+from torchvision.datasets import ImageFolder
+from torchvision.transforms import ToTensor
+from modelinversion.attack.Lokt.surrogate_trainer import LoktSurrogateTrainArgs, LoktSurrogateTrainer
+from development_config import get_dirs
+from modelinversion.models import *
+
+from modelinversion.foldermanager import FolderManager
+os.environ["CUDA_VISIBLE_DEVICES"] = '1'
+
+if __name__ == '__main__':
+    
+    model_name = 'densenet121'
+    
+    dirs = get_dirs(f'lokt_surrogate_{model_name}')
+    cache_dir, result_dir, ckpt_dir, dataset_dir, defense_ckpt_dir = dirs['work_dir'], dirs['result_dir'], dirs['ckpt_dir'], dirs['dataset_dir'], dirs['defense_ckpt_dir']
+    
+    folder_manager = FolderManager(ckpt_dir, dataset_dir, cache_dir, result_dir, defense_ckpt_dir, 'no_defense')
+    
+    
+    dataset_name = 'celeba'
+    epoch_num = 200
+    lr = 0.01
+    device = 'cuda'
+    batch_size = 256
+    
+    train_args = LoktSurrogateTrainArgs(
+        model_name, 
+        dataset_name, 
+        epoch_num,
+        device=device,
+        target_name='vgg16',
+        target_dataset_name='celeba',
+    )
+    
+    model = get_model(model_name, dataset_name, device=device, backbone_pretrain=True)
+    
+    optimizer = torch.optim.SGD(model.parameters(), lr=lr,
+                        momentum=0.9, weight_decay=5e-4)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max = epoch_num)
+    
+    trainer = LoktSurrogateTrainer(train_args, folder_manager=folder_manager, model=model, optimizer=optimizer, lr_scheduler=scheduler)
+    
+    trainset = trainer.get_trainset()
+    
+    trainloader = DataLoader(trainset, batch_size, shuffle=True)
+    
+    trainer.train(trainloader, None)

src/modelinversion/attack/GMI/gan_trainer.py

@@ -70,6 +70,8 @@ def _gradient_penalty(self, real: torch.Tensor, fake: torch.Tensor):
 
     def train_dis_step(self, batch) -> OrderedDict:
         args = self.args
+        
+        batch = batch[0].to(args.device)
         bs = len(batch)
         z = torch.randn(bs, args.z_dim).to(args.device)
 
@@ -90,7 +92,7 @@ def train_dis_step(self, batch) -> OrderedDict:
 
     def train_gen_step(self, batch) -> OrderedDict:
         args = self.args
-        bs = len(batch)
+        bs = len(batch[0])
         z = torch.randn(bs, args.z_dim).to(args.device)
 
         fake = self.G(z)

src/modelinversion/attack/KEDMI/gan_trainer.py

@@ -81,6 +81,8 @@ def softXEnt(self, input, target):
 
     def train_dis_step(self, batch) -> OrderedDict:
         args = self.args
+        
+        batch = batch[0].to(args.device)
         bs = len(batch)
         z = torch.randn(bs, args.z_dim).to(args.device)
 
@@ -114,6 +116,7 @@ def train_dis_step(self, batch) -> OrderedDict:
 
     def train_gen_step(self, batch) -> OrderedDict:
         args = self.args
+        batch = batch[0].to(args.device)
         bs = len(batch)
         z = torch.randn(bs, args.z_dim).to(args.device)
 

src/modelinversion/attack/Lokt/attacker.py

@@ -0,0 +1,51 @@
+import os
+from dataclasses import dataclass, field
+
+from ...foldermanager import FolderManager
+from ..PLGMI.attacker import PLGMIAttackConfig, PLGMIAttacker
+from ...models import *
+
+@dataclass
+class LoktAttackConfig(PLGMIAttackConfig):
+    surrogate_names: list = field(default_factory=lambda: ['vgg16'])
+
+@dataclass
+class LoktAttacker(PLGMIAttacker):
+    
+    def __init__(self, config: LoktAttackConfig) -> None:
+        self._tag = f'{config.dataset_name}_{config.target_name}_{config.gan_dataset_name}_{config.gan_target_name}'
+        super().__init__(config)
+        self.config: LoktAttackConfig
+        
+        self.surrogate_models = []
+        
+    
+    def prepare_attack(self):
+        super().prepare_attack()
+        
+        config = self.config
+        folder_manager = self.folder_manager
+        
+        self.surrogate_models = []
+        
+        for S_name in config.surrogate_names:
+            surrogate_model = get_model(S_name, config.dataset_name, device=config.device)
+            # folder_manager.load_target_model_state_dict(surrogate_model, config.dataset_name, S_name, device=config.device)
+            folder_manager.load_state_dict(surrogate_model, ['lokt', f'{S_name}_{config.dataset_name}_{self.config.target_name}_{self.config.defense_type}.pt'], device=config.device)
+            surrogate_model.eval()
+            self.surrogate_models.append(surrogate_model)
+            
+    def get_loss(self, fake, iden):
+        aug_list = self.config.attack_transform
+        
+        inv_loss = 0
+        
+        for S in self.surrogate_models:
+            out1 = S(aug_list(fake)).result
+            out2 = S(aug_list(fake)).result
+
+            inv_loss += self.loss_fn(out1, iden) + self.loss_fn(out2, iden)
+        return inv_loss / len(self.surrogate_models)
+            
+    def attack_step(self, iden):
+        return super().attack_step(iden)