PyTorch-BigGraph项目：图嵌入下游任务应用指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00015/article/details/148577787

PyTorch-BigGraph项目：图嵌入下游任务应用指南

PyTorch-BigGraph Generate embeddings from large-scale graph-structured data. 项目地址: https://gitcode.com/gh_mirrors/py/PyTorch-BigGraph

引言

PyTorch-BigGraph（简称PBG）是一个强大的图嵌入工具，能够将图中的实体转换为低维向量表示。这些向量表示（即嵌入）可以用于各种下游任务，如链接预测、实体排名和最近邻搜索等。本文将详细介绍如何解析PBG的输出数据，并展示几种典型下游任务的应用方法。

数据格式解析

PBG采用文件作为输入输出接口，主要支持三种数据格式：

HDF5格式（原生格式）

HDF5是PBG的原生二进制格式，具有高效存储和快速读取的特点。读取示例：

import json
import h5py

# 加载实体名称列表
with open("data/FB15k/entity_names_all_0.json", "rt") as tf:
    names = json.load(tf)
offset = names.index("/m/05hf_5")  # 查找目标实体偏移量

# 加载对应偏移量的嵌入向量
with h5py.File("model/fb15k/embeddings_all_0.v50.h5", "r") as hf:
    embedding = hf["embeddings"][offset, :]

HDF5支持部分读取，可以高效地获取特定数据而不必加载整个文件。

TSV格式（文本格式）

TSV是纯文本格式，便于人类阅读和调试，但解析速度较慢。格式说明：

实体嵌入文件：每行包含实体ID和对应的嵌入向量值（制表符分隔）
关系参数文件：包含关系类型、操作符及其参数信息

读取示例：

import numpy as np

# 加载Wikidata嵌入（跳过注释行，仅加载前78404883行）
embeddings = np.loadtxt(
    "wikidata_translation_v1.tsv",
    dtype=np.float32,
    delimiter="\t",
    skiprows=1,
    max_rows=78404883,
    usecols=range(1, 201),
    comments=None,
)

NPY格式（NumPy二进制格式）

NPY格式是NumPy的二进制格式，解析速度快，适合生产环境：

import numpy as np

# 直接加载预解析的嵌入数据
embeddings = np.load("wikidata_translation_v1_vectors.npy")

# 使用内存映射模式（适用于大数据）
embeddings = np.load("large_embeddings.npy", mmap_mode="r")

下游任务应用

1. 边得分预测

边得分预测是评估给定三元组（源实体、关系类型、目标实体）存在可能性的任务。

示例代码（预测"巴黎是法国首都"的得分）：

import torch
from torchbiggraph.model import ComplexDiagonalDynamicOperator, DotComparator

# 初始化操作符和比较器
operator = ComplexDiagonalDynamicOperator(400, dynamic_rel_count)
operator.load_state_dict(operator_state_dict)
comparator = DotComparator()

# 加载实体和关系类型
src_entity_offset = entity_names.index("/m/0f8l9c")  # 法国
dest_entity_offset = entity_names.index("/m/05qtj")  # 巴黎
rel_type_index = rel_type_names.index("/location/country/capital")

# 计算得分
scores, _, _ = comparator(
    comparator.prepare(src_embedding.view(1, 1, 400)),
    comparator.prepare(
        operator(
            dest_embedding.view(1, 400),
            torch.tensor([rel_type_index]),
        ).view(1, 1, 400),
    ),
    torch.empty(1, 0, 400),  # 左侧负样本（不需要）
    torch.empty(1, 0, 400),  # 右侧负样本（不需要）
)

2. 实体排名

给定源实体和关系类型，对所有目标实体进行可能性排名：

# 计算所有实体的得分
scores, _, _ = comparator(
    comparator.prepare(src_embedding.view(1, 1, 400)).expand(1, entity_count, 400),
    comparator.prepare(
        operator(
            dest_embeddings,
            torch.tensor([rel_type_index]).expand(entity_count),
        ).view(1, entity_count, 400),
    ),
    torch.empty(1, 0, 400),
    torch.empty(1, 0, 400),
)

# 获取排名前5的实体
permutation = scores.flatten().argsort(descending=True)
top5_entities = [entity_names[index] for index in permutation[:5]]

3. 最近邻搜索

使用FAISS库进行高效的最近邻搜索：

import faiss

# 创建FAISS索引
index = faiss.IndexFlatL2(400)  # L2距离度量
index.add(embeddings)  # 添加所有嵌入

# 搜索巴黎的最近邻
target_embedding = embeddings[entity_names.index("/m/05qtj")]
_, neighbors = index.search(target_embedding.reshape((1, 400)), 5)

# 映射回实体名称
top5_entities = [entity_names[index] for index in neighbors[0]]