WN18数据集

class dgl.data.WN18Dataset(reverse=True, raw_dir=None, force_reload=False, verbose=True, transform=None)

基类：KnowledgeGraphDataset

WN18 链接预测数据集。

WN18 数据集是在 Translating Embeddings for Modeling Multi-relational Data 中介绍的。它包含了从 WordNet 中抓取的全部 18 种关系，对应大约 41,000 个同义词集。默认情况下，创建数据集时，为每条边创建一条具有反向关系类型的反向边。

WN18 数据集统计信息

• 节点数：40943

• 关系类型数量：18

• 反向关系类型数量：18

• 数据集划分

  • 训练集：141442

  • 验证集：5000

  • 测试集：5000

参数:

• reverse (bool) – 是否添加反向边。默认值为 True。

• raw_dir (str) – 下载原始文件或包含输入数据目录的目录。默认值：~/.dgl/

• force_reload (bool) – 是否重新加载数据集。默认值：False

• verbose (bool) – 是否打印进度信息。默认值：True。

• transform (callable, optional) – 一个转换函数，接受一个 DGLGraph 对象并返回一个转换后的版本。该 DGLGraph 对象将在每次访问前进行转换。

num_nodes

节点数量

类型:

int

num_rels

关系类型数量

类型:

int

示例

>>> dataset = WN18Dataset()
>>> g = dataset.graph
>>> e_type = g.edata['e_type']
>>>
>>> # get data split
>>> train_mask = g.edata['train_mask']
>>> val_mask = g.edata['val_mask']
>>>
>>> train_set = th.arange(g.num_edges())[train_mask]
>>> val_set = th.arange(g.num_edges())[val_mask]
>>>
>>> # build train_g
>>> train_edges = train_set
>>> train_g = g.edge_subgraph(train_edges,
                              relabel_nodes=False)
>>> train_g.edata['e_type'] = e_type[train_edges];
>>>
>>> # build val_g
>>> val_edges = th.cat([train_edges, val_edges])
>>> val_g = g.edge_subgraph(val_edges,
                            relabel_nodes=False)
>>> val_g.edata['e_type'] = e_type[val_edges];
>>>
>>> # Train, Validation and Test
>>>

__getitem__(idx)

获取图对象

参数:

idx (int) – 项索引，WN18Dataset 只有一个图对象

返回值:

图中包含

• edata['e_type']: 边关系类型

• edata['train_edge_mask']: 正向训练边掩码

• edata['val_edge_mask']: 正向验证边掩码

• edata['test_edge_mask']: 正向测试边掩码

• edata['train_mask']: 训练边集掩码（包含反向训练边）

• edata['val_mask']: 验证边集掩码（包含反向验证边）

• edata['test_mask']: 测试边集掩码（包含反向测试边）

• ndata['ntype']: 节点类型。在此数据集中均为 0

返回类型:

dgl.DGLGraph

__len__()

数据集中的图数量。