From b12accee6aa1b4c99fe19104fa95f1f2f65ebab2 Mon Sep 17 00:00:00 2001 From: "lukyanov_kirya@bk.ru" Date: Mon, 25 Nov 2024 13:53:34 +0300 Subject: [PATCH 1/5] add QuantizationDefender --- experiments/attack_defense_test.py | 17 +++++++--- metainfo/evasion_defense_parameters.json | 2 +- src/defense/evasion_defense.py | 43 ++++++++++++++++++------ 3 files changed, 47 insertions(+), 15 deletions(-) diff --git a/experiments/attack_defense_test.py b/experiments/attack_defense_test.py index 6d2b6b6..046f888 100644 --- a/experiments/attack_defense_test.py +++ b/experiments/attack_defense_test.py @@ -131,7 +131,7 @@ def test_attack_defense(): _import_path=POISON_ATTACK_PARAMETERS_PATH, _config_class="PoisonAttackConfig", _config_kwargs={ - "n_edges_percent": 0.5, + "n_edges_percent": 1.0, } ) @@ -210,7 +210,16 @@ def test_attack_defense(): _import_path=EVASION_DEFENSE_PARAMETERS_PATH, _config_class="EvasionDefenseConfig", _config_kwargs={ - "regularization_strength": 0.1 * 10 + "regularization_strength": 0.1 * 1000 + } + ) + + quantization_evasion_defense_config = ConfigPattern( + _class_name="QuantizationDefender", + _import_path=EVASION_DEFENSE_PARAMETERS_PATH, + _config_class="EvasionDefenseConfig", + _config_kwargs={ + "num_levels": 2 } ) @@ -234,8 +243,8 @@ def test_attack_defense(): # gnn_model_manager.set_poison_attacker(poison_attack_config=random_poison_attack_config) # gnn_model_manager.set_poison_defender(poison_defense_config=gnnguard_poison_defense_config) - # gnn_model_manager.set_evasion_attacker(evasion_attack_config=netattackgroup_evasion_attack_config) - # gnn_model_manager.set_evasion_defender(evasion_defense_config=at_evasion_defense_config) + gnn_model_manager.set_evasion_attacker(evasion_attack_config=fgsm_evasion_attack_config) + gnn_model_manager.set_evasion_defender(evasion_defense_config=gradientregularization_evasion_defense_config) warnings.warn("Start training") dataset.train_test_split() diff --git a/metainfo/evasion_defense_parameters.json b/metainfo/evasion_defense_parameters.json index 682a4d2..b9514be 100644 --- a/metainfo/evasion_defense_parameters.json +++ b/metainfo/evasion_defense_parameters.json @@ -5,7 +5,7 @@ "regularization_strength": ["regularization_strength", "float", 0.1, {"min": 0.0001, "step": 0.01}, "?"] }, "QuantizationDefender": { - "qbit": ["qbit", "int", 8, {"min": 1, "step": 1}, "?"] + "num_levels": ["num_levels", "int", 32, {"min": 2, "step": 1}, "?"] }, "AdvTraining": { "attack_name": ["attack_name", "str", "FGSM", {}, "?"] diff --git a/src/defense/evasion_defense.py b/src/defense/evasion_defense.py index 18d2cca..9cd9c5a 100644 --- a/src/defense/evasion_defense.py +++ b/src/defense/evasion_defense.py @@ -54,14 +54,25 @@ def post_batch( pass -class GradientRegularizationDefender(EvasionDefender): +class GradientRegularizationDefender( + EvasionDefender +): name = "GradientRegularizationDefender" - def __init__(self, regularization_strength=0.1): + def __init__( + self, + regularization_strength: float = 0.1 + ): super().__init__() self.regularization_strength = regularization_strength - def post_batch(self, model_manager, batch, loss, **kwargs): + def post_batch( + self, + model_manager, + batch, + loss: torch.Tensor, + **kwargs + ) -> dict: batch.x.requires_grad = True outputs = model_manager.gnn(batch.x, batch.edge_index) loss_loc = model_manager.loss_function(outputs, batch.y) @@ -69,10 +80,10 @@ def post_batch(self, model_manager, batch, loss, **kwargs): grad_outputs=torch.ones_like(loss_loc), create_graph=True, retain_graph=True, only_inputs=True)[0] gradient_penalty = torch.sum(gradients ** 2) + batch.x.requires_grad = False return {"loss": loss + self.regularization_strength * gradient_penalty} -# TODO Kirill, add code in pre_batch class QuantizationDefender( EvasionDefender ): @@ -80,17 +91,31 @@ class QuantizationDefender( def __init__( self, - qbit: int = 8 + num_levels: int = 32 ): super().__init__() - self.regularization_strength = qbit + self.num_levels = num_levels def pre_batch( self, + model_manager, + batch, **kwargs ): - # TODO Kirill - pass + x = batch.x + batch.x = self.quantize(x) + return batch + + def quantize( + self, + x + ): + x_min = x.min() + x_max = x.max() + x_normalized = (x - x_min) / (x_max - x_min) + x_quantized = torch.round(x_normalized * (self.num_levels - 1)) / (self.num_levels - 1) + x_quantized = x_quantized * (x_max - x_min) + x_min + return x_quantized class AdvTraining( @@ -104,10 +129,8 @@ def __init__( attack_name: str = None, attack_config: EvasionAttackConfig = None, attack_type: str = None, - device: str = 'cpu' ): super().__init__() - assert device is not None, "Please specify 'device'!" if not attack_config: # build default config assert attack_name is not None From 9e51ade2e504349098bf807f53c31fbdbdb457d1 Mon Sep 17 00:00:00 2001 From: "lukyanov_kirya@bk.ru" Date: Mon, 25 Nov 2024 14:41:14 +0300 Subject: [PATCH 2/5] DistillationDefender --- experiments/attack_defense_test.py | 13 ++++++-- metainfo/evasion_defense_parameters.json | 3 ++ src/defense/evasion_defense.py | 36 +++++++++++++++++++++- src/models_builder/gnn_constructor.py | 38 +++++++++++++++++------- 4 files changed, 77 insertions(+), 13 deletions(-) diff --git a/experiments/attack_defense_test.py b/experiments/attack_defense_test.py index 046f888..eb8cf4e 100644 --- a/experiments/attack_defense_test.py +++ b/experiments/attack_defense_test.py @@ -173,7 +173,7 @@ def test_attack_defense(): _import_path=EVASION_ATTACK_PARAMETERS_PATH, _config_class="EvasionAttackConfig", _config_kwargs={ - "epsilon": 0.01 * 1, + "epsilon": 0.001 * 12, } ) @@ -223,6 +223,15 @@ def test_attack_defense(): } ) + distillation_evasion_defense_config = ConfigPattern( + _class_name="DistillationDefender", + _import_path=EVASION_DEFENSE_PARAMETERS_PATH, + _config_class="EvasionDefenseConfig", + _config_kwargs={ + "temperature": 0.5 * 20 + } + ) + fgsm_evasion_attack_config0 = ConfigPattern( _class_name="FGSM", _import_path=EVASION_ATTACK_PARAMETERS_PATH, @@ -244,7 +253,7 @@ def test_attack_defense(): # gnn_model_manager.set_poison_attacker(poison_attack_config=random_poison_attack_config) # gnn_model_manager.set_poison_defender(poison_defense_config=gnnguard_poison_defense_config) gnn_model_manager.set_evasion_attacker(evasion_attack_config=fgsm_evasion_attack_config) - gnn_model_manager.set_evasion_defender(evasion_defense_config=gradientregularization_evasion_defense_config) + gnn_model_manager.set_evasion_defender(evasion_defense_config=distillation_evasion_defense_config) warnings.warn("Start training") dataset.train_test_split() diff --git a/metainfo/evasion_defense_parameters.json b/metainfo/evasion_defense_parameters.json index b9514be..e9ff771 100644 --- a/metainfo/evasion_defense_parameters.json +++ b/metainfo/evasion_defense_parameters.json @@ -7,6 +7,9 @@ "QuantizationDefender": { "num_levels": ["num_levels", "int", 32, {"min": 2, "step": 1}, "?"] }, + "DistillationDefender": { + "temperature": ["temperature", "float", 5.0, {"min": 1, "step": 0.01}, "?"] + }, "AdvTraining": { "attack_name": ["attack_name", "str", "FGSM", {}, "?"] } diff --git a/src/defense/evasion_defense.py b/src/defense/evasion_defense.py index 9cd9c5a..2eea79e 100644 --- a/src/defense/evasion_defense.py +++ b/src/defense/evasion_defense.py @@ -108,7 +108,7 @@ def pre_batch( def quantize( self, - x + x: torch.Tensor ): x_min = x.min() x_max = x.max() @@ -118,6 +118,40 @@ def quantize( return x_quantized +class DistillationDefender( + EvasionDefender +): + name = "DistillationDefender" + + def __init__( + self, + temperature: float = 5.0 + ): + """ + """ + super().__init__() + self.temperature = temperature + + def post_batch( + self, + model_manager, + batch, + loss: torch.Tensor + ): + """ + """ + model = model_manager.gnn + logits = model(batch) + soft_targets = torch.softmax(logits / self.temperature, dim=1) + distillation_loss = torch.nn.functional.kl_div( + torch.log_softmax(logits / self.temperature, dim=1), + soft_targets, + reduction='batchmean' + ) * (self.temperature ** 2) + modified_loss = loss + distillation_loss + return {"loss": modified_loss} + + class AdvTraining( EvasionDefender ): diff --git a/src/models_builder/gnn_constructor.py b/src/models_builder/gnn_constructor.py index 3df22e8..bf8b325 100644 --- a/src/models_builder/gnn_constructor.py +++ b/src/models_builder/gnn_constructor.py @@ -617,17 +617,35 @@ def arguments_read( edge_weight = kwargs.get('edge_weight', None) if batch is None: batch = torch.zeros(kwargs['x'].shape[0], dtype=torch.int64, device=x.device) - elif len(args) == 2: - x, edge_index = args[0], args[1] - batch = torch.zeros(args[0].shape[0], dtype=torch.int64, device=x.device) - edge_weight = None - elif len(args) == 3: - x, edge_index, batch = args[0], args[1], args[2] - edge_weight = None - elif len(args) == 4: - x, edge_index, batch, edge_weight = args[0], args[1], args[2], args[3] else: - raise ValueError(f"forward's args should take 2 or 3 arguments but got {len(args)}") + if len(args) == 1: + args = args[0] + if 'x' in args and 'edge_index' in args: + x, edge_index = args.x, args.edge_index + else: + raise ValueError(f"forward's args should contain x and 3" + f" edge_index Tensors but {args.keys} doesn't content this Tensors") + if 'batch' in args: + batch = args.batch + else: + batch = torch.zeros(args.x.shape[0], dtype=torch.int64, device=x.device) + if 'edge_weight' in args: + edge_weight = args.edge_weight + else: + edge_weight = None + else: + if len(args) == 2: + x, edge_index = args[0], args[1] + batch = torch.zeros(args[0].shape[0], dtype=torch.int64, device=x.device) + edge_weight = None + elif len(args) == 3: + x, edge_index, batch = args[0], args[1], args[2] + edge_weight = None + elif len(args) == 4: + x, edge_index, batch, edge_weight = args[0], args[1], args[2], args[3] + else: + raise ValueError(f"forward's args should take 2 or 3 arguments but got {len(args)}") + else: if hasattr(data, "edge_weight"): x, edge_index, batch, edge_weight = data.x, data.edge_index, data.batch, data.edge_weight From aeed17675e8cd8dd722485ef0181ed45888b4dd3 Mon Sep 17 00:00:00 2001 From: "lukyanov_kirya@bk.ru" Date: Mon, 25 Nov 2024 15:49:01 +0300 Subject: [PATCH 3/5] add AutoEncoderDefender --- experiments/attack_defense_test.py | 13 +++- metainfo/evasion_defense_parameters.json | 5 ++ src/defense/evasion_defense.py | 92 ++++++++++++++++++++++++ 3 files changed, 109 insertions(+), 1 deletion(-) diff --git a/experiments/attack_defense_test.py b/experiments/attack_defense_test.py index eb8cf4e..8b386d5 100644 --- a/experiments/attack_defense_test.py +++ b/experiments/attack_defense_test.py @@ -232,6 +232,17 @@ def test_attack_defense(): } ) + autoencoder_evasion_defense_config = ConfigPattern( + _class_name="AutoEncoderDefender", + _import_path=EVASION_DEFENSE_PARAMETERS_PATH, + _config_class="EvasionDefenseConfig", + _config_kwargs={ + "hidden_dim": 300, + "bottleneck_dim": 100, + "reconstruction_loss_weight": 0.1, + } + ) + fgsm_evasion_attack_config0 = ConfigPattern( _class_name="FGSM", _import_path=EVASION_ATTACK_PARAMETERS_PATH, @@ -253,7 +264,7 @@ def test_attack_defense(): # gnn_model_manager.set_poison_attacker(poison_attack_config=random_poison_attack_config) # gnn_model_manager.set_poison_defender(poison_defense_config=gnnguard_poison_defense_config) gnn_model_manager.set_evasion_attacker(evasion_attack_config=fgsm_evasion_attack_config) - gnn_model_manager.set_evasion_defender(evasion_defense_config=distillation_evasion_defense_config) + gnn_model_manager.set_evasion_defender(evasion_defense_config=autoencoder_evasion_defense_config) warnings.warn("Start training") dataset.train_test_split() diff --git a/metainfo/evasion_defense_parameters.json b/metainfo/evasion_defense_parameters.json index e9ff771..da30af1 100644 --- a/metainfo/evasion_defense_parameters.json +++ b/metainfo/evasion_defense_parameters.json @@ -10,6 +10,11 @@ "DistillationDefender": { "temperature": ["temperature", "float", 5.0, {"min": 1, "step": 0.01}, "?"] }, + "AutoEncoderDefender": { + "hidden_dim": ["hidden_dim", "int", 5, {"min": 3, "step": 1}, "?"], + "bottleneck_dim": ["bottleneck_dim", "int", 3, {"min": 1, "step": 1}, "?"], + "reconstruction_loss_weight": ["reconstruction_loss_weight", "float", 0.1, {"min": 0.0001, "step": 0.01}, "?"] + }, "AdvTraining": { "attack_name": ["attack_name", "str", "FGSM", {}, "?"] } diff --git a/src/defense/evasion_defense.py b/src/defense/evasion_defense.py index 2eea79e..003c7b2 100644 --- a/src/defense/evasion_defense.py +++ b/src/defense/evasion_defense.py @@ -234,3 +234,95 @@ def post_batch( outputs = model_manager.gnn(self.perturbed_gen_dataset.data.x, self.perturbed_gen_dataset.data.edge_index) loss_loc = model_manager.loss_function(outputs, batch.y) return {"loss": loss + loss_loc} + + +class SimpleAutoEncoder( + torch.nn.Module +): + def __init__( + self, + input_dim: int, + hidden_dim: int, + bottleneck_dim: int, + device: str = 'cpu' + ): + """ + """ + super(SimpleAutoEncoder, self).__init__() + self.device = device + self.encoder = torch.nn.Sequential( + torch.nn.Linear(input_dim, hidden_dim), + torch.nn.ReLU(), + torch.nn.Linear(hidden_dim, bottleneck_dim), + torch.nn.ReLU() + ).to(self.device) + self.decoder = torch.nn.Sequential( + torch.nn.Linear(bottleneck_dim, hidden_dim), + torch.nn.ReLU(), + torch.nn.Linear(hidden_dim, input_dim) + ).to(self.device) + + def forward( + self, + x: torch.Tensor + ): + x = x.to(self.device) + encoded = self.encoder(x) + decoded = self.decoder(encoded) + return decoded + + +class AutoEncoderDefender( + EvasionDefender +): + name = "AutoEncoderDefender" + + def __init__( + self, + hidden_dim: int, + bottleneck_dim: int, + reconstruction_loss_weight=0.1, + ): + """ + """ + super().__init__() + self.autoencoder = None + self.hidden_dim = hidden_dim + self.bottleneck_dim = bottleneck_dim + self.reconstruction_loss_weight = reconstruction_loss_weight + + def post_batch(self, model_manager, batch, loss): + """ + """ + model_manager.gnn.eval() + if self.autoencoder is None: + self.init_autoencoder(batch.x) + self.autoencoder.train() + reconstructed_x = self.autoencoder(batch.x) + reconstruction_loss = torch.nn.functional.mse_loss(reconstructed_x, batch.x) + modified_loss = loss + self.reconstruction_loss_weight * reconstruction_loss.detach().clone() + autoencoder_optimizer = torch.optim.Adam(self.autoencoder.parameters(), lr=0.001) + autoencoder_optimizer.zero_grad() + reconstruction_loss.backward() + autoencoder_optimizer.step() + return {"loss": modified_loss} + + def denoise_with_autoencoder(self, x): + """ + """ + self.autoencoder.eval() + with torch.no_grad(): + x_denoised = self.autoencoder(x) + return x_denoised + + def init_autoencoder( + self, + x + ): + self.autoencoder = SimpleAutoEncoder( + input_dim=x.shape[1], + bottleneck_dim=self.bottleneck_dim, + hidden_dim=self.hidden_dim, + device=x.device + ) + From 6c694c37917d4da5a5b9f71d6ebb469ef7d3cc90 Mon Sep 17 00:00:00 2001 From: "lukyanov_kirya@bk.ru" Date: Mon, 25 Nov 2024 15:53:35 +0300 Subject: [PATCH 4/5] + --- src/defense/evasion_defense.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/src/defense/evasion_defense.py b/src/defense/evasion_defense.py index 003c7b2..52a9a0f 100644 --- a/src/defense/evasion_defense.py +++ b/src/defense/evasion_defense.py @@ -281,7 +281,7 @@ def __init__( self, hidden_dim: int, bottleneck_dim: int, - reconstruction_loss_weight=0.1, + reconstruction_loss_weight: float = 0.1, ): """ """ @@ -291,7 +291,12 @@ def __init__( self.bottleneck_dim = bottleneck_dim self.reconstruction_loss_weight = reconstruction_loss_weight - def post_batch(self, model_manager, batch, loss): + def post_batch( + self, + model_manager, + batch, + loss: torch.Tensor + ) -> dict: """ """ model_manager.gnn.eval() @@ -307,7 +312,10 @@ def post_batch(self, model_manager, batch, loss): autoencoder_optimizer.step() return {"loss": modified_loss} - def denoise_with_autoencoder(self, x): + def denoise_with_autoencoder( + self, + x: torch.Tensor + ) -> torch.Tensor: """ """ self.autoencoder.eval() @@ -317,12 +325,11 @@ def denoise_with_autoencoder(self, x): def init_autoencoder( self, - x - ): + x: torch.Tensor + ) -> None: self.autoencoder = SimpleAutoEncoder( input_dim=x.shape[1], bottleneck_dim=self.bottleneck_dim, hidden_dim=self.hidden_dim, device=x.device ) - From 26e0a51657ba39f3d46b83b401ea00e997c0b419 Mon Sep 17 00:00:00 2001 From: "lukyanov_kirya@bk.ru" Date: Tue, 26 Nov 2024 15:51:06 +0300 Subject: [PATCH 5/5] add assert --- src/defense/evasion_defense.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/defense/evasion_defense.py b/src/defense/evasion_defense.py index 52a9a0f..7604edc 100644 --- a/src/defense/evasion_defense.py +++ b/src/defense/evasion_defense.py @@ -94,6 +94,7 @@ def __init__( num_levels: int = 32 ): super().__init__() + assert num_levels > 1 self.num_levels = num_levels def pre_batch( @@ -112,9 +113,12 @@ def quantize( ): x_min = x.min() x_max = x.max() - x_normalized = (x - x_min) / (x_max - x_min) - x_quantized = torch.round(x_normalized * (self.num_levels - 1)) / (self.num_levels - 1) - x_quantized = x_quantized * (x_max - x_min) + x_min + if x_min != x_max: + x_normalized = (x - x_min) / (x_max - x_min) + x_quantized = torch.round(x_normalized * (self.num_levels - 1)) / (self.num_levels - 1) + x_quantized = x_quantized * (x_max - x_min) + x_min + else: + x_quantized = x return x_quantized