PYG教程【二】使用networkx创建一个PYG图

上面通过torch_geometric自带的方式创建了一个图，本节使用networkx包来创建一个图，代码如下：

import numpy as np
import networkx as nx
import matplotlib.pyplot as plt

G = nx.Graph()
edgelist = [(0,1),(0,2),(0,3),(1,3)]
G.add_edges_from(edgelist)

fig = plt.figure()
ax = fig.add_subplot(1,1,1)
option = {'font_family':'serif', 'font_size':'15', 'font_weight':'semibold'}
nx.draw_networkx(G, node_size=400, **option)
plt.show()

创建的边：

首先需要安装我们的测试数据pip install python-louvain -i https://pypi.douban.com/simple/

import torch
from torch_geometric.data import InMemoryDataset, Data
import community as community_louvain

# x是节点特征矩阵，这里设为单位矩阵。
x = torch.eye(G.number_of_nodes(), dtype=torch.float)

# adj是图G的邻接矩阵的稀疏表示，左边节点对代表一条边，右边是边的值，adj是对称矩阵。
adj = nx.to_scipy_sparse_matrix(G).tocoo()

# row是adj中非零元素所在的行索引
row = torch.from_numpy(adj.row.astype(np.int64)).to(torch.long)
# col是adj中非零元素所在的列索引。
col = torch.from_numpy(adj.col.astype(np.int64)).to(torch.long)

# 将行和列进行拼接，shape变为[2, num_edges], 包含两个列表，第一个是row, 第二个是col
edge_index = torch.stack([row, col], dim=0)

# Compute communities.用louvain算法对图G进行社区划分后的结果，可以看到0和2属于一个社区，1和3属于另一个社区。
partition = community_louvain.best_partition(G)
# y就是节点的社区标签。tensor([0, 1, 0, 1])
y = torch.tensor([partition[i] for i in range(G.number_of_nodes())])

# 选择训练集. train_mask是训练集的标签，用于半监督节点分类任务，每类节点中只有一个节点的标签设置为已知True，其他为未知False。
train_mask = torch.zeros(y.size(0), dtype=torch.bool)
for i in range(int(y.max()) + 1):
    train_mask[(y == i).nonzero(as_tuple=False)[0]] = True

将数据加入到geometric中：

data = Data(x=x, edge_index=edge_index, y=y, train_mask=train_mask)

# 将data.train_mask取反，然后非0元素的索引
remaining = (~data.train_mask).nonzero(as_tuple=False).view(-1)

# torch.randperm会随机得到0到2的一个数字序列
remaining = remaining[torch.randperm(remaining.size(0))]

data.test_mask = torch.zeros(y.size(0), dtype=torch.bool)
data.test_mask.fill_(False)

# tensor([False, False,  True,  True])
data.test_mask[remaining[:]] = True

# ['edge_index', 'train_mask', 'x', 'y', 'test_mask']
print(data.keys)

从上面例子看到可以继续添加test_mask属性，设置某些节点为测试集。

利用上面这个简单的图实现节点分类任务，类别就是上面louvain算法给出的社区类别。训练数据为节点0和1，测试数据为节点2和3。

构建一个图卷积神经网络，包含两个卷积层，第一层输入维度为4，输出维度为16；第二层输入维度为16，输出维度为2；第一层后面接上一个激活函数，并进行dropout操作。

import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = GCNConv(data.num_node_features, 16)
        self.conv2 = GCNConv(16, 2)

    def forward(self):
        x, edge_index = data.x, data.edge_index

        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)

        return F.log_softmax(x, dim=1)

device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
model = Net().to(device)
data = data.to(device)

optimizer = torch.optim.Adam([
    dict(params=model.conv1.parameters(), weight_decay=5e-4),
    dict(params=model.conv2.parameters(), weight_decay=0)
], lr=0.01)  # only perform weight-decay on first convolution.


def train():
    optimizer.zero_grad()  
    out = model()
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    print('train loss', loss.item())
    loss.backward()
    optimizer.step()

def test():
    model.eval()
    logits, accs = model(), []
    for _, mask in data('train_mask', 'test_mask'):
        pred = logits[mask].max(1)[1]
        acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item()
        accs.append(acc)
    return accs

for epoch in range(1, 51):
    train()
    log = 'Epoch: {:03d}, Train: {:.4f}, Test: {:.4f}'
    print(log.format(epoch, *test()))