CITED/exp1.py at main · LabRAI/CITED · GitHub

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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
from models.factory import get_model_by_name
from pipline.cited import CITED
from pipline.defense import RandomWMPipeline, BackdoorWMPipeline, SurviveWMPipeline
from utils.dataset import CustomDataset, OriginDataset


def run_once_randomwm(config, trial_id=0):
    print(f'\n========== Trial {trial_id + 1} ==========')

    dataset = CustomDataset(config['ds_name'])
    data = dataset.get()
    dataset.stats()

    print('[Defense] Start defense')
    target_model = get_model_by_name(config['model_name'], data, config['hidden_dim'])
    defense_pipe = RandomWMPipeline(target_model, data, device=config['device'])
    wm_data = defense_pipe.embed_watermark_trigger(random_node_num=config['random_node_num'],
                                                   random_edge_prob=config['random_edge_prob'],
                                                   random_feat_ratio=config['random_feat_ratio'])
    defense_pipe.finetune_on_watermarked_data(wm_data, epochs=config['finetune_epochs'], lr=config['lr'],
                                              weight_decay=config['weight_decay'])


def run_once_backdoorwm(config, trial_id=0):
    print(f'\n========== Trial {trial_id + 1} ==========')

    dataset = CustomDataset(config['ds_name'])
    data = dataset.get()
    dataset.stats()

    print('[Defense] Start defense')
    target_model = get_model_by_name(config['model_name'], data, config['hidden_dim'])
    defense_pipe = BackdoorWMPipeline(target_model, data, device=config['device'])
    wm_data = defense_pipe.embed_backdoor(backdoor_ratio=config['backdoor_ratio'], backdoor_len=config['backdoor_len'])
    print('wm num: ', wm_data.wm_mask.sum(), 'train num: ', wm_data.train_mask.sum())
    defense_pipe.finetune_on_backdoor_data(wm_data, epochs=config['finetune_epochs'], lr=config['lr'],
                                           weight_decay=config['weight_decay'])


def run_once_surviveWM(config, trial_id=0):
    print(f'\n========== Trial {trial_id + 1} ==========')

    dataset = CustomDataset(config['ds_name'])
    data = dataset.get()
    dataset.stats()

    print('[Defense] Start defense')
    target_model = get_model_by_name(config['model_name'], data, config['hidden_dim'])
    defense_pipe = SurviveWMPipeline(target_model, data, device=config['device'])
    wm_data = defense_pipe.embed_wm(survive_node_num=config['survive_node_num'],
                                    survive_edge_prob=config['survive_edge_prob'], )
    print('wm num: ', wm_data.wm_mask.sum(), 'train num: ', wm_data.train_mask.sum())
    defense_pipe.finetune_on_wm_data(wm_data, epochs=config['finetune_epochs'], lr=config['lr'],
                                     weight_decay=config['weight_decay'])


def run_once_cited(config, trial_id=0):
    print(f'\n========== Trial {trial_id + 1} ==========')

    dataset = CustomDataset(config['ds_name'])
    data = dataset.get()
    dataset.stats()

    print('[Defense] Start defense')
    target_model = get_model_by_name(config['model_name'], data, config['hidden_dim'])
    defense_pipe = CITED(target_model, data, device=config['device'])
    wm_data = defense_pipe.signature(cited_boundary_ratio=config['cited_boundary_ratio'],
                                     cited_signature_ratio=config['cited_signature_ratio'])
    defense_pipe.finetune_signature(wm_data, epochs=config['finetune_epochs'], lr=config['lr'],
                                    weight_decay=config['weight_decay'])


def run_target_pipeline(model_name, ds_name):
    from pipline.target import TargetPipeline
    from models.fagcn import FAGCN
    from models.gat import GAT
    from models.factory import get_model_by_name
    device = 'cuda:0'
    dataset = OriginDataset(ds_name)
    data = dataset.get()
    dataset.stats()
    # target_model = GCN(in_feats=data.num_features, out_feats=data.num_classes, hidden_dim=128)
    target_model = get_model_by_name(model_name, data, hidden_dim=128)
    target_model = GAT(in_feats=data.num_features, out_feats=data.num_classes, hidden_dim=128)
    target_pipeline = TargetPipeline(target_model, data, device=device, lr=0.001, weight_decay=1e-5)
    target_pipeline.run(3)


if __name__ == '__main__':
    config = {
        # exp setting
        'level': 'label',
        'variant_num': 15,
        # model setting
        'model_name': 'gat',
        'hidden_dim': 128,
        # train setting
        'train_epochs': 200,
        'lr': 0.001,
        'weight_decay': 1e-5,
        # attack setting
        'query_ratio': 0.5,
        # misc
        'fixed_seed': 42,
        'device': 'cuda:0',
        # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        'ds_name': 'pubmed',
        'finetune_epochs': 10,
        # Defense RandomWM
        'random_node_num': 100,
        'random_edge_prob': 0.5,
        'random_feat_ratio': 0.1,
        # Defense BackdoorWM
        'backdoor_ratio': 0.2,
        'backdoor_len': 40,  # 20-40
        # Defense SurviveWM
        'survive_node_num': 300,
        'survive_edge_prob': 0.5,  # 0.5
        # Defense CITED
        'cited_boundary_ratio': 0.1,
        'cited_signature_ratio': 0.5,
    }

    run_once_randomwm(config)
    # run_once_backdoorwm(config)
    # run_once_surviveWM(config)
    # run_once_cited(config)

    model_name = 'gat'
    ds_name = 'pubmed'
    # run_target_pipeline(model_name, ds_name)