帮我写一段代码，实现简单的图神经网络进行链接预测，要求使用构造的数据集

好的，我们可以使用Python中的DGL库来实现这个简单的图神经网络进行链接预测。以下是一个简单的代码示例：

import dgl
import torch
import torch.nn as nn

# 构造数据集
g = dgl.graph(([0, 1, 2, 3], [1, 2, 3, 4]))  # 定义一个包含5个节点和4条边的图
g.ndata['feat'] = torch.randn(5, 10)  # 每个节点有一个10维的特征向量
g.edata['feat'] = torch.randn(4, 5)  # 每条边有一个5维的特征向量
g.apply_edges(lambda edges: {'feat': torch.cat([edges.src['feat'], edges.dst['feat']], dim=1)})  # 将每条边的特征向量定义为源节点和目标节点的特征向量的拼接

# 定义一个简单的图神经网络模型
class GNN(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim):
        super(GNN, self).__init__()
        self.conv1 = dgl.nn.GraphConv(in_dim, hidden_dim)
        self.conv2 = dgl.nn.GraphConv(hidden_dim, out_dim)

    def forward(self, g):
        h = g.ndata['feat']
        h = self.conv1(g, h)
        h = torch.relu(h)
        h = self.conv2(g, h)
        return h

# 训练模型
model = GNN(20, 16, 1)  # 输入特征向量是源节点和目标节点特征向量拼接，所以输入维度为20（10+10）
criterion = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for epoch in range(100):
    logits = model(g)
    labels = g.edata['label']
    loss = criterion(logits, labels)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print('Epoch %d | Loss: %.4f' % (epoch, loss.item()))

# 预测链接
model.eval()
with torch.no_grad():
    logits = model(g)
    scores = torch.sigmoid(logits)
    _, indices = torch.topk(scores.view(-1), 2)  # 输出得分最高的两条边
    src, dst = g.find_edges(indices)
    print('Top 2 edges are:')
    for i in range(len(src)):
        print('%d -> %d' % (src[i].item(), dst[i].item())))

在这个代码示例中，我们首先构造了一个包含5个节点和4条边的图，并使用随机向量作为节点和边的特征向量。我们将每条边的特征向量定义为源节点和目标节点的特征向量的拼接。接下来，我们定义了一个简单的图神经网络模型，模型的输入特征向量是源节点和目标节点特征向量拼接，所以输入维度为20（10+10）。我们使用BCEWithLogitsLoss作为损失函数，Adam作为优化器进行模型训练。最后，我们使用训练好的模型进行链接预测，输出得分最高的两条边。