图神经网络开源框架DGL

图神经网络开源框架DGL学习

• 还有另外一个开源框架PyTorch Geometric, 在这里不做详细介绍，有兴趣的同学自己去了解一下，个人推荐DGL，简单易用，扩展性好。

• 个人公众号
  

1. 使用教程

1.1 pip安装

pip install dgl           # For CPU Build
pip install dgl-cu90      # For CUDA 9.0 Build
pip install dgl-cu92      # For CUDA 9.2 Build
pip install dgl-cu100     # For CUDA 10.0 Build
pip install dgl-cu101     # For CUDA 10.1 Build

1.2 CUDA安装

conda install -c dglteam dgl              # For CPU Build
conda install -c dglteam dgl-cuda9.0      # For CUDA 9.0 Build
conda install -c dglteam dgl-cuda10.0     # For CUDA 10.0 Build
conda install -c dglteam dgl-cuda10.1     # For CUDA 10.1 Build

1.3 官方文档

• https://docs.dgl.ai/

1.4 代码仓库

• https://github.com/dmlc/dgl

2. 视频教程

• https://www.bilibili.com/video/BV1JE411K7aT?from=search&seid=15647099632944979039

3. 代码示例

import dgl
import scipy.io
import urllib.request

data_url = 'https://data.dgl.ai/dataset/ACM.mat'
data_file_path = '/tmp/ACM.mat'

urllib.request.urlretrieve(data_url, data_file_path)
data = scipy.io.loadmat(data_file_path)
print(list(data.keys()))

###############################################################################
# The dataset stores node information by their types: ``P`` for paper, ``A``
# for author, ``C`` for conference, ``L`` for subject code, and so on. The relationships
# are stored as SciPy sparse matrix under key ``XvsY``, where ``X`` and ``Y``
# could be any of the node type code.
#
# The following code prints out some statistics about the paper-author relationships.

print(type(data['PvsA']))
print('#Papers:', data['PvsA'].shape[0])
print('#Authors:', data['PvsA'].shape[1])
print('#Links:', data['PvsA'].nnz)

###############################################################################
# Converting this SciPy matrix to a heterograph in DGL is straightforward.

pa_g = dgl.heterograph({('paper', 'written-by', 'author') : data['PvsA']})
# equivalent (shorter) API for creating heterograph with two node types:
pa_g = dgl.bipartite(data['PvsA'], 'paper', 'written-by', 'author')

###############################################################################
# You can easily print out the type names and other structural information.

print('Node types:', pa_g.ntypes)
print('Edge types:', pa_g.etypes)
print('Canonical edge types:', pa_g.canonical_etypes)

# Nodes and edges are assigned integer IDs starting from zero and each type has its own counting.
# To distinguish the nodes and edges of different types, specify the type name as the argument.
print(pa_g.number_of_nodes('paper'))
# Canonical edge type name can be shortened to only one edge type name if it is
# uniquely distinguishable.
print(pa_g.number_of_edges(('paper', 'written-by', 'author')))
print(pa_g.number_of_edges('written-by'))
print(pa_g.successors(1, etype='written-by'))  # get the authors that write paper #1

# Type name argument could be omitted whenever the behavior is unambiguous.
print(pa_g.number_of_edges())  # Only one edge type, the edge type argument could be omitted

###############################################################################
# A homogeneous graph is just a special case of a heterograph with only one type
# of node and edge. In this case, all the APIs are exactly the same as in
# :class:`DGLGraph`.

# Paper-citing-paper graph is a homogeneous graph
pp_g = dgl.heterograph({('paper', 'citing', 'paper') : data['PvsP']})
# equivalent (shorter) API for creating homogeneous graph
pp_g = dgl.graph(data['PvsP'], 'paper', 'cite')

# All the ntype and etype arguments could be omitted because the behavior is unambiguous.
print(pp_g.number_of_nodes())
print(pp_g.number_of_edges())
print(pp_g.successors(3))

###############################################################################
# Create a subset of the ACM graph using the paper-author, paper-paper, 
# and paper-subject relationships.  Meanwhile, also add the reverse
# relationship to prepare for the later sections.

G = dgl.heterograph({
        ('paper', 'written-by', 'author') : data['PvsA'],
        ('author', 'writing', 'paper') : data['PvsA'].transpose(),
        ('paper', 'citing', 'paper') : data['PvsP'],
        ('paper', 'cited', 'paper') : data['PvsP'].transpose(),
        ('paper', 'is-about', 'subject') : data['PvsL'],
        ('subject', 'has', 'paper') : data['PvsL'].transpose(),
    })

print(G)

###############################################################################
# **Metagraph** (or network schema) is a useful summary of a heterograph.
# Serving as a template for a heterograph, it tells how many types of objects
# exist in the network and where the possible links exist.
#
# DGL provides easy access to the metagraph, which could be visualized using
# external tools.

# Draw the metagraph using graphviz.
import pygraphviz as pgv
def plot_graph(nxg):
    ag = pgv.AGraph(strict=False, directed=True)
    for u, v, k in nxg.edges(keys=True):
        ag.add_edge(u, v, label=k)
    ag.layout('dot')
    ag.draw('graph.png')

plot_graph(G.metagraph)

###############################################################################
# Learning tasks associated with heterographs
# -------------------------------------------
# Some of the typical learning tasks that involve heterographs include:
#
# * *Node classification and regression* to predict the class of each node or
#   estimate a value associated with it.
#
# * *Link prediction* to predict if there is an edge of a certain
#   type between a pair of nodes, or predict which other nodes a particular
#   node is connected with (and optionally the edge types of such connections).
#
# * *Graph classification/regression* to assign an entire
#   heterograph into one of the target classes or to estimate a numerical
#   value associated with it.
#
# In this tutorial, we designed a simple example for the first task.
#
# A semi-supervised node classification example
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Our goal is to predict the publishing conference of a paper using the ACM
# academic graph we just created. To further simplify the task, we only focus
# on papers published in three conferences: *KDD*, *ICML*, and *VLDB*. All
# the other papers are not labeled, making it a semi-supervised setting.
#
# The following code extracts those papers from the raw dataset and prepares 
# the training, validation, testing split.

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

pvc = data['PvsC'].tocsr()
# find all papers published in KDD, ICML, VLDB
c_selected = [0, 11, 13]  # KDD, ICML, VLDB
p_selected = pvc[:, c_selected].tocoo()
# generate labels
labels = pvc.indices
labels[labels == 11] = 1
labels[labels == 13] = 2
labels = torch.tensor(labels).long()

# generate train/val/test split
pid = p_selected.row
shuffle = np.random.permutation(
	)
train_idx = torch.tensor(shuffle[0:800]).long()
val_idx = torch.tensor(shuffle[800:900]).long()
test_idx = torch.tensor(shuffle[900:]).long()

###############################################################################
# Relational-GCN on heterograph
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# We use `Relational-GCN `_ to learn the
# representation of nodes in the graph. Its message-passing equation is as
# follows:
#
# .. math::
#
#    h_i^{(l+1)} = \sigma\left(\sum_{r\in \mathcal{R}}
#    \sum_{j\in\mathcal{N}_r(i)}W_r^{(l)}h_j^{(l)}\right)
#
# Breaking down the equation, you see that there are two parts in the
# computation.
#
# (i) Message computation and aggregation within each relation :math:`r`
#
# (ii) Reduction that merges the results from multiple relationships
#
# Following this intuition, perform message passing on a heterograph in
# two steps.
#
# (i) Per-edge-type message passing
#
# (ii) Type wise reduction

import dgl.function as fn

class HeteroRGCNLayer(nn.Module):
    def __init__(self, in_size, out_size, etypes):
        super(HeteroRGCNLayer, self).__init__()
        # W_r for each relation
        self.weight = nn.ModuleDict({
                name : nn.Linear(in_size, out_size) for name in etypes
            })

    def forward(self, G, feat_dict):
        # The input is a dictionary of node features for each type
        funcs = {}
        for srctype, etype, dsttype in G.canonical_etypes:
            # Compute W_r * h
            Wh = self.weight[etype](feat_dict[srctype])
            # Save it in graph for message passing
            G.nodes[srctype].data['Wh_%s' % etype] = Wh
            # Specify per-relation message passing functions: (message_func, reduce_func).
            # Note that the results are saved to the same destination feature 'h', which
            # hints the type wise reducer for aggregation.
            funcs[etype] = (fn.copy_u('Wh_%s' % etype, 'm'), fn.mean('m', 'h'))
        # Trigger message passing of multiple types.
        # The first argument is the message passing functions for each relation.
        # The second one is the type wise reducer, could be "sum", "max",
        # "min", "mean", "stack"
        G.multi_update_all(funcs, 'sum')
        # return the updated node feature dictionary
        return {ntype : G.nodes[ntype].data['h'] for ntype in G.ntypes}

###############################################################################
# Create a simple GNN by stacking two ``HeteroRGCNLayer``. Since the
# nodes do not have input features, make their embeddings trainable.

class HeteroRGCN(nn.Module):
    def __init__(self, G, in_size, hidden_size, out_size):
        super(HeteroRGCN, self).__init__()
        # Use trainable node embeddings as featureless inputs.
        embed_dict = {ntype : nn.Parameter(torch.Tensor(G.number_of_nodes(ntype), in_size))
                      for ntype in G.ntypes}
        for key, embed in embed_dict.items():
            nn.init.xavier_uniform_(embed)
        self.embed = nn.ParameterDict(embed_dict)
        # create layers
        self.layer1 = HeteroRGCNLayer(in_size, hidden_size, G.etypes)
        self.layer2 = HeteroRGCNLayer(hidden_size, out_size, G.etypes)

    def forward(self, G):
        h_dict = self.layer1(G, self.embed)
        h_dict = {k : F.leaky_relu(h) for k, h in h_dict.items()}
        h_dict = self.layer2(G, h_dict)
        # get paper logits
        return h_dict['paper']

###############################################################################
# Train and evaluate
# ~~~~~~~~~~~~~~~~~~
# Train and evaluate this network.

# Create the model. The output has three logits for three classes.
model = HeteroRGCN(G, 10, 10, 3)

opt = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

best_val_acc = 0
best_test_acc = 0

for epoch in range(100):
    logits = model(G)
    # The loss is computed only for labeled nodes.
    loss = F.cross_entropy(logits[train_idx], labels[train_idx])

    pred = logits.argmax(1)
    train_acc = (pred[train_idx] == labels[train_idx]).float().mean()
    val_acc = (pred[val_idx] == labels[val_idx]).float().mean()
    test_acc = (pred[test_idx] == labels[test_idx]).float().mean()

    if best_val_acc < val_acc:
        best_val_acc = val_acc
        best_test_acc = test_acc

    opt.zero_grad()
    loss.backward()
    opt.step()

    if epoch % 5 == 0:
        print('Loss %.4f, Train Acc %.4f, Val Acc %.4f (Best %.4f), Test Acc %.4f (Best %.4f)' % (
            loss.item(),
            train_acc.item(),
            val_acc.item(),
            best_val_acc.item(),
            test_acc.item(),
            best_test_acc.item(),
        ))