FB15k237Dataset

class dgl.data.FB15k237Dataset(reverse=True, raw_dir=None, force_reload=False, verbose=True, transform=None)[源码]

基类：KnowledgeGraphDataset

FB15k237 链接预测数据集。

FB15k-237 是 FB15k 的一个子集，其中移除了逆关系。创建数据集时，默认会为每条边创建一个带有反向关系类型的反向边。

FB15k237 数据集统计信息

  • 节点数：14541

  • 关系类型数：237

  • 反向关系类型数：237

  • 数据集划分

    • 训练集：272115

    • 验证集：17535

    • 测试集：20466

参数：

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

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

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

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

  • transform (可调用对象, 可选) – 一个转换函数，接受一个 DGLGraph 对象并返回一个转换后的版本。每次访问 DGLGraph 对象之前都会进行转换。

num_nodes

节点数

类型：

int

num_rels

关系类型数

类型：

int

示例

>>> dataset = FB15k237Dataset()
>>> g = dataset.graph
>>> e_type = g.edata['e_type']
>>>
>>> # get data split
>>> train_mask = g.edata['train_mask']
>>> val_mask = g.edata['val_mask']
>>> test_mask = g.edata['test_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) – 项目索引，FB15k237Dataset 只有一个图对象

返回：

图包含以下信息：

  • 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__()[源码]

数据集中的图数量。