Add files via upload

Author: Djerry-h

models/detector.py

@@ -0,0 +1,95 @@
+from models.gnn import *
+
+from sklearn.metrics import roc_auc_score, average_precision_score
+import numpy as np
+
+import time
+
+class BaseDetector(object):
+    def __init__(self, train_config, model_config, data):
+        self.model_config = model_config
+        self.train_config = train_config
+        self.data = data
+        model_config['in_feats'] = self.data.graph.ndata['feature'].shape[1]
+        graph = self.data.graph.to(self.train_config['device'])
+        self.labels = graph.ndata['label']
+        self.train_mask = graph.ndata['train_mask'].bool()
+        self.val_mask = graph.ndata['val_mask'].bool()
+        self.test_mask = graph.ndata['test_mask'].bool()
+        self.weight = (1 - self.labels[self.train_mask]).sum().item() / self.labels[self.train_mask].sum().item()
+        self.source_graph = graph
+        print(train_config['inductive'])
+        if train_config['inductive'] == False:
+            self.train_graph = graph
+            self.val_graph = graph
+        else:
+            self.train_graph = graph.subgraph(self.train_mask)
+            self.val_graph = graph.subgraph(self.train_mask+self.val_mask)
+        self.best_score = -1
+        # self.patience_knt = 0
+        
+    def train(self):
+        pass
+
+    def eval(self, labels, probs):
+        score = {}
+        with torch.no_grad():
+            if torch.is_tensor(labels):
+                labels = labels.cpu().numpy()
+            if torch.is_tensor(probs):
+                probs = probs.cpu().numpy()
+            score['AUROC'] = roc_auc_score(labels, probs)
+            score['AUPRC'] = average_precision_score(labels, probs)
+            labels = np.array(labels)
+            k = labels.sum()
+        score['RecK'] = sum(labels[probs.argsort()[-k:]]) / sum(labels)
+        return score
+
+
+class BaseGNNDetector(BaseDetector):
+    def __init__(self, train_config, model_config, data):
+        super().__init__(train_config, model_config, data)
+        gnn = globals()[model_config['model']]
+        model_config['in_feats'] = self.data.graph.ndata['feature'].shape[1]
+        self.model = gnn(**model_config).to(train_config['device'])
+       
+       
+    def train(self):
+        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.model_config['lr'])
+        train_labels, val_labels, test_labels = self.labels[self.train_mask], self.labels[self.val_mask], self.labels[self.test_mask]
+      
+        for e in range(self.train_config['epochs']):
+            self.model.train()
+            logits = self.model(self.train_graph)
+            loss = F.cross_entropy(logits[self.train_graph.ndata['train_mask']], train_labels,
+                                   weight=torch.tensor([1., self.weight], device=self.labels.device))
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            if self.model_config['drop_rate'] > 0 or self.train_config['inductive']:
+                self.model.eval()
+                logits = self.model(self.val_graph)
+            probs = logits.softmax(1)[:, 1]
+            val_score = self.eval(val_labels, probs[self.val_graph.ndata['val_mask']])
+           
+            if val_score[self.train_config['metric']] > self.best_score:
+                if self.train_config['inductive']:
+                    logits = self.model(self.source_graph)
+                    probs = logits.softmax(1)[:, 1]
+                self.patience_knt = 0
+                self.best_score = val_score[self.train_config['metric']]
+                test_score = self.eval(test_labels, probs[self.test_mask])
+                print('Epoch {}, Loss {:.4f}, Val AUC {:.4f}, PRC {:.4f}, RecK {:.4f}, test AUC {:.4f}, PRC {:.4f}, RecK {:.4f}'.format(
+                    e, loss, val_score['AUROC'], val_score['AUPRC'], val_score['RecK'],
+                    test_score['AUROC'], test_score['AUPRC'], test_score['RecK']))
+
+            else:
+                self.patience_knt += 1
+                if self.patience_knt > self.train_config['patience']:
+                    break
+        
+        return test_score
+
+

models/gnn.py

@@ -1 +1,97 @@
+import torch
+import torch.nn.functional as F
+import dgl.function as fn
+import dgl.nn.pytorch.conv as dglnn
+from torch import nn
+from transformer import Decoder
+from dgl.nn.pytorch.conv import EdgeConv
 
+class GCN(nn.Module):
+    def __init__(self, in_feats, h_feats=64, num_classes=2, num_layers=2, mlp_layers=1, dropout_rate=0.,
+                 activation='ReLU', **kwargs):
+        super().__init__()
+        self.h_feats = h_feats
+        self.layers = nn.ModuleList()
+        self.act = getattr(nn, activation)()
+        self.layers.append(dglnn.GraphConv(in_feats, h_feats, activation=self.act))
+        for i in range(num_layers-1):
+            self.layers.append(dglnn.GraphConv(h_feats, h_feats, activation=self.act))
+        self.mlp = MLP(h_feats, h_feats, num_classes, mlp_layers, dropout_rate)
+        self.dropout = nn.Dropout(dropout_rate) if dropout_rate > 0 else nn.Identity()
+
+    def forward(self, graph):
+        h = graph.ndata['feature']
+        for i, layer in enumerate(self.layers):
+            if i != 0:
+                h = self.dropout(h)
+            h = layer(graph, h)
+        h = self.mlp(h, False)
+        return h
+
+class MLP(nn.Module):
+    def __init__(self, in_feats, h_feats=32, num_classes=2, num_layers=2, dropout_rate=0, activation='ReLU', **kwargs):
+        super(MLP, self).__init__()
+        self.layers = nn.ModuleList()
+        self.act = getattr(nn, activation)()
+        if num_layers == 0:
+            return
+        if num_layers == 1:
+            self.layers.append(nn.Linear(in_feats, num_classes))
+        else:
+            self.layers.append(nn.Linear(in_feats, h_feats))
+            for i in range(1, num_layers-1):
+                self.layers.append(nn.Linear(h_feats, h_feats))
+            self.layers.append(nn.Linear(h_feats, num_classes))
+        self.dropout = nn.Dropout(dropout_rate) if dropout_rate > 0 else nn.Identity()
+
+    def forward(self, h, is_graph=True):
+        if is_graph:
+            h = h.ndata['feature']
+        for i, layer in enumerate(self.layers):
+            if i != 0:
+                h = self.dropout(h)
+            h = layer(h)
+            if i != len(self.layers)-1:
+                h = self.act(h)
+        return h
+
+
+class MGADN(nn.Module):
+    def __init__(self, in_feats, h_feats=64, n_head=8, n_layers=4, dropout_rate=0., **kwargs):
+        super().__init__()
+        self.attn_fn = nn.Tanh()
+        self.act_fn = nn.ReLU()
+        self.decoder = Decoder(in_feats=in_feats, h_feats=h_feats, n_head=n_head, dropout_rate=0., n_layers=n_layers)
+        self.filters3 = GCN(in_feats, h_feats=h_feats, num_classes=h_feats, num_layers=2, mlp_layers=2, dropout_rate=0., activation='ReLU')
+        
+        self.DMGNN = EdgeConv(in_feats, out_feat=h_feats)
+        
+        self.linear1 = nn.Linear(h_feats*2, h_feats)
+                                    
+        self.linear = nn.Sequential(nn.Linear(h_feats, h_feats),
+                                     self.attn_fn,
+                                     nn.Linear(h_feats, 2))
+        
+        self.gate_layer = nn.Linear(h_feats, h_feats)
+    
+    def forward(self, graph):
+        x = graph.ndata['feature']
+        h_list = []
+        x = x.to(torch.float32)
+    
+        out1 = self.decoder(x, graph)
+        out2 = self.filters3(graph)
+        
+        F = self.DMGNN(graph,x)
+
+        h_list.append(out1)
+        h_list.append(out2)
+        
+        res = torch.cat((h_list[0], h_list[1]), dim=1)
+        output = self.linear1(res)
+        
+        gate = torch.sigmoid(self.gate_layer(output))
+        
+        out = gate * output + (1 - gate) * F
+        result = self.linear(out)
+        return result