PyG 2.6 正式发布，支持大模型的图神经网络来了！

最新推荐文章于 2024-10-01 23:51:03 发布

猴子2333

最新推荐文章于 2024-10-01 23:51:03 发布

阅读量44

点赞数

文章标签：神经网络 pytorch 人工智能

原文链接：https://mp.weixin.qq.com/s/ffJPoE3n82cZ0KSO0edZEw

版权

PyG 2.6 正式发布，支持大模型的图神经网络来了！🚀
大家好，今天为大家介绍的是最新发布的 PyG 2.6 版本。在图学习领域，PyTorch Geometric（PyG）一直是重要的工具库，本次 2.6 版本的更新带来了很多值得关注的新功能。
主要亮点 ✨

GNN 与大模型结合

这次更新最重要的功能之一就是引入了大模型与图神经网络（GNN）的结合。通过新模型 GRetriever，PyG 实现了 LLaMA2 与 GNN 的联合训练，大幅提升了知识图谱任务的效果。

在实际应用中，GRetriever 能处理包含大量节点和边的图数据，并利用大模型对复杂的图结构进行更智能的表示学习。这为知识图谱和推理任务带来了新的可能性！为了让大家更好地理解其应用方式，下面给出一个具体的代码示例：

import torch
from torch_geometric.nn import GATv2Conv
from torch_geometric.nn.models import GRetriever
from torch_geometric.nn.nlp import LLM

# 创建图神经网络
class GNN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GATv2Conv(in_channels, hidden_channels, heads=4)
        self.conv2 = GATv2Conv(hidden_channels * 4, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        return self.conv2(x, edge_index)

# 初始化GNN
gnn = GNN(in_channels=1, hidden_channels=16, out_channels=1024)

# 初始化LLM（这里使用PyG提供的LLM封装）
llm = LLM(model_name='meta-llama/Llama-2-7b-chat-hf', num_params=7)

# 初始化GRetriever
model = GRetriever(llm=llm, gnn=gnn)

# 创建图神经网络的数据示例
edge_index = torch.tensor([[0, 1], [1, 2]], dtype=torch.long)
x = torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.float)

# 假设我们有一个问题和描述
question = "What is the capital of France?"
desc = "This is a question about geography."

# 使用GRetriever进行推理
with torch.no_grad():
    output = model.inference(question, x, edge_index, batch=None, edge_attr=None, desc=desc)

print(output)

这个简单的例子展示了如何加载 LLaMA2 作为预训练模型，并将其与 GNN 结合使用。通过 GRetriever 模型实现了大模型与 GNN 的联合训练，具体来说，这里使用了一个小型的图结构（edge_index 表示图中的边，x 表示节点的特征），通过 GRetriever 将图数据传入并输出相应的结果。

可以根据自己的图数据和任务需求，对模型进行调整和优化，进一步提升知识图谱任务中的表现。想要更多的细节和代码，查看官方示例：GRetriever 示例。