update test_pgd.py for evasion attack

examples/graph/test_pgd.py

@@ -10,32 +10,30 @@
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--seed', type=int, default=15, help='Random seed.')
-parser.add_argument('--epochs', type=int, default=200,
+parser.add_argument('--epochs', type=int, default=100,
                     help='Number of epochs to train.')
-parser.add_argument('--lr', type=float, default=0.01,
-                    help='Initial learning rate.')
-parser.add_argument('--weight_decay', type=float, default=5e-4,
-                    help='Weight decay (L2 loss on parameters).')
-parser.add_argument('--hidden', type=int, default=16,
-                    help='Number of hidden units.')
-parser.add_argument('--dropout', type=float, default=0.5,
-                    help='Dropout rate (1 - keep probability).')
-
 parser.add_argument('--dataset', type=str, default='citeseer', choices=['cora', 'cora_ml', 'citeseer', 'polblogs', 'pubmed'], help='dataset')
 parser.add_argument('--ptb_rate', type=float, default=0.05,  help='pertubation rate')
 parser.add_argument('--model', type=str, default='PGD', choices=['PGD', 'min-max'], help='model variant')
 
 args = parser.parse_args()
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 np.random.seed(args.seed)
 torch.manual_seed(args.seed)
 if device != 'cpu':
     torch.cuda.manual_seed(args.seed)
 
-data = Dataset(root='/tmp/', name=args.dataset, setting='gcn')
+# data = Dataset(root='/tmp/', name=args.dataset, setting='gcn')
+
+from torch_geometric.datasets import Planetoid
+from deeprobust.graph.data import Pyg2Dpr
+dataset = Planetoid('./', name=args.dataset)
+data = Pyg2Dpr(dataset)
+
 adj, features, labels = data.adj, data.features, data.labels
+
 features = normalize_feature(features)
 
 idx_train, idx_val, idx_test = data.idx_train, data.idx_val, data.idx_test
@@ -44,37 +42,23 @@
 
 adj, features, labels = preprocess(adj, features, labels, preprocess_adj=False)
 
-# Setup Victim Model
-victim_model = GCN(nfeat=features.shape[1], nclass=labels.max().item()+1, nhid=32,
-        dropout=0.5, weight_decay=5e-4, device=device)
-
-victim_model = victim_model.to(device)
-# victim_model.fit(features, adj, labels, idx_train)
-victim_model.fit(features, adj, labels, idx_train, idx_val, patience=30)
-
-# Setup Attack Model
-
-model = PGDAttack(model=victim_model, nnodes=adj.shape[0], loss_type='CE', device=device)
-
-model = model.to(device)
-
-def test(adj, gcn=None):
+def test(new_adj, gcn=None):
     ''' test on GCN '''
 
     if gcn is None:
         # adj = normalize_adj_tensor(adj)
         gcn = GCN(nfeat=features.shape[1],
-                  nhid=32,
+                  nhid=16,
                   nclass=labels.max().item() + 1,
-                  dropout=args.dropout, device=device)
+                  dropout=0.5, device=device)
         gcn = gcn.to(device)
-        # gcn.fit(features, adj, labels, idx_train) # train without model picking
-        gcn.fit(features, adj, labels, idx_train, idx_val, patience=10) # train with validation model picking
+        # gcn.fit(features, new_adj, labels, idx_train) # train without model picking
+        gcn.fit(features, new_adj, labels, idx_train, idx_val, patience=30) # train with validation model picking
         gcn.eval()
         output = gcn.predict().cpu()
     else:
         gcn.eval()
-        output = gcn.predict(features.to(device), adj.to(device)).cpu()
+        output = gcn.predict(features.to(device), new_adj.to(device)).cpu()
 
     loss_test = F.nll_loss(output[idx_test], labels[idx_test])
     acc_test = accuracy(output[idx_test], labels[idx_test])
@@ -87,18 +71,32 @@ def test(adj, gcn=None):
 
 def main():
     target_gcn = GCN(nfeat=features.shape[1],
-              nhid=32,
+              nhid=16,
               nclass=labels.max().item() + 1,
-              dropout=0.5, device=device)
+              dropout=0.5, device=device, lr=0.01)
 
     target_gcn = target_gcn.to(device)
     target_gcn.fit(features, adj, labels, idx_train, idx_val, patience=30)
+    # target_gcn.fit(features, adj, labels, idx_train)
 
     print('=== testing GCN on clean graph ===')
     test(adj, target_gcn)
 
+    # Setup Attack Model
+    print('=== setup attack model ===')
+    model = PGDAttack(model=target_gcn, nnodes=adj.shape[0], loss_type='CE', device=device)
+    model = model.to(device)
+
+    # model.attack(features, adj, labels, idx_train, perturbations, epochs=args.epochs)
+    # Here for the labels we need to replace it with predicted ones
+    fake_labels = target_gcn.predict(features.to(device), adj.to(device))
+    fake_labels = torch.argmax(fake_labels, 1).cpu()
+    # Besides, we need to add the idx into the whole process
+    idx_fake = np.concatenate([idx_train,idx_test])
+    model.attack(features, adj, fake_labels, idx_fake, perturbations, epochs=args.epochs)
+
     print('=== testing GCN on Evasion attack ===')
-    model.attack(features, adj, labels, idx_train, perturbations, epochs=args.epochs)
+
     modified_adj = model.modified_adj
     test(modified_adj, target_gcn)
 
@@ -112,3 +110,4 @@ def main():
 
 if __name__ == '__main__':
     main()
+