gat/models.py at master · ChengyuSun/gat · 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
import sys
sys.path.append('../')
import torch
import torch.nn as nn
import torch.nn.functional as F
from layers import GraphAttentionLayer, SpGraphAttentionLayer,MyLayer,OneLayer,CovLayer,MotifentropyLayer
import numpy as np
from utils import read_entropy_attention_list,read_entropy_vector_list
from utils import read_csv
from utils import read_txt
class GAT(nn.Module):
    def __init__(self, nfeat, nhid, nclass, dropout, alpha, nheads):
        """Dense version of GAT."""
        super(GAT, self).__init__()
        self.dropout = dropout

        #trainable--attention
        #self.attentions = [GraphAttentionLayer(nfeat, nhid, dropout=dropout, alpha=alpha, concat=True) for _ in
         #                  range(nheads)]

        #entropy--attention
        attentionlist=read_entropy_attention_list()
        self.attentions = [MyLayer(nfeat, nhid, attentionlist[i] ,dropout=dropout,concat=True) for i in range(nheads)]

        #motif-entropy
        # motif_entropy_list=read_entropy_vector_list()
        # self.attentions = [MotifentropyLayer(nfeat, nhid,dropout=dropout, alpha=alpha,
        #                                      entropy_vector=motif_entropy_list[i],  concat=True) for i in range(nheads)]

        #adj_citeseer = read_txt()
        #adj_cora=read_csv()

        #gnn-layers
        #self.attentions = [OneLayer(nfeat, nhid, dropout=dropout, adj=adj_cora,concat=True) for i in range(nheads)]

        #simple--gnn
        #self.simpleLayer=OneLayer(nfeat, nclass, dropout=dropout, adj=read_csv(),concat=False)

        #simple--attenetion
        #self.simpleLayer=MyLayer(nfeat, nclass, attentionlist[0] ,dropout=dropout,concat=False)


        if hasattr(self,'attentions'):
            for i, attention in enumerate(self.attentions):
                print('add {} layer to model'.format(i))
                self.add_module('attention_{}'.format(i), attention)
            print('add out layer to model')
            #self.out_att = GraphAttentionLayer(nhid * nheads, nclass, dropout=dropout, alpha=alpha, concat=False)
            self.out_att = CovLayer(nhid * nheads, nclass)
        elif hasattr(self, 'simpleLayer'):
            print('add simple layer into model')
            self.add_module('simple', self.simpleLayer)


    def forward(self, x, adj):
        x = F.dropout(x, self.dropout, training=self.training)

        if hasattr(self, 'attentions'):
            x = torch.cat([att(x, adj) for att in self.attentions], dim=1)
            x = F.dropout(x, self.dropout, training=self.training)
            x = F.elu(self.out_att(x, adj))

        elif hasattr(self, 'simpleLayer'):
            x=F.elu(self.simpleLayer(x,adj))

        return F.log_softmax(x, dim=1)

    def show(self):
        w_list=[]
        a_list=[]
        for i, attention in enumerate(self.attentions):
            w_list.append(attention.W)
            a_list.append(attention.a)
        return w_list,a_list

class SpGAT(nn.Module):
    def __init__(self, nfeat, nhid, nclass, dropout, alpha, nheads):
        """Sparse version of GAT."""
        super(SpGAT, self).__init__()
        self.dropout = dropout

        self.attentions = [SpGraphAttentionLayer(nfeat,
                                                 nhid,
                                                 dropout=dropout,
                                                 alpha=alpha,
                                                 concat=True) for _ in range(nheads)]
        for i, attention in enumerate(self.attentions):
            self.add_module('attention_{}'.format(i), attention)

        self.out_att = SpGraphAttentionLayer(nhid * nheads,
                                             nclass,
                                             dropout=dropout,
                                             alpha=alpha,
                                             concat=False)

    def forward(self, x, adj):
        x = F.dropout(x, self.dropout, training=self.training)
        x = torch.cat([att(x, adj) for att in self.attentions], dim=1)
        x = F.dropout(x, self.dropout, training=self.training)
        x = F.elu(self.out_att(x, adj))
        return F.log_softmax(x, dim=1)