【节点分类】python实现：4种(GNN,GAN,SAGE,APPNP)图神经网络（-dgl库-pytorch-cuda-）_sagegnn结构示意图

信息系统建模作业，要求是使用四种不同的节点表征方法两个3k+数据集
环境：pytorch cuda11.1 dgl-0.6.1
（cuda环境配置指路：我发的第一篇文章）

dgl库各种网络功能介绍：
https://zhuanlan.zhihu.com/p/161853672

参考代码资料：
GNN
GAT
SAGE

数据集：CiteSeer PubMed

GraphConv基础图卷积

import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl.data import citation_graph as citegrh
from dgl.nn.pytorch.conv.graphconv import GraphConv

#加载数据集，另一个数据集使用citegrh.load_pubmed()
dataset = citegrh.load_citeseer()
print('Number of categories:', dataset.num_classes)

#dataset是list，数据集为单图，g是graph类
g = dataset[0]
#数据放在gpu上
g = g.to('cuda')

#构造网络
class GCN(nn.Module):
    def __init__(self, in_feats, h_feats, num_classes):
        super(GCN, self).__init__()
        self.conv1 = GraphConv(in_feats, h_feats) #特征维度映射
        self.conv2 = GraphConv(h_feats, num_classes) #映射到节点类别

    def forward(self, g, in_feat):
        # g：数据Graph信息，经过归一化的邻接矩阵
        # in_feat：节点初始化特征信息
        h = self.conv1(g, in_feat)
        h = F.relu(h)
        h = self.conv2(g, h)
        return h

# 实例化GCN模型在gpu上，hidden维度设定为16，输入维度和输出维度由数据集确定。
model = GCN(g.ndata['feat'].shape[1], 16, dataset.num_classes).to('cuda')

#训练网络
def train(g, model):
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    best_val_acc = 0
    best_test_acc = 0

    #.ndata属性是一个dict，含五个key
    features = g.ndata['feat']
    labels = g.ndata['label']
    train_mask = g.ndata['train_mask']
    val_mask = g.ndata['val_mask']
    test_mask = g.ndata['test_mask']
    for e in range(100):
        # Forward
        logits = model(g, features)

        # Compute prediction
        pred = logits.argmax(1) #每个样本全连接输出中最大的值，预测值

        # Compute loss
        loss = F.cross_entropy(logits[train_mask], labels[train_mask])

        # Compute accuracy on training/validation/test
        train_acc = (pred[train_mask] == labels[train_mask]).float().mean()
        val_acc = (pred[val_mask] == labels[val_mask]).float().mean()
        test_acc = (pred[test_mask] == labels[test_mask]).float().mean()

        # Save the best validation accuracy and the corresponding test accuracy.
        if best_val_acc < val_acc:
            best_val_acc = val_acc
            best_test_acc = test_acc

        # Backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if e % 5 == 0:
            print('In epoch {}, loss: {:.3f}, val acc: {:.3f} (best {:.3f}), test acc: {:.3f} (best {:.3f})'.format(
                e, loss, val_acc, best_val_acc, test_acc, best_test_acc))


train(g, model)

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

GATConv注意力机制

import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl.data import citation_graph as citegrh
import time

def load_data():
    data = citegrh.load_pubmed() #list
    g = data[0] # 单图
    g = g.to('cuda') #放到gpu上
    features = g.ndata['feat']
    labels = g.ndata['label']
    train_mask = g.ndata['train_mask']
    test_mask = g.ndata['test_mask']
    val_mask = g.ndata['val_mask']

    return g, features, labels, train_mask, test_mask, val_mask

#构造注意力图网络
class GATLayer(nn.Module):
    def __init__(self, g, in_dim, out_dim):
        super(GATLayer, self).__init__()
        self.g = g
        # step1 : 下层嵌入的线性变换
        self.fc = nn.Linear(in_dim, out_dim, bias=False)
        # step2 : 计算邻居节点对的非标准化注意力得分
        self.attn_fc = nn.Linear(2 * out_dim, 1, bias=False)

    def edge_attention(self, edges):
        # 计算step2需要的edge UDF
        z2 = torch.cat([edges.src['z'], edges.dst['z']], dim=1)
        a = self.attn_fc(z2)
        return {'e': F.leaky_relu(a)}

    def message_func(self, edges):
        # step 3 & 4需要的message UDF
        return {'z': edges.src['z'], 'e': edges.data['e']}

    def reduce_func(self, nodes):
        # 3 & 4 reduce UDF
        # step3 : 标准化每个节点的入边的注意力得分
        alpha = F.softmax(nodes.mailbox['e'], dim=1)
        # step4 : 聚合邻居节点的嵌入信息，并按照注意力得分缩放
        h = torch.sum(alpha * nodes.mailbox['z'], dim=1)
        return {'h': h}

    def forward(self, h):
        # step1
        z = self.fc(h)
        self.g.ndata['z'] = z
        # step2
        self.g.apply_edges(self.edge_attention)
        # step 3 & 4
        self.g.update_all(self.message_func, self.reduce_func)
        return self.g.ndata.pop('h')

# 用单头注意力构建多头注意力层
class MultiHeadGATLayer(nn.Module):
    def __init__(self, g, in_dim, out_dim, num_heads, merge='cat'):
        super(MultiHeadGATLayer, self).__init__()
        self.heads = nn.ModuleList()
        for i in range(num_heads):
            self.heads.append(GATLayer(g, in_dim, out_dim))
        self.merge = merge

    def forward(self, h):
        head_outs = [attn_head(h) for attn_head in self.heads]
        if self.merge == 'cat':
            # concat 输出特征
            return torch.cat(head_outs, dim=1)
        else:
            # 聚合信息，用mean方法
            return torch.mean(torch.stack(head_outs))

class GAT(nn.Module):
    def __init__(self, g, in_dim, hidden_dim, out_dim, num_heads):
        super(GAT, self).__init__()
        self.layer1 = MultiHeadGATLayer(g, in_dim, hidden_dim, num_heads)
        # 输入维度是hidden_dim * num_heads，输出是多头连接的
        # 输出层是一个注意力头
        self.layer2 = MultiHeadGATLayer(g, hidden_dim * num_heads, out_dim, 1)

    def forward(self, h):
        h = self.layer1(h)
        h = F.elu(h)
        h = self.layer2(h)
        return h

# 加载数据
g, features, labels, train_mask, test_mask, val_mask = load_data()

# 实例化网络，两个注意力头，每个头的 hidden size = 8
net = GAT(g,in_dim=features.size()[1],hidden_dim=8,out_dim=7,num_heads=2).to('cuda')

# create optimizer
optimizer = torch.optim.Adam(net.parameters(), lr=1e-2)
best_val_acc = 0
best_test_acc = 0

# main loop
dur = []
for epoch in range(100):
    if epoch >= 3:
        t0 = time.time()

    logits = net(features)
    # Compute prediction
    pred = logits.argmax(1)
    # Compute loss
    loss = F.cross_entropy(logits[train_mask], labels[train_mask])
    # Compute accuracy on training/validation/test
    train_acc = (pred[train_mask] == labels[train_mask]).float().mean()
    val_acc = (pred[val_mask] == labels[val_mask]).float().mean()
    test_acc = (pred[test_mask] == labels[test_mask]).float().mean()

    # Save the best validation accuracy and the corresponding test accuracy.
    if best_val_acc < val_acc:
        best_val_acc = val_acc
        best_test_acc = test_acc

    # Backward
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch % 5 == 0:
        print('In epoch {}, loss: {:.3f}, val acc: {:.3f} (best {:.3f}), test acc: {:.3f} (best {:.3f})'.format(
            epoch, loss, val_acc, best_val_acc, test_acc, best_test_acc))
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

SAGEConv采样聚类适用于大规模图

import dgl
import numpy as np
import torch
import torch as th
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from dgl.nn.pytorch.conv.sageconv import SAGEConv
import time
from dgl.data import citation_graph
import tqdm

# 邻居采样器
class NeighborSampler(object):
    def __init__(self, g, fanouts):
        """
        g 为 DGLGraph；
        fanouts 为采样节点的数量，两个参数，分别指一阶邻居采样和二阶邻居采样数
        """
        self.g = g
        self.fanouts = fanouts

    def sample_blocks(self, seeds):
        seeds = th.LongTensor(np.asarray(seeds))
        blocks = []
        for fanout in self.fanouts:
            # 对每个节点实现邻居采样，返回对应子图
            # replace=True 用采样节点代替所有邻居节点
            frontier = dgl.sampling.sample_neighbors(g, seeds, fanout, replace=True)
            # 将子图frontier转为可以传递消息的二部图
            block = dgl.to_block(frontier, seeds)
            # 获取新图的源节点作为种子节点
            # srcdata表示源节点信息，因采样和聚合是相反的操作过程
            seeds = block.srcdata[dgl.NID]
            blocks.insert(0, block)
        return blocks

# 构建适用于大图的sage网络
class GraphSAGE(nn.Module):
    def __init__(self,
                 in_feats,
                 n_hidden,
                 n_classes,
                 n_layers,
                 activation,
                 dropout):
        super().__init__()
        self.n_layers = n_layers
        self.n_hidden = n_hidden
        self.n_classes = n_classes
        self.layers = nn.ModuleList()
        # 输入层，四种聚合方式，此处使用"gcn"
        # 参数 : 输入特征，输出特征，聚合方式（特征drop概率=0，输出层bias=True，归一化=None，更新节点特征的激活函数=None）
        self.layers.append(SAGEConv(in_feats, n_hidden, 'gcn'))
        # 隐藏层
        for i in range(1, n_layers - 1):
            self.layers.append(SAGEConv(n_hidden, n_hidden, 'gcn'))
        # 输出层
        self.layers.append(SAGEConv(n_hidden, n_classes, 'gcn'))
        self.dropout = nn.Dropout(dropout)
        self.activation = activation

    # block是采样获得的子图（二部图）
    # x表示节点特征
    def forward(self, blocks, x):
        h = x
        for l, (layer, block) in enumerate(zip(self.layers, blocks)):
            h_dst = h[:block.number_of_dst_nodes()]
            h = layer(block, (h, h_dst))
            if l != len(self.layers) - 1:
                h = self.activation(h)
                h = self.dropout(h)
        return h

    # 用于评估测试
    def inference(self, g, x, batch_size):
        nodes = th.arange(g.number_of_nodes())
        for l, layer in enumerate(self.layers):
            y = th.zeros(g.number_of_nodes(),
                         self.n_hidden if l != len(self.layers) - 1 else self.n_classes)
            for start in tqdm.trange(0, len(nodes), batch_size):
                end = start + batch_size
                batch_nodes = nodes[start:end]
                block = dgl.to_block(dgl.in_subgraph(g, batch_nodes), batch_nodes)
                input_nodes = block.srcdata[dgl.NID]
                h = x[input_nodes]
                h_dst = h[:block.number_of_dst_nodes()]
                h = layer(block, (h, h_dst))
                if l != len(self.layers) - 1:
                    h = self.activation(h)
                    h = self.dropout(h)
                y[start:end] = h.cpu()
            x = y
        return y

# 计算准确率
def evaluate(model, g, inputs, labels, val_mask, batch_size):
    """
    评估模型，调用 model 的 inference 函数
    """
    model.eval()
    with torch.no_grad():
        label_pred = model.inference(g, inputs, batch_size)
    model.train()
    return (torch.argmax(label_pred[val_mask], dim=1) == labels[val_mask]).float().sum() / len(label_pred[val_mask])

def load_subtensor(g, labels, seeds, input_nodes, device):
    """
    将一组节点的特征和标签复制到 GPU 上。
    """
    batch_inputs = g.ndata['features'][input_nodes].to(device)
    batch_labels = labels[seeds].to(device)
    return batch_inputs, batch_labels

# 参数设置
gpu = 0
num_epochs = 50
num_hidden = 16
num_layers = 2
fan_out = '5,10'# 一阶采样5，二阶10
batch_size = 50
log_every = 1  # 记录日志的频率
eval_every = 5
lr = 0.001
dropout = 0.5
num_workers = 0

if gpu >= 0:
    device = th.device('cuda:%d' % gpu)
else:
    device = th.device('cpu')

# 读取数据，另一个数据集使用citation_graph.CiteseerGraphDataset()
data = citation_graph.PubmedGraphDataset()
graph = data[0] # 取单图
train_mask = graph.ndata['train_mask']
val_mask = graph.ndata['val_mask']
test_mask = graph.ndata['test_mask']
features = graph.ndata['feat']
in_feats = features.shape[1]
labels = graph.ndata['label']
n_classes = data.num_classes

# Construct graph
g = data.__getitem__(0) #为DGL.Graph类
g.ndata['features'] = features

train_nid = th.LongTensor(np.nonzero(train_mask)[0])
val_nid = th.LongTensor(np.nonzero(val_mask)[0])
train_mask = th.BoolTensor(train_mask)
val_mask = th.BoolTensor(val_mask)
test_mask = th.BoolTensor(test_mask)

# Create sampler
sampler = NeighborSampler(g, [int(fanout) for fanout in fan_out.split(',')])

# Create PyTorch DataLoader for constructing blocks
dataloader = DataLoader(
    dataset=train_nid.numpy(),
    batch_size=batch_size,
    collate_fn=sampler.sample_blocks,
    shuffle=True,
    drop_last=False,
    num_workers=num_workers)

# Define model and optimizer
model = GraphSAGE(in_feats, num_hidden, n_classes, num_layers, F.relu, dropout)
loss_fcn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)

# Training loop
best_val_acc = 0

for epoch in range(num_epochs):
    for step, blocks in enumerate(dataloader):

        tic_step = time.time()

        input_nodes = blocks[0].srcdata[dgl.NID]  #特殊二部图，输入节点的特征向量
        seeds = blocks[-1].dstdata[dgl.NID] #输出节点的标签

        # Load the input features as well as output labels
        batch_inputs, batch_labels = load_subtensor(g, labels, seeds, input_nodes, device) #输入节点的特征和标签

        # Compute loss and prediction
        batch_inputs = batch_inputs.to("cpu")
        batch_labels = batch_labels.to("cpu")
        batch_pred = model(blocks, batch_inputs).to(device)
        batch_pred = batch_pred.to("cpu")
        loss = loss_fcn(batch_pred, batch_labels).to(device)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        print('| Loss {:.4f}'.format(loss.item()))

        if epoch % 5 == 0:
            train_acc = evaluate(model, g, g.ndata['features'], labels, train_mask, batch_size)
            val_acc = evaluate(model, g, g.ndata['features'], labels, val_mask, batch_size)
            print('Epoch {:05d} | Val ACC {:.4f} | Train ACC {:.4f}'.format(epoch, val_acc.item(), train_acc.item()))

# 模型训练完毕，检查test集合的acc
acc_test = evaluate(model, g, g.ndata['features'], labels, test_mask, batch_size)
print('Test ACC: %.4f' % (acc_test.item()))






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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212

APPNPConv-PageRank的改良版本

import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl.data import citation_graph as citegrh
from dgl.nn.pytorch.conv.graphconv import GraphConv
from dgl.nn.pytorch.conv.appnpconv import APPNPConv

dataset = citegrh.load_citeseer() #更换数据集使用citegrh.load_pubmed()
print('Number of categories:', dataset.num_classes)

# 单图
g = dataset[0]
# 数据放在gpu上
g = g.to('cuda')

# 构造APPNP网络，该网络使用类似PageRank框架更新节点嵌入
class APPNP(nn.Module):
    def __init__(self, in_feats, h_feats, num_classes):
        super(APPNP, self).__init__()
        self.conv1 = GraphConv(in_feats, h_feats)
        # APPNP层，特征维度不变。
        '''
         k : int, The number of iterations 
         alpha : float, The teleport probability
        '''
        self.conv2 = APPNPConv(k=3, alpha=0.5)
        # 维度映射到节点类别
        self.conv3 = GraphConv(h_feats, num_classes)

    def forward(self, g, in_feat):
        h1 = self.conv1(g, in_feat)
        h2 = self.conv2(g, h1)
        h = self.conv2(g, h2)
        return h

# 实例化到gpu，hidden维度设定为16，输入维度和输出维度由数据集确定。
model = APPNP(g.ndata['feat'].shape[1], 16, dataset.num_classes).to('cuda')


def train(g, model):
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    best_val_acc = 0
    best_test_acc = 0

    features = g.ndata['feat']
    labels = g.ndata['label']
    train_mask = g.ndata['train_mask']
    val_mask = g.ndata['val_mask']
    test_mask = g.ndata['test_mask']
    for e in range(3000):
        # Forward
        logits = model(g, features)

        # Compute prediction
        pred = logits.argmax(1) #每个样本全连接输出中最大的值，预测值

        # Compute loss
        loss = F.cross_entropy(logits[train_mask], labels[train_mask])

        # Compute accuracy on training/validation/test
        train_acc = (pred[train_mask] == labels[train_mask]).float().mean()
        val_acc = (pred[val_mask] == labels[val_mask]).float().mean()
        test_acc = (pred[test_mask] == labels[test_mask]).float().mean()

        # Save the best validation accuracy and the corresponding test accuracy.
        if best_val_acc < val_acc:
            best_val_acc = val_acc
            best_test_acc = test_acc

        # Backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if e % 5 == 0:
            print('In epoch {}, loss: {:.3f}, val acc: {:.3f} (best {:.3f}), test acc: {:.3f} (best {:.3f})'.format(
                e, loss, val_acc, best_val_acc, test_acc, best_test_acc))


train(g, model)

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

分类结果