FB15kDataset¶

class dgl.data.FB15kDataset(reverse=True, raw_dir=None, force_reload=False, verbose=True, transform=None)[source]¶

基类: KnowledgeGraphDataset

FB15k 链接预测数据集。

FB15K 数据集在 Translating Embeddings for Modeling Multi-relational Data 这篇论文中被引入。它是 Freebase 的一个子集，包含约 14,951 个实体和 1,345 种不同的关系。默认情况下，创建数据集时，会为每条边创建一个具有反向关系类型的反向边。

FB15k 数据集统计

节点数: 14,951
关系类型数: 1,345
反向关系类型数: 1,345
标签划分
- 训练集: 483142
- 验证集: 50000
- 测试集: 59071

参数:

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 = FB15kDataset()
>>> 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)[source]¶

获取图对象

参数:

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

返回值:

图包含:

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__()[source]¶: 数据集中图的数量。