第二十一课.DeepGraphLibrary(二)

构建GNN模块

官方SAGEConv和HeteroGraphConv用法

SAGEConv

在自主构建GNN模块前,先了解两个官方实现的GNN模块用法,首先学习 SAGEConv 的用法,GraphSAGE 来自论文 " Inductive Representation Learning on Large Graphs",消息传递过程为: h N ( i ) ( l + 1 ) = a g g r e g a t e ( { h j ( l ) , ∀ j ∈ N ( i ) } ) h_{N(i)}^{(l+1)}=aggregate(\left\{h_{j}^{(l)},\forall j\in N(i)\right\}) hN(i)(l+1)=aggregate({hj(l),jN(i)}) h i ( l + 1 ) = σ ( W ⋅ c o n c a t ( h i ( l ) , h N ( i ) ( l + 1 ) ) ) h_{i}^{(l+1)}=\sigma(W\cdot concat(h_{i}^{(l)},h_{N(i)}^{(l+1)})) hi(l+1)=σ(Wconcat(hi(l),hN(i)(l+1))) h i ( l + 1 ) = n o r m ( h i ( l + 1 ) ) = h i ( l + 1 ) ∣ ∣ h i ( l + 1 ) ∣ ∣ 2 h_{i}^{(l+1)}=norm(h_{i}^{(l+1)})=\frac{h_{i}^{(l+1)}}{||h_{i}^{(l+1)}||_{2}} hi(l+1)=norm(hi(l+1))=hi(l+1)2hi(l+1)如果在edge上存在权重 e e e,则第一步聚合表示为: h N ( i ) ( l + 1 ) = a g g r e g a t e ( { e j i ⋅ h j ( l ) , ∀ j ∈ N ( i ) } ) h_{N(i)}^{(l+1)}=aggregate(\left\{e_{ji}\cdot h_{j}^{(l)},\forall j\in N(i)\right\}) hN(i)(l+1)=aggregate({ejihj(l),jN(i)})在官方的api中,SAGEConv为:

dgl.nn.pytorch.conv.SAGEConv(in_feats, 
							 out_feats, 
							 aggregator_type, 
							 feat_drop=0.0, 
							 bias=True, 
							 norm=None, 
							 activation=None)

参数说明如下:

  • in_feats:输入特征的size,比如 h i ( l ) h_{i}^{(l)} hi(l)的维数,SAGEConv可用于同构图和单向二分图;对于单向二分图,in_feats 指定源节点和目标节点上的输入特征大小,如果给出标量,则源节点和目标节点特征大小将采用相同的值;
    如果聚合类型aggregator_type为gcn,则要求源节点和目标节点的特征大小相同;
  • out_feats:输出特征的size,比如 h i ( l + 1 ) h_{i}^{(l+1)} hi(l+1)的维数;
  • feat_drop:特征丢失率;
  • aggregator_type:聚合类型 (mean, gcn, pool, lstm)
  • norm:用于节点更新后的标准化;
  • activation:用于节点更新后的激活函数;

SAGEConv的forward函数:

forward(graph, feat, edge_weight=None)

参数说明:

  • graph:要计算的图;
  • feat:如果输入是一个tensor,则表示形状为 ( N , D i n ) (N,D_{in}) (N,Din)的输入特征, N N N是节点数量, D i n D_{in} Din是输入特征的维数;
    如果输入是一对tensor,这对tensor的形状应为: ( N i n , D i n s r c ) (N_{in},D_{in}^{src}) (Nin,Dinsrc) ( N o u t , D i n d s t ) (N_{out},D_{in}^{dst}) (Nout,Dindst)
  • edge_weight:在edge上的权重,如果给定,图卷积将对消息加权;

SAGEConv在前向计算后返回一个张量,形状为 ( N , D o u t ) (N,D_{out}) (N,Dout),其中, D o u t D_{out} Dout是输出特征的维数;

同构图的实例如下:

import dgl
import numpy as np
import torch as th
from dgl.nn.pytorch.conv import SAGEConv

# 同构图 Homogeneous graph
# 6个节点,6条边
g=dgl.graph(([0,1,2,3,2,5],[1,2,3,4,0,3]))
# 每个节点添加自环后:6个节点,12条边
g=dgl.add_self_loop(g)

feat=th.ones(6,10)
conv=SAGEConv(10,2,'pool')

res=conv(g,feat)

# 在同构图中,源节点也是目标节点
print(res.size()) # torch.Size([6, 2])

单向二分图的实例如下:

# 单向二分图 Unidirectional bipartite graph (一种异构图)
u=[0,1,0,0,1]
v=[0,1,2,3,2]
graph_data={('u','link','v'):(u,v)}

g=dgl.heterograph(graph_data)
print(g)
"""
Graph(num_nodes={'u': 2, 'v': 4},
      num_edges={('u', 'link', 'v'): 5},
      metagraph=[('u', 'v', 'link')])
"""

u_feat=th.randn(2,5)
v_feat=th.randn(4,10)

conv=SAGEConv((5,10),2,'mean')
res=conv(g,(u_feat,v_feat))

# 结果为二分图中目标节点的特征
print(res.size()) # torch.Size([4, 2])

关于SAGE和GCN的讨论

SAGE相比GCN,真正意义上做到了逐个局部区域的消息聚合,GCN的计算可以回顾图神经网络专栏:第二课.图卷积神经网络,通过对比二者的算法流程就能发现SAGE真正实现了在大图数据上应用GNN


HeteroGraphConv

下面学习HeteroGraphConv的用法,这是一个用于计算异构图卷积的通用模块。

异构图卷积将不同的GNN模块应用到各个关系上,从源节点读取特征并将更新的特征写入目标节点。如果多个关系具有相同的目标节点类型,则目标节点的结果将按照指定方法聚合。

HeteroGraphConv定义为:

dgl.nn.pytorch.HeteroGraphConv(mods, aggregate='sum')

参数说明:

  • mods:一个字典dict{str : nn.Module},表示对各个规范边类型(也称为关系,或不同的子图)应用不同的GNN模块;
  • aggregate:定义不同关系下的同类节点特征的聚合方法,有“sum”、“max”、“min”、“mean”、“stack”

关于forward函数:

forward(g, inputs)

参数说明:

  • g:要计算的图;
  • inputs:一个字典dict[str : Tensor]或者一对字典,一个字典代表输入节点(源节点)特征,如果是一对字典(一般为tuple(src_feat,dst_feat))代表源节点特征和目标节点特征;
    注意在异构图中,有些类型下的目标节点是其他类型的源节点,对这类节点,我们统一把更新前的这类节点也叫做源节点,更新后即变成目标节点;

前向计算返回的是目标节点的特征(字典dict[str : Tensor]);

下面是HeteroGraphConv的伪代码例子,首先建立异构图:

import dgl
g = dgl.heterograph({
    ('user', 'follows', 'user') : edges1,
    ('user', 'plays', 'game') : edges2,
    ('store', 'sells', 'game')  : edges3})

创建一个 HeteroGraphConv,将不同的卷积模块应用于不同的关系(注意,'follows''plays' 的模块不共享权重):

import dgl.nn.pytorch as dglnn
conv = dglnn.HeteroGraphConv({
    'follows' : dglnn.GraphConv(...),
    'plays' : dglnn.GraphConv(...),
    'sells' : dglnn.SAGEConv(...)},
    aggregate='sum')

'user'调用forward

import torch as th
h1 = {'user' : th.randn((g.number_of_nodes('user'), 5))}
h2 = conv(g, h1)
print(h2.keys()) # dict_keys(['user', 'game'])

'user''store'调用forward,因为 'plays''sells' 关系都会更新 'game' 特征,所以'game'的结果由指定的聚合方法aggregateHeteroGraphConv(mods, aggregate='sum')中的参数)进行聚合:

f1 = {'user' : th.randn((g.number_of_nodes('user'), 5)), 
	  'store' : th.randn((g.number_of_nodes('user'), 3))}
f2 = conv(g, f1)
print(f2.keys()) # dict_keys(['user', 'game'])

'store'调用forward,由于包含源节点'store'的关系只有('store', 'sells', 'game'),所以只有目标节点'game'被更新,所以forward返回的目标节点特征只有'game'的:

g1 = {'store' : ...}
g2 = conv(g, g1)
print(g2.keys()) # dict_keys(['game'])

也可以对输入的一对数据调用forward,每个子模块也会被这一对输入调用:

x_src = {'user' : ..., 'store' : ...}
x_dst = {'user' : ..., 'game' : ...}
y_dst = conv(g, (x_src, x_dst))
print(y_dst.keys()) # dict_keys(['user', 'game'])

利用DGL构建SAGE

首先定义构造函数:

import torch.nn as nn
from dgl.utils import expand_as_pair

class SAGEConv(nn.Module):
    def __init__(self,
                 in_feats,
                 out_feats,
                 aggregator_type,
                 bias=True,
                 norm=None,
                 activation=None):
        super().__init__()
        self._in_src_feats,self._in_dst_feats=expand_as_pair(in_feats)
        self._out_feats=out_feats
        self._aggre_type=aggregator_type
        self.norm=norm
        self.activation=activation

        # 聚合类型:mean、max_pool、lstm、gcn
        if aggregator_type not in ['mean', 'max_pool', 'lstm', 'gcn']:
            raise KeyError('Aggregator type {} not supported.'.format(aggregator_type))
        if aggregator_type == 'max_pool':
            self.fc_pool = nn.Linear(self._in_src_feats, self._in_src_feats)
        if aggregator_type == 'lstm':
            self.lstm = nn.LSTM(self._in_src_feats, self._in_src_feats, batch_first=True)
        if aggregator_type in ['mean', 'max_pool', 'lstm']:
            self.fc_self = nn.Linear(self._in_dst_feats, out_feats, bias=bias)
        self.fc_neigh = nn.Linear(self._in_src_feats, out_feats, bias=bias)

        # 参数初始化
        self.reset_parameters()

这些参数可以回顾上面关于官方SAGEConv的参数说明,额外需要注意函数expand_as_pair(input,g=None)

  • 源节点特征 _in_src_feats 和目标节点特征 _in_dst_feats 需要根据图类型被指定。 用于指定图类型并将 in_feats 扩展为 _in_src_feats_in_dst_feats 的函数是 expand_as_pair
  • 对于同构图,源节点和目标节点相同,它们都是图中的所有节点;
  • 当输入特征是一个元组时,图将会被视为二分图。元组中的第一个元素为源节点特征,第二个元素为目标节点特征,expand_as_pair可以解码输入从而返回源节点特征维数和目标节点特征维数;
  • expand_as_pair也可以直接对输入特征解码;

除了数据维度,图神经网络的一个典型选项是聚合类型(self._aggre_type)。对于特定目标节点,聚合类型决定了如何聚合不同边上的信息。 常用的聚合类型包括 mean、 sum、 max 和 min。一些模块可能会使用更加复杂的聚合函数,比如 lstm。

权重初始化是类中的一个实例方法:

    def reset_parameters(self):
        """重新初始化可学习的参数"""
        gain = nn.init.calculate_gain('relu')
        if self._aggre_type == 'max_pool':
            nn.init.xavier_uniform_(self.fc_pool.weight, gain=gain)
        if self._aggre_type == 'lstm':
            self.lstm.reset_parameters()
        if self._aggre_type != 'gcn':
            nn.init.xavier_uniform_(self.fc_self.weight, gain=gain)
        nn.init.xavier_uniform_(self.fc_neigh.weight, gain=gain)

定义前向计算函数为2个操作:

  • 消息传递和聚合;
  • 聚合后,更新特征作为输出;

回顾GraphSAGE,消息传递过程为: h N ( i ) ( l + 1 ) = a g g r e g a t e ( { h j ( l ) , ∀ j ∈ N ( i ) } ) h_{N(i)}^{(l+1)}=aggregate(\left\{h_{j}^{(l)},\forall j\in N(i)\right\}) hN(i)(l+1)=aggregate({hj(l),jN(i)}) h i ( l + 1 ) = σ ( W ⋅ c o n c a t ( h i ( l ) , h N ( i ) ( l + 1 ) ) ) h_{i}^{(l+1)}=\sigma(W\cdot concat(h_{i}^{(l)},h_{N(i)}^{(l+1)})) hi(l+1)=σ(Wconcat(hi(l),hN(i)(l+1))) h i ( l + 1 ) = n o r m ( h i ( l + 1 ) ) = h i ( l + 1 ) ∣ ∣ h i ( l + 1 ) ∣ ∣ 2 h_{i}^{(l+1)}=norm(h_{i}^{(l+1)})=\frac{h_{i}^{(l+1)}}{||h_{i}^{(l+1)}||_{2}} hi(l+1)=norm(hi(l+1))=hi(l+1)2hi(l+1)因此,实现如下:

    def forward(self,graph,feat):
    	"""
    	关于local_scope():
    	在这个范围中,可以引用、修改图中的特征值 ,但是只是临时的,出了这个范围,一切恢复原样;
    	这样做的目的是方便计算。毕竟我们在图上作消息传递和聚合,有时仅是为了计算值,并不想改变原始图
    	"""
        with graph.local_scope():
            # 指定图类型,然后根据图类型扩展输入特征
            feat_src,feat_dst=expand_as_pair(feat,graph)

            import dgl.function as fn
            import torch.nn.functional as F
            from dgl.utils import check_eq_shape

            h_self = feat_dst
			"""消息传递和聚合"""
            if self._aggre_type == 'mean':
                # graph.srcdata利用了描述符(一个特殊的装饰器@property)
                # 将类中的读写方法转为属性控制
                graph.srcdata['h'] = feat_src
                # 将源节点的特征拷贝到目标节点的 mailbox['m'], 对mailbox内的消息聚合
                graph.update_all(fn.copy_u('h', 'm'), fn.mean('m', 'neigh'))
                # 取出目标节点的更新特征
                h_neigh = graph.dstdata['neigh']
                # graph.srcdata['neigh']是索引不到的
            elif self._aggre_type == 'gcn':
                check_eq_shape(feat)
                graph.srcdata['h'] = feat_src
                graph.dstdata['h'] = feat_dst
                graph.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'neigh'))
                # 除以入度
                degs = graph.in_degrees().to(feat_dst)
                h_neigh = (graph.dstdata['neigh'] + graph.dstdata['h']) / (degs.unsqueeze(-1) + 1)
            elif self._aggre_type == 'max_pool':
                graph.srcdata['h'] = F.relu(self.fc_pool(feat_src))
                graph.update_all(fn.copy_u('h', 'm'), fn.max('m', 'neigh'))
                h_neigh = graph.dstdata['neigh']
            else:
                raise KeyError('Aggregator type {} not recognized.'.format(self._aggre_type))

            # GraphSAGE中gcn聚合不需要fc_self
            if self._aggre_type == 'gcn':
                rst = self.fc_neigh(h_neigh)
            else:
            	"""
				修改了一下SAGE的计算:
				W x concat(h_self,h_neigh) 变成 W_1 x h_self + W_2 x h_neigh 
				这不影响SAGE的原理, 它依然是可以用于 Large Graph 的算法
				"""
                rst = self.fc_self(h_self) + self.fc_neigh(h_neigh)

			"""聚合后,更新特征作为输出"""
            # 激活函数
            if self.activation is not None:
                rst = self.activation(rst)
            # 归一化
            if self.norm is not None:
                rst = self.norm(rst)
            return rst

然后输入一个同构图和单向二分图:

import dgl
import torch as th

# 同构图 Homogeneous graph
# 6个节点,6条边
g=dgl.graph(([0,1,2,3,2,5],[1,2,3,4,0,3]))
# 每个节点添加自环后:6个节点,12条边
g=dgl.add_self_loop(g)

feat=th.ones(6,10)
conv=SAGEConv(10,2,'max_pool')

res=conv(g,feat)

print(res.size()) # torch.Size([6, 2])

# 单向二分图 Unidirectional bipartite graph (一种异构图)
u=[0,1,0,0,1]
v=[0,1,2,3,2]
graph_data={('u','link','v'):(u,v)}

g=dgl.heterograph(graph_data)
"""
Graph(num_nodes={'u': 2, 'v': 4},
      num_edges={('u', 'link', 'v'): 5},
      metagraph=[('u', 'v', 'link')])
"""

u_feat=th.randn(2,5)
v_feat=th.randn(4,10)

conv=SAGEConv((5,10),2,'mean')
res=conv(g,(u_feat,v_feat))

print(res.size()) # torch.Size([4, 2])

可以发现,输出的张量形状和官方SAGE的计算一致;


思考

关于mailbox,mailbox本质是目标节点的成员,但我们在第二十课.DeepGraphLibrary(一)的单独调用Edge-wise更新边特征时有这样一段内容:

import dgl.function as fn

graph.apply_edges(fn.u_add_v('el', 'er', 'e'))

apply_edges() 的参数是一个消息函数。并且在默认情况下, apply_edges()将更新所有的边。

上面内容中,消息'e'仿佛保存在edge上,其实消息依然是在目标节点上,只是在聚合前,mailbox['e']每行中的张量还包括了来自源节点的排序,这使得apply_edges可以根据 “(源节点,边,目标节点)” 的模式将这些消息赋值到每条边上


DGL图数据集

DGL在 dgl.data 里实现了很多常用的图数据集。它们遵循了由 dgl.data.DGLDataset 类定义的标准的数据处理管道。 DGL推荐用户将图数据处理为 dgl.data.DGLDataset 的子类。该类为导入、处理和保存图数据提供了简单而干净的解决方案。

DGLDataset

DGLDataset 是处理、导入和保存 dgl.data 中定义的图数据集的基类。 它实现了用于处理图数据的基本模版。下面的流程图展示了这个模版的工作方式:
fig1
为了处理位于远程服务器或本地磁盘上的图数据集,下面的例子中定义了一个类,称为 MyDataset, 它继承自 dgl.data.DGLDataset:

from dgl.data import DGLDataset

class MyDataset(DGLDataset):
    """ 
    用于在DGL中自定义图数据集的模板:
    Parameters
    ----------
    url : str
        下载原始数据集的url。
    raw_dir : str
        指定下载数据的存储目录或已下载数据的存储目录。默认: ~/.dgl/
    save_dir : str
        处理完成的数据集的保存目录。默认:raw_dir指定的值
    force_reload : bool
        是否重新导入数据集。默认:False
    verbose : bool
        是否打印进度信息。
    """
    def __init__(self,
                 url=None,
                 raw_dir=None,
                 save_dir=None,
                 force_reload=False,
                 verbose=False):
        super(MyDataset, self).__init__(name='dataset_name',
                                        url=url,
                                        raw_dir=raw_dir,
                                        save_dir=save_dir,
                                        force_reload=force_reload,
                                        verbose=verbose)

    def download(self):
        # 将原始数据下载到本地磁盘
        pass

    def process(self):
        # 将原始数据处理为图、标签和数据集划分的掩码
        pass

    def __getitem__(self, idx):
        # 通过idx得到与之对应的一个样本
        pass

    def __len__(self):
        # 数据样本的数量
        pass

    def save(self):
        # 将处理后的数据保存至 `self.save_path`
        pass

    def load(self):
        # 从 `self.save_path` 导入处理后的数据
        pass

    def has_cache(self):
        # 检查在 `self.save_path` 中是否存有处理后的数据
        pass

DGLDataset 类有方法 process()__getitem__(idx)__len__()。子类必须实现这些函数。同时DGL也建议实现保存和导入函数, 因为对于处理后的大型数据集,这么做可以节省大量的时间。

请注意, DGLDataset 的目的是提供一种标准且方便的方式来导入图数据。 用户可以存储有关数据集的图、特征、标签、掩码,以及诸如类别数、标签数等基本信息。 诸如采样、划分或特征归一化等操作建议在 DGLDataset 子类之外完成。

下载原始数据

如果用户的数据集已经在本地磁盘中,请确保它被存放在目录 raw_dir 中。 如果用户想在任何地方运行代码而又不想自己下载数据并将其移动到正确的目录中,则可以通过实现函数 download() 来自动完成。

如果数据集是一个zip文件,可以直接继承 dgl.data.DGLBuiltinDataset 类。后者支持解压zip文件。 否则用户需要自己实现 download(),下面是download()的举例定义:

import os
from dgl.data.utils import download

def download(self):
    # 存储文件的路径
    file_path = os.path.join(self.raw_dir, self.name + '.mat')
    # 下载文件
    download(self.url, path=file_path)

上面的代码将一个.mat文件下载到目录 self.raw_dir。如果文件是.gz.tar.tar.gz.tgz文件,需要使用 extract_archive() 函数进行解压。以下代码展示了如何在 BitcoinOTCDataset 类中下载一个.gz文件:

from dgl.data.utils import download, check_sha1

def download(self):
    # 存储文件的路径,请确保使用与原始文件名相同的后缀
    gz_file_path = os.path.join(self.raw_dir, self.name + '.csv.gz')
    # 下载文件
    download(self.url, path=gz_file_path)
    # 检查 SHA-1:
    # 一个可选项是用户可以按照上面的示例检查下载后文件的SHA-1字符串,以防他人在远程服务器上更改了文件
    if not check_sha1(gz_file_path, self._sha1_str):
        raise UserWarning('File {} is downloaded but the content hash does not match.'
                          'The repo may be outdated or download may be incomplete. '
                          'Otherwise you can create an issue for it.'.format(self.name + '.csv.gz'))
    # 将文件解压缩到目录self.raw_dir下的self.name目录中
    self._extract_gz(gz_file_path, self.raw_path)

上面的代码会将文件解压缩到 self.raw_dir 下的目录 self.name 中。 如果该类继承自 dgl.data.DGLBuiltinDataset 来处理zip文件, 则它也会将文件解压缩到目录 self.name 中。

处理数据

用户可以在 process() 函数中实现数据处理。该函数假定原始数据已经位于 self.raw_dir 目录中。

图上的机器学习任务通常有三种类型:整图分类、节点分类和链接预测。下面将展示如何处理与这些任务相关的数据集。

处理整图分类数据集

整图分类数据集与用小批次训练的典型机器学习任务中的大多数数据集类似。 因此,需要将原始数据处理为 dgl.DGLGraph 对象的列表和标签张量的列表。 此外,如果原始数据已被拆分为多个文件,则可以添加参数 split 以导入数据的特定部分。

下面以 QM7bDataset 为例,QM7b是整图分类的数据集,用于预测分子的属性,一共包含7211个分子,预测分类到14种属性,每个分子中的edge上具有特征(size=1,表示原子之间的库仑力),在官方的QM7bDataset类中,有下面几个需要了解的方法:

from dgl.data import DGLDataset

class QM7bDataset(DGLDataset):
    _url = 'http://deepchem.io.s3-website-us-west-1.amazonaws.com/' \
           'datasets/qm7b.mat'
    _sha1_str = '4102c744bb9d6fd7b40ac67a300e49cd87e28392'

    def __init__(self, raw_dir=None, force_reload=False, verbose=False):
        super(QM7bDataset, self).__init__(name='qm7b',
                                          url=self._url,
                                          raw_dir=raw_dir,
                                          force_reload=force_reload,
                                          verbose=verbose)

    def process(self):
        mat_path = self.raw_path + '.mat'
        # 将数据处理为图列表和标签列表
        # _load_graph是此类的一个实例方法
        self.graphs, self.label = self._load_graph(mat_path)

    def __getitem__(self, idx):
        """ 通过idx获取对应的图和标签
        Parameters
        ----------
        idx : int
            Item index

        Returns
        -------
        (dgl.DGLGraph, Tensor)
        """
        return self.graphs[idx], self.label[idx]

    def __len__(self):
        """数据集中图的数量"""
        return len(self.graphs)
        
	@property
	def num_labels(self):
	    """每个图的标签数,即预测类别数"""
	    return 14

函数 process() 将原始数据处理为图列表和标签列表。用户必须实现 __getitem__(idx)__len__() 以进行迭代。 DGL建议让 __getitem__(idx) 返回如上面代码所示的元组 (图,标签);

实例化QM7bDataset,可以将其用于训练:

import dgl
import torch

from dgl.dataloading import GraphDataLoader

# 数据导入
dataset = QM7bDataset()
num_labels = dataset.num_labels

# 创建 dataloaders
dataloader = GraphDataLoader(dataset, batch_size=1, shuffle=True)

# 训练
for epoch in range(100):
    for g, labels in dataloader:
        # 用户自己的训练代码
        pass

处理节点分类数据集

与整图分类不同,节点分类通常在单个图上进行。因此数据集的划分是在图的节点集上进行。 DGL建议使用节点掩码来指定数据集的划分。 下面的内容将以内置数据集 CitationGraphDataset 展开。

此外,DGL推荐重新排列图的节点/边,使得相邻节点/边的ID位于邻近区间内。这个过程可以提高节点/边的邻居的局部性,为后续在图上进行的计算与分析的性能改善提供可能。 DGL提供了名为 dgl.reorder_graph() 的API用于此优化。下面分析 CitationGraphDataset 中的部分代码:

# DGLBuiltinDataset支持解压zip
from dgl.data import DGLBuiltinDataset
from dgl.data.utils import _get_dgl_url

class CitationGraphDataset(DGLBuiltinDataset):
    _urls = {
        'cora_v2' : 'dataset/cora_v2.zip',
        'citeseer' : 'dataset/citeseer.zip',
        'pubmed' : 'dataset/pubmed.zip',
    }

    def __init__(self, name, raw_dir=None, force_reload=False, verbose=True):
        assert name.lower() in ['cora', 'citeseer', 'pubmed']
        if name.lower() == 'cora':
            name = 'cora_v2'
        url = _get_dgl_url(self._urls[name])
        super(CitationGraphDataset, self).__init__(name,
                                                   url=url,
                                                   raw_dir=raw_dir,
                                                   force_reload=force_reload,
                                                   verbose=verbose)

    def process(self):
        # 跳过一些处理的代码
        # === 跳过数据处理 ===

        # 构建图
        g = dgl.graph(graph)

        # 划分掩码
        g.ndata['train_mask'] = train_mask
        g.ndata['val_mask'] = val_mask
        g.ndata['test_mask'] = test_mask

        # 节点的标签
        g.ndata['label'] = torch.tensor(labels)

        # 节点的特征
        g.ndata['feat'] = torch.tensor(_preprocess_features(features),
                                       dtype=F.data_type_dict['float32'])
        self._num_labels = onehot_labels.shape[1]
        self._labels = labels
        # 重排图以获得更优的局部性
        self._g = dgl.reorder_graph(g)

    def __getitem__(self, idx):
        assert idx == 0, "这个数据集里只有一个图"
        return self._g

    def __len__(self):
        return 1

为简便起见,这里省略了 process() 中的一些代码,以突出展示用于处理节点分类数据集的关键部分:划分掩码。 节点特征和节点的labels被存储在 g.ndata 中。请注意,这里 __getitem__(idx)__len__() 的实现也发生了变化, 这是因为节点分类任务通常只用一个图。

下面使用 dgl.data.CiteseerGraphDataset 来演示如何用于节点分类:

# 导入数据
dataset = CiteseerGraphDataset(raw_dir='...')
graph = dataset[0]

# 获取划分的掩码
train_mask = graph.ndata['train_mask']
val_mask = graph.ndata['val_mask']
test_mask = graph.ndata['test_mask']

# 获取节点特征
feats = graph.ndata['feat']

# 获取标签
labels = graph.ndata['label']

处理edge预测数据集

链接预测数据集的处理与节点分类相似,数据集中通常只有一个图。以内置的数据集 KnowledgeGraphDataset 为例,同时省略了详细的数据处理代码以突出展示处理链接预测数据集的关键部分:

# 创建链接预测数据集示例
class KnowledgeGraphDataset(DGLBuiltinDataset):
    def __init__(self, name, reverse=True, raw_dir=None, force_reload=False, verbose=True):
        self._name = name
        self.reverse = reverse
        url = _get_dgl_url('dataset/') + '{}.tgz'.format(name)
        super(KnowledgeGraphDataset, self).__init__(name,
                                                    url=url,
                                                    raw_dir=raw_dir,
                                                    force_reload=force_reload,
                                                    verbose=verbose)

    def process(self):
        # 跳过一些处理的代码
        # === 跳过数据处理 ===

        # 划分掩码
        g.edata['train_mask'] = train_mask
        g.edata['val_mask'] = val_mask
        g.edata['test_mask'] = test_mask

        # 边类型
        g.edata['etype'] = etype

        # 节点类型
        g.ndata['ntype'] = ntype
        self._g = g

    def __getitem__(self, idx):
        assert idx == 0, "这个数据集只有一个图"
        return self._g

    def __len__(self):
        return 1

如代码所示,图的 edata 存储了划分掩码。下面使用 KnowledgeGraphDataset的子类 dgl.data.FB15k237Dataset 来做演示如何使用用于链路预测的数据集:

from dgl.data import FB15k237Dataset
import torch

# 导入数据
dataset = FB15k237Dataset()
graph = dataset[0]

# 获取训练集掩码
train_mask = graph.edata['train_mask']
print(train_mask,len(train_mask))
# tensor([ True,  True,  True,  ..., False, False, False]) 620232

"""
squeeze()去除冗余维度
torch.nonzero 返回非零元素的索引
"""
train_idx = torch.nonzero(train_mask, as_tuple=False).squeeze()
print(train_idx,len(train_idx))
# tensor(0, 1, 2, ..., 620214, 620215, 620216]) 544230

# 获取训练集中的边类型
rel = graph.edata['etype'][train_idx]
print(rel) # tensor([183, 183, 183,  ..., 329, 329, 329])

保存和加载数据

DGL建议用户实现保存和加载数据的函数,将处理后的数据缓存在本地磁盘中。 这样在多数情况下可以帮用户节省大量的数据处理时间。DGL提供了4个函数让任务变得简单:

  • dgl.save_graphs()dgl.load_graphs():保存DGLGraph对象和标签到本地磁盘和从本地磁盘读取它们;
  • dgl.data.utils.save_info()dgl.data.utils.load_info():将数据集的有用信息(python dict对象)保存到本地磁盘和从本地磁盘读取它们;

下面的示例显示了如何保存和读取图和数据集信息的列表:

import os
from dgl import save_graphs, load_graphs
from dgl.data.utils import makedirs, save_info, load_info

def save(self):
    # 保存图和标签
    graph_path = os.path.join(self.save_path, self.mode + '_dgl_graph.bin')
    save_graphs(graph_path, self.graphs, {'labels': self.labels})
    # 在Python字典里保存其他信息
    info_path = os.path.join(self.save_path, self.mode + '_info.pkl')
    save_info(info_path, {'num_classes': self.num_classes})

def load(self):
    # 从目录 `self.save_path` 里读取处理过的数据
    graph_path = os.path.join(self.save_path, self.mode + '_dgl_graph.bin')
    self.graphs, label_dict = load_graphs(graph_path)
    self.labels = label_dict['labels']
    info_path = os.path.join(self.save_path, self.mode + '_info.pkl')
    self.num_classes = load_info(info_path)['num_classes']

def has_cache(self):
    # 检查在 `self.save_path` 里是否有处理过的数据文件
    graph_path = os.path.join(self.save_path, self.mode + '_dgl_graph.bin')
    info_path = os.path.join(self.save_path, self.mode + '_info.pkl')
    return os.path.exists(graph_path) and os.path.exists(info_path)

使用ogb包导入OGB数据集

Open Graph Benchmark (OGB) 是一个图深度学习的基准数据集。 官方的 ogb 包提供了用于下载和处理OGB数据集到 dgl.data.DGLGraph 对象的API。下面介绍它们的基本用法;

首先安装ogb:

pip install ogb

下面是整图预测数据集的加载示例:

import dgl
import torch
from ogb.graphproppred import DglGraphPropPredDataset
from dgl.dataloading import GraphDataLoader

def _collate_fn(batch):
    # 小批次是一个元组(graph, label)列表 [(graph),(label)]
    graphs = [e[0] for e in batch]
    """
    dgl.batch(graphs, ndata='__ALL__', edata='__ALL__')
    将一组 DGLGraph 分批组合到一个图中,以实现更高效的图计算
    """
    g = dgl.batch(graphs)
    labels = [e[1] for e in batch]
    """
    outputs = torch.stack(inputs, dim=0) → Tensor
    把多个2维张量堆叠成一个3维张量,多个3维张量堆叠成一个4维张量,以此类推
    inputs : 待堆叠的张量序列
    dim : 新的维度
    cat和stack的区别在于cat会增加现有维度的内容,即拼接,stack会直接增加一个维度,即堆叠
    """
    labels = torch.stack(labels,dim=0)
    # print(labels.size()) # (batch_size,1)
    return g, labels

# 载入数据集
# ogbg-molhiv 和 ogbg-molpcba 数据集是两个不同大小的分子特性预测数据集;
# ogbg-molhiv(小)和 ogbg-molpcba(中)
# 下载并处理数据保存到 './dataset/ogbg_molhiv/'
dataset = DglGraphPropPredDataset(name='ogbg-molhiv',root='./dataset/')
# split_idx用于划分数据集
split_idx = dataset.get_idx_split()

# dataloader
"""
dgl.dataloading.pytorch.GraphDataLoader(dataset, collate_fn=None, **kwargs)
PyTorch数据加载器,用于对一组图进行批量迭代,生成所述小批量的批量图和相应的标签张量

collate_fn:collate function,如果没有给出,将使用默认的collate函数,
           collate函数指定如何重新整理batch中的样本
"""
train_loader = GraphDataLoader(dataset[split_idx["train"]],
                               batch_size=32,
                               shuffle=True,
                               collate_fn=_collate_fn)
valid_loader = GraphDataLoader(dataset[split_idx["valid"]],
                               batch_size=32,
                               shuffle=False,
                               collate_fn=_collate_fn)
test_loader = GraphDataLoader(dataset[split_idx["test"]],
                              batch_size=32,
                              shuffle=False,
                              collate_fn=_collate_fn)

g,labels=next(iter(train_loader))
print(g)
print(labels)
"""
Graph(num_nodes=772, num_edges=1664,
      ndata_schemes={'feat': Scheme(shape=(9,), dtype=torch.int64)}
      edata_schemes={'feat': Scheme(shape=(3,), dtype=torch.int64)})

tensor([[1],
        [0],
        [0],
        ...
        [0]]) size=(32,1)
"""

加载 Node Property Prediction(节点属性预测) 数据集类似,但要注意的是这种数据集只有一个图对象:

from ogb.nodeproppred import DglNodePropPredDataset

"""
ogbn-proteins 数据集是一个无向加权和类型化(根据物种)图,节点表示蛋白质,
边表示蛋白质之间不同类型的具有生物学意义的关联,例如物理相互作用、共表达或同源性,
所有的边都带有 8 维特征,蛋白质来自8个物种;

预测任务:该任务是在多标签二元分类问题中预测蛋白质功能的存在,其中总共有 112 种标签要预测
"""
dataset = DglNodePropPredDataset(name='ogbn-proteins',root='./dataset/')
split_idx = dataset.get_idx_split()

# 在Node Property Prediction数据集里只有一个图
g, labels = dataset[0]
# 获取划分的标签
train_label = dataset.labels[split_idx['train']]
valid_label = dataset.labels[split_idx['valid']]
test_label = dataset.labels[split_idx['test']]

print(train_label)
print(train_label[0])
print(train_label.size())
"""
tensor([[1, 1, 1,  ..., 0, 0, 0],
        [1, 1, 0,  ..., 0, 0, 0],
        [0, 0, 0,  ..., 0, 0, 0],
        ...,
        [0, 0, 0,  ..., 0, 0, 0],
        [1, 0, 0,  ..., 0, 0, 0],
        [1, 0, 1,  ..., 0, 0, 0]])
tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 0, 1, 1, 1, 1, 0, 1, 1, 1,
        1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 1, 1,
        1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0,
        0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
torch.Size([86619, 112])
"""

每个 Link Property Prediction (边属性预测)数据集也只包括一个图:

from ogb.linkproppred import DglLinkPropPredDataset

"""
ogbl-ppa 数据集是一个无向未加权的图,节点代表来自 58 个不同物种的蛋白质,
边表示蛋白质之间具有生物学意义的关联,例如物理相互作用,共表达,同源性或基因组邻域, 
OGB提供了一个由训练edge构建的图(不包含验证和测试edge),
每个节点包含一个 58 维的 one-hot 特征向量,指示相应蛋白质来自的物种;

预测任务:任务是在给定训练边的情况下预测新的关联边
"""
dataset = DglLinkPropPredDataset(name='ogbl-ppa')
split_edge = dataset.get_edge_split()

graph = dataset[0]

print(split_edge['train'].keys())
print(split_edge['valid'].keys())
print(split_edge['test'].keys())
"""
dict_keys(['edge'])
dict_keys(['edge', 'edge_neg'])
dict_keys(['edge', 'edge_neg'])

edge_neg是随机采样的3000000个negative edge,用于计算positive edge的比例,
从而评价模型的edge预测效果
"""

print(split_edge['test']['edge'].size(),split_edge['test']['edge_neg'].size())
# torch.Size([3031780, 2]) torch.Size([3000000, 2])
# dim=1维度size=2表示(源节点,目标节点), 即边的表示方式
  • 3
    点赞
  • 10
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值