-
Notifications
You must be signed in to change notification settings - Fork 11
/
random_injection.py
57 lines (48 loc) · 2.52 KB
/
random_injection.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
import os.path as osp
import torch
import torch_geometric.transforms as T
from greatx.attack.injection import RandomInjection
from greatx.datasets import GraphDataset
from greatx.nn.models import GCN
from greatx.training import Trainer
from greatx.training.callbacks import ModelCheckpoint
from greatx.utils import split_nodes
dataset = 'Cora'
root = osp.join(osp.dirname(osp.realpath(__file__)), '../../..', 'data')
dataset = GraphDataset(root=root, name=dataset,
transform=T.LargestConnectedComponents())
data = dataset[0]
splits = split_nodes(data.y, random_state=15)
num_features = data.x.size(-1)
num_classes = data.y.max().item() + 1
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# ================================================================== #
# Before Attack #
# ================================================================== #
trainer_before = Trainer(GCN(num_features, num_classes), device=device)
ckp = ModelCheckpoint('model_before.pth', monitor='val_acc')
trainer_before.fit(data, mask=(splits.train_nodes, splits.val_nodes),
callbacks=[ckp])
logs = trainer_before.evaluate(data, splits.test_nodes)
print(f"Before attack\n {logs}")
# ================================================================== #
# Attacking #
# ================================================================== #
attacker = RandomInjection(data, device=device)
attacker.reset()
attacker.attack(10, feat_limits=(0, 1)) # for continuous features
# attacker.attack(10, feat_budgets=10) # for binary features
# ================================================================== #
# After evasion Attack #
# ================================================================== #
logs = trainer_before.evaluate(attacker.data(), splits.test_nodes)
print(f"After evasion attack\n {logs}")
# ================================================================== #
# After poisoning Attack #
# ================================================================== #
trainer_after = Trainer(GCN(num_features, num_classes), device=device)
ckp = ModelCheckpoint('model_after.pth', monitor='val_acc')
trainer_after.fit(attacker.data(), mask=(splits.train_nodes, splits.val_nodes),
callbacks=[ckp])
logs = trainer_after.evaluate(attacker.data(), splits.test_nodes)
print(f"After poisoning attack\n {logs}")