如何创建并训练异构图GCN神经网络(以论文引用数据集为例)

0. 案例背景：论文的引用网络图：

节点的种类分为：研究领域，论文，作者和研究机构。

相互之间的关系也分为：具有主题，引用，撰写，隶属XX机构

具体如图：

数据来源：Open Graph Benchmark | A collection of benchmark datasets, data-loaders and evaluators for graph machine learning in PyTorch. (stanford.edu)https://ogb.stanford.edu/

1. 创建异构图，用HeteroData创建所有的数据

首先导入HetroData类，并实例化。

在实例化后，分别对其中每种类型的节点传入节点特征：

比如data['paper'].x=...  就是把paper这种类型节点的所有特征全部输入进去。

同理‘author', 'institution','field_of_study' 也是这种情况。

对边的处理略复杂。需要传入两组信息，一个是边的索引(也就是邻接矩阵，以2*编号对 的形式传入)， 一个是边的特征(有多少条边，每条边的特征是什么)

注意一下边的构造形式，以data['paper', 'cites', 'paper'].edge_index为例，在描述paper到paper的边的时候，要按照“起始点”——“边的名称”——“终点”的格式传入。

from torch_geometric.data import HeteroData

data = HeteroData()

data['paper'].x = ... # [num_papers, num_features_paper]
data['author'].x = ... # [num_authors, num_features_author]
data['institution'].x = ... # [num_institutions, num_features_institution]
data['field_of_study'].x = ... # [num_field, num_features_field]

data['paper', 'cites', 'paper'].edge_index = ... # [2, num_edges_cites]
data['author', 'writes', 'paper'].edge_index = ... # [2, num_edges_writes]
data['author', 'affiliated_with', 'institution'].edge_index = ... # [2, num_edges_affiliated]
data['paper', 'has_topic', 'field_of_study'].edge_index = ... # [2, num_edges_topic]

data['paper', 'cites', 'paper'].edge_attr = ... # [num_edges_cites, num_features_cites]
data['author', 'writes', 'paper'].edge_attr = ... # [num_edges_writes, num_features_writes]
data['author', 'affiliated_with', 'institution'].edge_attr = ... # [num_edges_affiliated, num_features_affiliated]
data['paper', 'has_topic', 'field_of_study'].edge_attr = ... # [num_edges_topic, num_features_topic]

data可以被直接打印出来，本例子中的data打印出来后如下图：

HeteroData(
  paper={
    x=[736389, 128],
    y=[736389],
    train_mask=[736389],
    val_mask=[736389],
    test_mask=[736389]
  },
  author={ x=[1134649, 128] },
  institution={ x=[8740, 128] },
  field_of_study={ x=[59965, 128] },
  (author, affiliated_with, institution)={ edge_index=[2, 1043998] },
  (author, writes, paper)={ edge_index=[2, 7145660] },
  (paper, cites, paper)={ edge_index=[2, 5416271] },
  (paper, has_topic, field_of_study)={ edge_index=[2, 7505078] }
)

2. 将创建的所有数据导入模型之中

用HeteroGNN()创建模型，并把data.x_dict, data.edge_index_dict, data.edge_attr_dict三者作为参数输入模型之中。

注意，此时输入的不是data[‘...’].x等，而是把所有种类的节点信息全部按照字典的形式输入，所以x的输入是用.x_dict，同理，边的信息也是用的..._dict的形式。

model = HeteroGNN(...)

output = model(data.x_dict, data.edge_index_dict, data.edge_attr_dict)

3. *一些附带小功能

下面是torch_geometric.data.HeteroData这个类附带的一些功能性函数：

3.1 直接用索引把单独的点信息或者边信息提出出来

paper_node_data = data['paper']
cites_edge_data = data['paper', 'cites', 'paper']

cites_edge_data = data['paper', 'paper']
cites_edge_data = data['cites']

3.2 可以增加新的节点类型或者张量，并把他们移除：

data['paper'].year = ...    # Setting a new paper attribute
del data['field_of_study']  # Deleting 'field_of_study' node type
del data['has_topic']       # Deleting 'has_topic' edge type

3.3 可以获得data的“元数据”，并得到所有的节点及边的信息(原文:We can access the meta-data of the data object, holding information of all present node and edge types:)：

node_types, edge_types = data.metadata()
print(node_types)
['paper', 'author', 'institution']
print(edge_types)
[('paper', 'cites', 'paper'),
('author', 'writes', 'paper'),
('author', 'affiliated_with', 'institution')]

3.4 data可以在CPU和GPU之间自由转换：

data = data.to('cuda:0')
data = data.cpu()

3.5 其他辅助函数：

data.has_isolated_nodes()
data.has_self_loops()
data.is_undirected()

3.6 将异构图数据转化为同构图：

homogeneous_data = data.to_homogeneous()
print(homogeneous_data)
Data(x=[1879778, 128], edge_index=[2, 13605929], edge_type=[13605929])

4. *异构图转换功能

绝大部分在同构图中用于前处理数据的transformations方法都可以在异构图中使用：

import torch_geometric.transforms as T

data = T.ToUndirected()(data)
data = T.AddSelfLoops()(data)
data = T.NormalizeFeatures()(data)

ToUndirected()ToUndirected()方法把一个有向图转化为一个无向图。方法是为每一条边创建一个反向边。

AddSelfLoops()方法是给每个点增加指向自己的边

NormalizeFeatures()方法用于做归一化，和同构图一样

5.创建异构图神经网络，有三种方法。最常用的是Automatically Converting GNN Models

用如下两条命令之一，即可自动转换为异构图模型：torch_geometric.nn.to_hetero()和torch_geometric.nn.to_hetero_with_bases()

以下代码是例子：

import torch_geometric.transforms as T
from torch_geometric.datasets import OGB_MAG
from torch_geometric.nn import SAGEConv, to_hetero


dataset = OGB_MAG(root='./data', preprocess='metapath2vec', transform=T.ToUndirected())
data = dataset[0]

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels)
        self.conv2 = SAGEConv((-1, -1), out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x


model = GNN(hidden_channels=64, out_channels=dataset.num_classes)
model = to_hetero(model, data.metadata(), aggr='sum')

很明显，上面的例子里创建的GNN模型是一个同构图模型。利用to_hetero()方法，把实例化的GNN模型作为参数传入，同时把数据集也传入，to_hetero()方法可以直接返回一个异构图模型。

其底层逻辑如下图所示：

to_hetero()和to_hetero_with_bases()都非常灵活，诸如残差链接、条约知识等技术都可以直接从同构图自动转化为异构图。下面是一个是一个实施了可学习的残差链接的同构图转化为异构图的例子：

from torch_geometric.nn import GATConv, Linear, to_hetero

class GAT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GATConv((-1, -1), hidden_channels, add_self_loops=False)
        self.lin1 = Linear(-1, hidden_channels)
        self.conv2 = GATConv((-1, -1), out_channels, add_self_loops=False)
        self.lin2 = Linear(-1, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index) + self.lin1(x)
        x = x.relu()
        x = self.conv2(x, edge_index) + self.lin2(x)
        return x


model = GAT(hidden_channels=64, out_channels=dataset.num_classes)
model = to_hetero(model, data.metadata(), aggr='sum')

但是要注意一下，这里我们把add_self_loops=False设置为false，原因是自循环的概念在二分图中没有明确的定义。这段解释document中的原文如下：

6. 在全部创建完毕后，就可以正常训练了：

def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x_dict, data.edge_index_dict)
    mask = data['paper'].train_mask
    loss = F.cross_entropy(out['paper'][mask], data['paper'].y[mask])
    loss.backward()
    optimizer.step()
    return float(loss)