wehub-resource-sync
95 lines
4.5 KiB
ReStructuredText
95 lines
4.5 KiB
ReStructuredText
Build Model
|
|
===========
|
|
|
|
**GraphTransformer** is a graph neural network that uses multi-head self-attention (sparse or dense) to encode the graph structure and node features. It is a generalization of the `Transformer <https://arxiv.org/abs/1706.03762>`_ architecture to arbitrary graphs.
|
|
|
|
In this tutorial, we will show how to build a graph transformer model with DGL using the `Graphormer <https://arxiv.org/abs/2106.05234>`_ model as an example.
|
|
|
|
Graphormer is a Transformer model designed for graph-structured data, which encodes the structural information of a graph into the standard Transformer. Specifically, Graphormer utilizes degree encoding to measure the importance of nodes, spatial and path Encoding to measure the relation between node pairs. The degree encoding and the node features serve as input to Graphormer, while the spatial and path encoding act as bias terms in the self-attention module.
|
|
|
|
Degree Encoding
|
|
-------------------
|
|
The degree encoder is a learnable embedding layer that encodes the degree of each node into a vector. It takes as input the batched input and output degrees of graph nodes, and outputs the degree embeddings of the nodes.
|
|
|
|
.. code:: python
|
|
|
|
degree_encoder = dgl.nn.DegreeEncoder(
|
|
max_degree=8, # the maximum degree to cut off
|
|
embedding_dim=512 # the dimension of the degree embedding
|
|
)
|
|
|
|
Path Encoding
|
|
-------------
|
|
The path encoder encodes the edge features on the shortest path between two nodes to get attention bias for the self-attention module. It takes as input the batched edge features in shape and outputs the attention bias based on path encoding.
|
|
|
|
.. code:: python
|
|
|
|
path_encoder = PathEncoder(
|
|
max_len=5, # the maximum length of the shortest path
|
|
feat_dim=512, # the dimension of the edge feature
|
|
num_heads=8, # the number of attention heads
|
|
)
|
|
|
|
Spatial Encoding
|
|
----------------
|
|
The spatial encoder encodes the shortest distance between two nodes to get attention bias for the self-attention module. It takes as input the shortest distance between two nodes and outputs the attention bias based on spatial encoding.
|
|
|
|
.. code:: python
|
|
|
|
spatial_encoder = SpatialEncoder(
|
|
max_dist=5, # the maximum distance between two nodes
|
|
num_heads=8, # the number of attention heads
|
|
)
|
|
|
|
|
|
Graphormer Layer
|
|
----------------
|
|
The Graphormer layer is like a Transformer encoder layer with the Multi-head Attention part replaced with :class:`~dgl.nn.BiasedMHA`. It takes in not only the input node features, but also the attention bias computed computed above, and outputs the updated node features.
|
|
|
|
We can stack multiple Graphormer layers as a list just like implementing a Transformer encoder in PyTorch.
|
|
|
|
.. code:: python
|
|
|
|
layers = th.nn.ModuleList([
|
|
GraphormerLayer(
|
|
feat_size=512, # the dimension of the input node features
|
|
hidden_size=1024, # the dimension of the hidden layer
|
|
num_heads=8, # the number of attention heads
|
|
dropout=0.1, # the dropout rate
|
|
activation=th.nn.ReLU(), # the activation function
|
|
norm_first=False, # whether to put the normalization before attention and feedforward
|
|
)
|
|
for _ in range(6)
|
|
])
|
|
|
|
Model Forward
|
|
-------------
|
|
Grouping the modules above defines the primary components of the Graphormer model. We then can define the forward process as follows:
|
|
|
|
.. code:: python
|
|
|
|
node_feat, in_degree, out_degree, attn_mask, path_data, dist = \
|
|
next(iter(dataloader)) # we will use the first batch as an example
|
|
num_graphs, max_num_nodes, _ = node_feat.shape
|
|
deg_emb = degree_encoder(th.stack((in_degree, out_degree)))
|
|
|
|
# node feature + degree encoding as input
|
|
node_feat = node_feat + deg_emb
|
|
|
|
# spatial encoding and path encoding serve as attention bias
|
|
path_encoding = path_encoder(dist, path_data)
|
|
spatial_encoding = spatial_encoder(dist)
|
|
attn_bias[:, 1:, 1:, :] = path_encoding + spatial_encoding
|
|
|
|
# graphormer layers
|
|
for layer in layers:
|
|
x = layer(
|
|
x,
|
|
attn_mask=attn_mask,
|
|
attn_bias=attn_bias,
|
|
)
|
|
|
|
For simplicity, we omit some details in the forward process. For the complete implementation, please refer to the `Graphormer example <https://github.com/dmlc/dgl/tree/master/examples/core/Graphormer>`_.
|
|
|
|
You can also explore other `utility modules <https://docs.dgl.ai/api/python/nn-pytorch.html#utility-modules-for-graph-transformer>`_ to customize your own graph transformer model. In the next section, we will show how to prepare the data for training.
|