From 564f66a4d53445027983c412558882526a8c8b6e Mon Sep 17 00:00:00 2001 From: anurag0013 Date: Fri, 8 Aug 2025 11:59:54 +0545 Subject: [PATCH] feat: Data Free Attack --- models/attack/attack1.py | 120 ++++++++++++++ models/attack/attack2.py | 103 ++++++++++++ models/attack/attack3.py | 115 +++++++++++++ models/gnn_algo/generator.py | 106 ++++++++++++ models/gnn_algo/surrogate.py | 43 +++++ models/gnn_algo/victim.py | 49 ++++++ models/main.py | 311 +++++++++++++++++++++++++++++++++++ requirements.txt | 11 ++ 8 files changed, 858 insertions(+) create mode 100644 models/attack/attack1.py create mode 100644 models/attack/attack2.py create mode 100644 models/attack/attack3.py create mode 100644 models/gnn_algo/generator.py create mode 100644 models/gnn_algo/surrogate.py create mode 100644 models/gnn_algo/victim.py create mode 100644 models/main.py create mode 100644 requirements.txt diff --git a/models/attack/attack1.py b/models/attack/attack1.py new file mode 100644 index 0000000..3f81619 --- /dev/null +++ b/models/attack/attack1.py @@ -0,0 +1,120 @@ +import torch +import torch.nn as nn +import torch.optim as optim +from tqdm import tqdm + +class TypeIAttack: + def __init__(self, generator, surrogate_model, victim_model, device, + noise_dim, num_nodes, feature_dim, + generator_lr=1e-6, surrogate_lr=0.001, + n_generator_steps=2, n_surrogate_steps=5): + self.generator = generator + self.surrogate_model = surrogate_model + self.victim_model = victim_model + self.device = device + self.noise_dim = noise_dim + self.num_nodes = num_nodes + self.feature_dim = feature_dim + + self.generator_optimizer = optim.Adam(self.generator.parameters(), lr=generator_lr) + self.surrogate_optimizer = optim.Adam(self.surrogate_model.parameters(), lr=surrogate_lr) + + self.criterion = nn.CrossEntropyLoss() + self.n_generator_steps = n_generator_steps + self.n_surrogate_steps = n_surrogate_steps + + def generate_graph(self): + z = torch.randn(1, self.noise_dim).to(self.device) + features, adj = self.generator(z) + edge_index = self.generator.adj_to_edge_index(adj) + return features, edge_index + + def train_generator(self): + self.generator.train() + self.surrogate_model.eval() + + total_loss = 0 + for _ in range(self.n_generator_steps): + self.generator_optimizer.zero_grad() + + features, edge_index = self.generate_graph() + + with torch.no_grad(): + victim_output = self.victim_model(features, edge_index) + surrogate_output = self.surrogate_model(features, edge_index) + + loss = -self.criterion(surrogate_output, victim_output.argmax(dim=1)) + + # Zeroth-order optimization with multiple random directions + epsilon = 1e-6 + num_directions = 2 + estimated_gradient = torch.zeros_like(features) + + for _ in range(num_directions): + u = torch.randn_like(features) + perturbed_features = features + epsilon * u + + with torch.no_grad(): + perturbed_victim_output = self.victim_model(perturbed_features, edge_index) + perturbed_surrogate_output = self.surrogate_model(perturbed_features, edge_index) + perturbed_loss = -self.criterion(perturbed_surrogate_output, perturbed_victim_output.argmax(dim=1)) + + estimated_gradient += (perturbed_loss - loss) / epsilon * u + + estimated_gradient /= num_directions + features.grad = estimated_gradient + + self.generator_optimizer.step() + total_loss += loss.item() + + return total_loss / self.n_generator_steps + + def train_surrogate(self): + self.generator.eval() + self.surrogate_model.train() + + total_loss = 0 + for _ in range(self.n_surrogate_steps): + self.surrogate_optimizer.zero_grad() + + features, edge_index = self.generate_graph() + + with torch.no_grad(): + victim_output = self.victim_model(features, edge_index) + surrogate_output = self.surrogate_model(features, edge_index) + + loss = self.criterion(surrogate_output, victim_output.argmax(dim=1)) + + loss.backward() + torch.nn.utils.clip_grad_norm_(self.surrogate_model.parameters(), max_norm=1.0) + self.surrogate_optimizer.step() + + total_loss += loss.item() + + return total_loss / self.n_surrogate_steps + + def attack(self, num_queries, log_interval=10): + generator_losses = [] + surrogate_losses = [] + + pbar = tqdm(range(num_queries), desc="Attacking") + for query in pbar: + gen_loss = self.train_generator() + surr_loss = self.train_surrogate() + + generator_losses.append(gen_loss) + surrogate_losses.append(surr_loss) + + if (query + 1) % log_interval == 0: + pbar.set_postfix({ + 'Gen Loss': f"{gen_loss:.4f}", + 'Surr Loss': f"{surr_loss:.4f}" + }) + + return self.surrogate_model, generator_losses, surrogate_losses + +def run_attack(generator, surrogate_model, victim_model, num_queries, device, + noise_dim, num_nodes, feature_dim): + attack = TypeIAttack(generator, surrogate_model, victim_model, device, + noise_dim, num_nodes, feature_dim) + return attack.attack(num_queries) diff --git a/models/attack/attack2.py b/models/attack/attack2.py new file mode 100644 index 0000000..8ec06f4 --- /dev/null +++ b/models/attack/attack2.py @@ -0,0 +1,103 @@ +import torch +import torch.nn as nn +import torch.optim as optim +from tqdm import tqdm + +class TypeIIAttack: + def __init__(self, generator, surrogate_model, victim_model, device, + noise_dim, num_nodes, feature_dim, + generator_lr=1e-6, surrogate_lr=0.001, + n_generator_steps=2, n_surrogate_steps=5): + self.generator = generator + self.surrogate_model = surrogate_model + self.victim_model = victim_model + self.device = device + self.noise_dim = noise_dim + self.num_nodes = num_nodes + self.feature_dim = feature_dim + + self.generator_optimizer = optim.Adam(self.generator.parameters(), lr=generator_lr) + self.surrogate_optimizer = optim.Adam(self.surrogate_model.parameters(), lr=surrogate_lr) + + self.criterion = nn.CrossEntropyLoss() + self.n_generator_steps = n_generator_steps + self.n_surrogate_steps = n_surrogate_steps + + def generate_graph(self): + z = torch.randn(1, self.noise_dim).to(self.device) + features, adj = self.generator(z) + edge_index = self.generator.adj_to_edge_index(adj) + return features, edge_index + + def train_generator(self): + self.generator.train() + self.surrogate_model.eval() + + total_loss = 0 + for _ in range(self.n_generator_steps): + self.generator_optimizer.zero_grad() + + features, edge_index = self.generate_graph() + + with torch.no_grad(): + victim_output = self.victim_model(features, edge_index) + surrogate_output = self.surrogate_model(features, edge_index) + + # In Type II, we use the surrogate model's gradient directly + loss = -self.criterion(surrogate_output, victim_output.argmax(dim=1)) + loss.backward() + + self.generator_optimizer.step() + total_loss += loss.item() + + return total_loss / self.n_generator_steps + + def train_surrogate(self): + self.generator.eval() + self.surrogate_model.train() + + total_loss = 0 + for _ in range(self.n_surrogate_steps): + self.surrogate_optimizer.zero_grad() + + features, edge_index = self.generate_graph() + + with torch.no_grad(): + victim_output = self.victim_model(features, edge_index) + surrogate_output = self.surrogate_model(features, edge_index) + + loss = self.criterion(surrogate_output, victim_output.argmax(dim=1)) + + loss.backward() + torch.nn.utils.clip_grad_norm_(self.surrogate_model.parameters(), max_norm=1.0) + self.surrogate_optimizer.step() + + total_loss += loss.item() + + return total_loss / self.n_surrogate_steps + + def attack(self, num_queries, log_interval=10): + generator_losses = [] + surrogate_losses = [] + + pbar = tqdm(range(num_queries), desc="Attacking") + for query in pbar: + gen_loss = self.train_generator() + surr_loss = self.train_surrogate() + + generator_losses.append(gen_loss) + surrogate_losses.append(surr_loss) + + if (query + 1) % log_interval == 0: + pbar.set_postfix({ + 'Gen Loss': f"{gen_loss:.4f}", + 'Surr Loss': f"{surr_loss:.4f}" + }) + + return self.surrogate_model, generator_losses, surrogate_losses + +def run_attack(generator, surrogate_model, victim_model, num_queries, device, + noise_dim, num_nodes, feature_dim): + attack = TypeIIAttack(generator, surrogate_model, victim_model, device, + noise_dim, num_nodes, feature_dim) + return attack.attack(num_queries) diff --git a/models/attack/attack3.py b/models/attack/attack3.py new file mode 100644 index 0000000..e57692d --- /dev/null +++ b/models/attack/attack3.py @@ -0,0 +1,115 @@ +import torch +import torch.nn as nn +import torch.optim as optim +from tqdm import tqdm + +class TypeIIIAttack: + def __init__(self, generator, surrogate_model1, surrogate_model2, victim_model, device, + noise_dim, num_nodes, feature_dim, + generator_lr=1e-6, surrogate_lr=0.001, + n_generator_steps=2, n_surrogate_steps=5): + self.generator = generator + self.surrogate_model1 = surrogate_model1 + self.surrogate_model2 = surrogate_model2 + self.victim_model = victim_model + self.device = device + self.noise_dim = noise_dim + self.num_nodes = num_nodes + self.feature_dim = feature_dim + + self.generator_optimizer = optim.Adam(self.generator.parameters(), lr=generator_lr) + self.surrogate_optimizer1 = optim.Adam(self.surrogate_model1.parameters(), lr=surrogate_lr) + self.surrogate_optimizer2 = optim.Adam(self.surrogate_model2.parameters(), lr=surrogate_lr) + + self.criterion = nn.CrossEntropyLoss() + self.n_generator_steps = n_generator_steps + self.n_surrogate_steps = n_surrogate_steps + + def generate_graph(self): + z = torch.randn(1, self.noise_dim).to(self.device) + features, adj = self.generator(z) + edge_index = self.generator.adj_to_edge_index(adj) + return features, edge_index + + def train_generator(self): + self.generator.train() + self.surrogate_model1.eval() + self.surrogate_model2.eval() + + total_loss = 0 + for _ in range(self.n_generator_steps): + self.generator_optimizer.zero_grad() + + features, edge_index = self.generate_graph() + + surrogate_output1 = self.surrogate_model1(features, edge_index) + surrogate_output2 = self.surrogate_model2(features, edge_index) + + # Compute disagreement loss + loss = -torch.mean(torch.std(torch.stack([surrogate_output1, surrogate_output2]), dim=0)) + loss.backward() + + self.generator_optimizer.step() + total_loss += loss.item() + + return total_loss / self.n_generator_steps + + def train_surrogate(self): + self.generator.eval() + self.surrogate_model1.train() + self.surrogate_model2.train() + + total_loss = 0 + for _ in range(self.n_surrogate_steps): + self.surrogate_optimizer1.zero_grad() + self.surrogate_optimizer2.zero_grad() + + features, edge_index = self.generate_graph() + + with torch.no_grad(): + victim_output = self.victim_model(features, edge_index) + surrogate_output1 = self.surrogate_model1(features, edge_index) + surrogate_output2 = self.surrogate_model2(features, edge_index) + + loss1 = self.criterion(surrogate_output1, victim_output.argmax(dim=1)) + loss2 = self.criterion(surrogate_output2, victim_output.argmax(dim=1)) + + # Combine losses and backpropagate once + combined_loss = loss1 + loss2 + combined_loss.backward() + + torch.nn.utils.clip_grad_norm_(self.surrogate_model1.parameters(), max_norm=1.0) + torch.nn.utils.clip_grad_norm_(self.surrogate_model2.parameters(), max_norm=1.0) + + self.surrogate_optimizer1.step() + self.surrogate_optimizer2.step() + + total_loss += combined_loss.item() / 2 + + return total_loss / self.n_surrogate_steps + + def attack(self, num_queries, log_interval=10): + generator_losses = [] + surrogate_losses = [] + + pbar = tqdm(range(num_queries), desc="Attacking") + for query in pbar: + gen_loss = self.train_generator() + surr_loss = self.train_surrogate() + + generator_losses.append(gen_loss) + surrogate_losses.append(surr_loss) + + if (query + 1) % log_interval == 0: + pbar.set_postfix({ + 'Gen Loss': f"{gen_loss:.4f}", + 'Surr Loss': f"{surr_loss:.4f}" + }) + + return (self.surrogate_model1, self.surrogate_model2), generator_losses, surrogate_losses + +def run_attack(generator, surrogate_model1, surrogate_model2, victim_model, num_queries, device, + noise_dim, num_nodes, feature_dim): + attack = TypeIIIAttack(generator, surrogate_model1, surrogate_model2, victim_model, device, + noise_dim, num_nodes, feature_dim) + return attack.attack(num_queries) diff --git a/models/gnn_algo/generator.py b/models/gnn_algo/generator.py new file mode 100644 index 0000000..ae8637e --- /dev/null +++ b/models/gnn_algo/generator.py @@ -0,0 +1,106 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch_geometric.nn import GCNConv + +class GraphGenerator(nn.Module): + def __init__(self, noise_dim, num_nodes, feature_dim, generator_type='cosine', threshold=0.1): + super(GraphGenerator, self).__init__() + self.noise_dim = noise_dim + self.num_nodes = num_nodes + self.feature_dim = feature_dim + self.generator_type = generator_type + self.threshold = threshold + + # Feature generator + self.feature_gen = nn.Sequential( + nn.Linear(noise_dim, 128), + nn.ReLU(), + nn.Linear(128, 256), + nn.ReLU(), + nn.Linear(256, num_nodes * feature_dim), + nn.Tanh() + ) + + # Full parameterization structure generator + if generator_type == 'full_param': + self.structure_gen = nn.Sequential( + nn.Linear(noise_dim, 128), + nn.ReLU(), + nn.Linear(128, 256), + nn.ReLU(), + nn.Linear(256, num_nodes * num_nodes), + nn.Sigmoid() + ) + + def forward(self, z): + # Generate features + features = self.feature_gen(z).view(self.num_nodes, self.feature_dim) + + # Generate adjacency matrix + if self.generator_type == 'cosine': + adj = self.cosine_similarity_generator(features) + elif self.generator_type == 'full_param': + adj = self.full_param_generator(z) + else: + raise ValueError("Invalid generator type. Choose 'cosine' or 'full_param'.") + + # Normalize adjacency matrix + adj = adj / adj.sum(1, keepdim=True).clamp(min=1) + + return features, adj + + def cosine_similarity_generator(self, features): + # Compute cosine similarity + norm_features = F.normalize(features, p=2, dim=1) + adj = torch.mm(norm_features, norm_features.t()) + + # Apply threshold + adj = (adj > self.threshold).float() + + # Remove self-loops + adj = adj * (1 - torch.eye(self.num_nodes, device=adj.device)) + + return adj + + def full_param_generator(self, z): + adj = self.structure_gen(z).view(self.num_nodes, self.num_nodes) + + # Make symmetric + adj = (adj + adj.t()) / 2 + + # Remove self-loops + adj = adj * (1 - torch.eye(self.num_nodes, device=adj.device)) + + return adj + + def adj_to_edge_index(self, adj): + return adj.nonzero().t() + + def self_supervised_training(self, x, adj, model): + # Implement self-supervised denoising task + self.train() + + # Add noise to features + noise = torch.randn_like(x) * 0.1 + noisy_x = x + noise + + # Use the model to denoise + edge_index = self.adj_to_edge_index(adj) + denoised_x = model(noisy_x, edge_index) + + # Compute reconstruction loss + loss = F.mse_loss(denoised_x, x) + + return loss + +class DenoisingModel(nn.Module): + def __init__(self, input_dim, hidden_dim): + super(DenoisingModel, self).__init__() + self.conv1 = GCNConv(input_dim, hidden_dim) + self.conv2 = GCNConv(hidden_dim, input_dim) + + def forward(self, x, edge_index): + x = F.relu(self.conv1(x, edge_index)) + x = self.conv2(x, edge_index) + return x diff --git a/models/gnn_algo/surrogate.py b/models/gnn_algo/surrogate.py new file mode 100644 index 0000000..7d7e4c1 --- /dev/null +++ b/models/gnn_algo/surrogate.py @@ -0,0 +1,43 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch_geometric.nn import GCNConv + +class SurrogateModel(nn.Module): + def __init__(self, input_dim, hidden_dim, output_dim, num_layers=2, dropout_rate=0.5): + super(SurrogateModel, self).__init__() + self.convs = nn.ModuleList() + self.convs.append(GCNConv(input_dim, hidden_dim)) + + for _ in range(num_layers - 2): + self.convs.append(GCNConv(hidden_dim, hidden_dim)) + + self.convs.append(GCNConv(hidden_dim, output_dim)) + self.dropout_rate = dropout_rate + + def forward(self, x, edge_index): + for i, conv in enumerate(self.convs[:-1]): + x = conv(x, edge_index) + x = F.relu(x) + x = F.dropout(x, p=self.dropout_rate, training=self.training) + + x = self.convs[-1](x, edge_index) + return F.softmax(x, dim=1) + + def train_step(self, generator, victim_model, optimizer, criterion, device): + self.train() + optimizer.zero_grad() + + z = torch.randn(1, generator.noise_dim).to(device) + features, adj = generator(z) + edge_index = generator.adj_to_edge_index(adj) + + with torch.no_grad(): + victim_output = victim_model(features, edge_index) + surrogate_output = self(features, edge_index) + + loss = criterion(surrogate_output, victim_output.argmax(dim=1)) + loss.backward() + optimizer.step() + + return loss.item() diff --git a/models/gnn_algo/victim.py b/models/gnn_algo/victim.py new file mode 100644 index 0000000..d15f2b0 --- /dev/null +++ b/models/gnn_algo/victim.py @@ -0,0 +1,49 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch_geometric.nn import GCNConv + +class VictimModel(nn.Module): + def __init__(self, input_dim, hidden_dim, output_dim, num_layers=2): + super(VictimModel, self).__init__() + self.convs = nn.ModuleList() + self.convs.append(GCNConv(input_dim, hidden_dim)) + + for _ in range(num_layers - 2): + self.convs.append(GCNConv(hidden_dim, hidden_dim)) + + self.convs.append(GCNConv(hidden_dim, output_dim)) + + def forward(self, x, edge_index): + for i, conv in enumerate(self.convs[:-1]): + x = conv(x, edge_index) + x = F.relu(x) + x = F.dropout(x, p=0.25, training=self.training) # Paper: p=0.5 + + x = self.convs[-1](x, edge_index) + return F.log_softmax(x, dim=1) + +def create_victim_model_cora(): + input_dim = 1433 + hidden_dim = 128 # Paper: 128 + output_dim = 7 + return VictimModel(input_dim, hidden_dim, output_dim) + +def create_victim_model_computers(): + input_dim = 767 + hidden_dim = 128 # Paper: 128 + output_dim = 10 + return VictimModel(input_dim, hidden_dim, output_dim) + +def create_victim_model_pubmed(): + input_dim = 500 + hidden_dim = 128 # Paper: 128 + output_dim = 3 + return VictimModel(input_dim, hidden_dim, output_dim) + +def create_victim_model_ogb_arxiv(): + input_dim = 128 + hidden_dim = 256 # Paper: 256 + output_dim = 40 + num_layers = 2 # Paper: 3 + return VictimModel(input_dim, hidden_dim, output_dim, num_layers) diff --git a/models/main.py b/models/main.py new file mode 100644 index 0000000..66c45e2 --- /dev/null +++ b/models/main.py @@ -0,0 +1,311 @@ +import sys +import torch +import torch.nn as nn +import torch.optim as optim +from torch_geometric.datasets import Planetoid, Amazon +from ogb.nodeproppred import PygNodePropPredDataset +from torch_geometric.transforms import NormalizeFeatures +from torch_geometric.utils import to_undirected +import matplotlib.pyplot as plt +import seaborn as sns +import numpy as np +from sklearn.metrics import accuracy_score, f1_score, confusion_matrix +from tqdm import tqdm +from reportlab.pdfgen import canvas +from reportlab.lib.pagesizes import letter +from reportlab.lib.units import inch + +# Import our custom modules +from gnn_algo.victim import create_victim_model_cora, create_victim_model_computers, create_victim_model_pubmed, create_victim_model_ogb_arxiv +from gnn_algo.generator import GraphGenerator +from attack.attack1 import TypeIAttack +from attack.attack2 import TypeIIAttack +from attack.attack3 import TypeIIIAttack + +def create_masks(num_nodes, train_ratio=0.6, val_ratio=0.2): + indices = np.random.permutation(num_nodes) + train_size = int(num_nodes * train_ratio) + val_size = int(num_nodes * val_ratio) + + train_mask = torch.zeros(num_nodes, dtype=torch.bool) + val_mask = torch.zeros(num_nodes, dtype=torch.bool) + test_mask = torch.zeros(num_nodes, dtype=torch.bool) + + train_mask[indices[:train_size]] = True + val_mask[indices[train_size:train_size+val_size]] = True + test_mask[indices[train_size+val_size:]] = True + + return train_mask, val_mask, test_mask + +def main(attack_type, dataset_name): + # Set device + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + # Load dataset and create victim model + if dataset_name == 'cora': + dataset = Planetoid(root='/tmp/Cora', name='Cora', transform=NormalizeFeatures()) + data = dataset[0].to(device) + victim_model = create_victim_model_cora().to(device) + elif dataset_name == 'computers': + dataset = Amazon(root='/tmp/Amazon', name='Computers', transform=NormalizeFeatures()) + data = dataset[0].to(device) + data.edge_index = to_undirected(data.edge_index) + train_mask, val_mask, test_mask = create_masks(data.num_nodes) + data.train_mask = train_mask.to(device) + data.val_mask = val_mask.to(device) + data.test_mask = test_mask.to(device) + victim_model = create_victim_model_computers().to(device) + elif dataset_name == 'pubmed': + dataset = Planetoid(root='/tmp/Pubmed', name='Pubmed', transform=NormalizeFeatures()) + data = dataset[0].to(device) + victim_model = create_victim_model_pubmed().to(device) + elif dataset_name == 'ogb-arxiv': + dataset = PygNodePropPredDataset(name='ogbn-arxiv', transform=NormalizeFeatures()) + data = dataset[0].to(device) + split_idx = dataset.get_idx_split() + data.train_mask = torch.zeros(data.num_nodes, dtype=torch.bool) + data.val_mask = torch.zeros(data.num_nodes, dtype=torch.bool) + data.test_mask = torch.zeros(data.num_nodes, dtype=torch.bool) + data.train_mask[split_idx['train']] = True + data.val_mask[split_idx['valid']] = True + data.test_mask[split_idx['test']] = True + data.train_mask = data.train_mask.to(device) + data.val_mask = data.val_mask.to(device) + data.test_mask = data.test_mask.to(device) + victim_model = create_victim_model_ogb_arxiv().to(device) + else: + raise ValueError("Invalid dataset name. Choose 'cora', 'computers', 'pubmed', or 'ogb-arxiv'.") + + # Train victim model + train_victim_model(victim_model, data, dataset_name) + + # Initialize generator and surrogate model(s) + noise_dim = 32 + num_nodes = 500 + feature_dim = dataset.num_features + output_dim = dataset.num_classes + + generator = GraphGenerator(noise_dim, num_nodes, feature_dim, generator_type='cosine').to(device) + + if dataset_name == 'cora': + surrogate_model1 = create_victim_model_cora().to(device) + elif dataset_name == 'computers': + surrogate_model1 = create_victim_model_computers().to(device) + elif dataset_name == 'pubmed': + surrogate_model1 = create_victim_model_pubmed().to(device) + elif dataset_name == 'ogb-arxiv': + surrogate_model1 = create_victim_model_ogb_arxiv().to(device) + + # Attack parameters + num_queries = 700 + generator_lr = 1e-6 + surrogate_lr = 0.001 + n_generator_steps = 2 + n_surrogate_steps = 5 + + # Run attack based on attack_type + print(f"Running attack type {attack_type} on {dataset_name} dataset...") + + if attack_type == 1: + attack = TypeIAttack(generator, surrogate_model1, victim_model, device, + noise_dim, num_nodes, feature_dim, generator_lr, surrogate_lr, + n_generator_steps, n_surrogate_steps) + elif attack_type == 2: + attack = TypeIIAttack(generator, surrogate_model1, victim_model, device, + noise_dim, num_nodes, feature_dim, generator_lr, surrogate_lr, + n_generator_steps, n_surrogate_steps) + elif attack_type == 3: + if dataset_name == 'cora': + surrogate_model2 = create_victim_model_cora().to(device) + elif dataset_name == 'computers': + surrogate_model2 = create_victim_model_computers().to(device) + elif dataset_name == 'pubmed': + surrogate_model2 = create_victim_model_pubmed().to(device) + elif dataset_name == 'ogb-arxiv': + surrogate_model2 = create_victim_model_ogb_arxiv().to(device) + + attack = TypeIIIAttack(generator, surrogate_model1, surrogate_model2, victim_model, device, + noise_dim, num_nodes, feature_dim, generator_lr, surrogate_lr, + n_generator_steps, n_surrogate_steps) + else: + raise ValueError("Invalid attack type. Choose 1, 2, or 3.") + + trained_surrogate, generator_losses, surrogate_losses = attack.attack(num_queries) + + # Evaluate models + accuracy, fidelity, f1, conf_matrix = evaluate_models(victim_model, trained_surrogate, data) + + # Calculate random baselines + random_accuracy, random_f1 = calculate_random_baselines(data) + + # Print and store stats + stats = { + "Dataset": dataset_name, + "Attack Type": attack_type, + "Accuracy": accuracy, + "Fidelity": fidelity, + "F1 Score": f1, + "Random Accuracy": random_accuracy, + "Random F1": random_f1, + "Accuracy Improvement": (accuracy - random_accuracy) / random_accuracy * 100, + "F1 Improvement": (f1 - random_f1) / random_f1 * 100 + } + + print_stats(stats) + + # Plot results + plot_confusion_matrix(conf_matrix, output_dim, attack_type, dataset_name) + plot_losses(generator_losses, surrogate_losses, attack_type, dataset_name) + + # Generate PDF report + generate_pdf_report(stats, conf_matrix, attack_type, dataset_name) + +def train_victim_model(model, data, dataset_name, epochs=200, lr=0.01, weight_decay=5e-4): + optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay) + model.train() + for epoch in range(epochs): + optimizer.zero_grad() + out = model(data.x, data.edge_index) + if dataset_name == 'ogb-arxiv': + loss = nn.functional.nll_loss(out[data.train_mask], data.y.squeeze()[data.train_mask]) + else: + loss = nn.functional.nll_loss(out[data.train_mask], data.y[data.train_mask]) + loss.backward() + optimizer.step() + + if (epoch + 1) % 10 == 0: + model.eval() + with torch.no_grad(): + val_out = model(data.x, data.edge_index) + if dataset_name == 'ogb-arxiv': + val_loss = nn.functional.nll_loss(val_out[data.val_mask], data.y.squeeze()[data.val_mask]) + val_acc = (val_out[data.val_mask].argmax(dim=1) == data.y.squeeze()[data.val_mask]).float().mean() + else: + val_loss = nn.functional.nll_loss(val_out[data.val_mask], data.y[data.val_mask]) + val_acc = (val_out[data.val_mask].argmax(dim=1) == data.y[data.val_mask]).float().mean() + model.train() + print(f'Epoch {epoch+1}/{epochs}, Train Loss: {loss.item():.4f}, Val Loss: {val_loss.item():.4f}, Val Acc: {val_acc.item():.4f}') + +def evaluate_models(victim_model, trained_surrogate, data): + victim_model.eval() + if isinstance(trained_surrogate, tuple): + surrogate_model1, surrogate_model2 = trained_surrogate + surrogate_model1.eval() + surrogate_model2.eval() + else: + surrogate_model = trained_surrogate + surrogate_model.eval() + + with torch.no_grad(): + victim_out = victim_model(data.x, data.edge_index) + if isinstance(trained_surrogate, tuple): + surrogate_out1 = surrogate_model1(data.x, data.edge_index) + surrogate_out2 = surrogate_model2(data.x, data.edge_index) + surrogate_out = (surrogate_out1 + surrogate_out2) / 2 # Simple ensemble + else: + surrogate_out = surrogate_model(data.x, data.edge_index) + + victim_preds = victim_out.argmax(dim=1) + surrogate_preds = surrogate_out.argmax(dim=1) + + accuracy = accuracy_score(victim_preds[data.test_mask].cpu(), surrogate_preds[data.test_mask].cpu()) + fidelity = accuracy_score(victim_preds[data.test_mask].cpu(), surrogate_preds[data.test_mask].cpu()) + f1 = f1_score(victim_preds[data.test_mask].cpu(), surrogate_preds[data.test_mask].cpu(), average='weighted') + conf_matrix = confusion_matrix(victim_preds[data.test_mask].cpu(), surrogate_preds[data.test_mask].cpu()) + + return accuracy, fidelity, f1, conf_matrix + +def plot_confusion_matrix(conf_matrix, num_classes, attack_type, dataset_name): + plt.figure(figsize=(10, 8)) + sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Blues') + plt.title(f'Confusion Matrix - Type {attack_type} Attack on {dataset_name}') + plt.xlabel('Predicted') + plt.ylabel('True') + plt.savefig(f'confusion_matrix_type{attack_type}_{dataset_name}.png') + plt.close() + +def plot_losses(generator_losses, surrogate_losses, attack_type, dataset_name): + plt.figure(figsize=(10, 5)) + plt.plot(generator_losses, label='Generator Loss') + plt.plot(surrogate_losses, label='Surrogate Loss') + plt.title(f'Losses over time - Type {attack_type} Attack on {dataset_name}') + plt.xlabel('Query') + plt.ylabel('Loss') + plt.legend() + plt.savefig(f'losses_over_time_type{attack_type}_{dataset_name}.png') + plt.close() + +def calculate_random_baselines(data): + num_classes = data.y.max().item() + 1 + random_preds = torch.randint(0, num_classes, data.y.shape).to(data.y.device) + random_accuracy = accuracy_score(data.y[data.test_mask].cpu(), random_preds[data.test_mask].cpu()) + random_f1 = f1_score(data.y[data.test_mask].cpu(), random_preds[data.test_mask].cpu(), average='weighted') + return random_accuracy, random_f1 + +def print_stats(stats): + for key, value in stats.items(): + if isinstance(value, float): + print(f"{key}: {value:.4f}") + else: + print(f"{key}: {value}") + +def generate_pdf_report(stats, conf_matrix, attack_type, dataset_name): + pdf_filename = f"type{attack_type}_attack_{dataset_name}_report.pdf" + c = canvas.Canvas(pdf_filename, pagesize=letter) + width, height = letter + + c.setFont("Helvetica-Bold", 16) + c.drawString(50, height - 50, f"Type {attack_type} Attack on {dataset_name} Report") + + c.setFont("Helvetica", 12) + y = height - 100 + for key, value in stats.items(): + if isinstance(value, float): + c.drawString(50, y, f"{key}: {value:.4f}") + else: + c.drawString(50, y, f"{key}: {value}") + y -= 20 + + # Add confusion matrix to the report + c.showPage() + c.setFont("Helvetica-Bold", 14) + c.drawString(50, height - 50, "Confusion Matrix") + + table_width = 400 + table_height = 300 + x_start = (width - table_width) / 2 + y_start = height - 100 - table_height + + cell_width = table_width / conf_matrix.shape[1] + cell_height = table_height / conf_matrix.shape[0] + + for i in range(conf_matrix.shape[0]): + for j in range(conf_matrix.shape[1]): + x = x_start + j * cell_width + y = y_start + (conf_matrix.shape[0] - 1 - i) * cell_height + c.rect(x, y, cell_width, cell_height) + c.setFont("Helvetica", 10) + c.drawString(x + 2, y + 2, str(conf_matrix[i, j])) + + c.save() + print(f"PDF report saved as {pdf_filename}") + +if __name__ == "__main__": + if len(sys.argv) != 3: + print("Usage: python main.py ") + print("attack_type: 1, 2, or 3") + print("dataset_name: cora, computers, pubmed, or ogb-arxiv") + sys.exit(1) + + try: + attack_type = int(sys.argv[1]) + if attack_type not in [1, 2, 3]: + raise ValueError + dataset_name = sys.argv[2] + if dataset_name not in ['cora', 'computers', 'pubmed', 'ogb-arxiv']: + raise ValueError + except ValueError: + print("Invalid input. Please choose attack type 1, 2, or 3 and dataset name 'cora', 'computers', 'pubmed', or 'ogb-arxiv'.") + sys.exit(1) + + main(attack_type, dataset_name) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..f21d561 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,11 @@ +torch==1.9.0 +torch-geometric==2.0.3 +torch-scatter==2.0.8 +torch-sparse==0.6.12 +numpy==1.21.2 +matplotlib==3.4.3 +seaborn==0.11.2 +scikit-learn==0.24.2 +tqdm==4.62.3 +reportlab==3.5.68 +ogb==1.3.2