使用DGL做自建异质图上的链接预测

本文是自定义的图。
首先自定义图的结构：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

n_users = 1000  # 定义 user 个数
n_items = 500   # 定义 item 个数
n_follows = 3000    # 定义 follow个数
n_clicks = 5000     # 定义 click 个数
n_dislikes = 500    # 定义 dislikes 个数
n_hetero_features = 10  # 定义 hetero_features 嵌入维数
n_user_classes = 5  # 用户也有 5 种类别
n_max_clicks = 10   # 定义 click 类别数
"""
异质关系如下：
('user', 'follow', 'user')
('user', 'click', 'item')
('user', 'dislike', 'item')
"""
follow_src = np.random.randint(0, n_users, n_follows)   # 边 follow 的源节点
follow_dst = np.random.randint(0, n_users, n_follows)   # 边 follow 的目标节点
click_src = np.random.randint(0, n_users, n_clicks)     # 边 click 的源节点
click_dst = np.random.randint(0, n_items, n_clicks)     # 边 click 的目标节点
dislike_src = np.random.randint(0, n_users, n_dislikes) # 边 dislike 的源节点
dislike_dst = np.random.randint(0, n_items, n_dislikes) # 边 dislike 的目标节点

# 将构造的节点和边都转移到 gpu 上
# 将 NumPy 数组转换为 Torch 张量并移至适当设备
follow_src = torch.tensor(follow_src, device=device)
follow_dst = torch.tensor(follow_dst, device=device)
click_src = torch.tensor(click_src, device=device)
click_dst = torch.tensor(click_dst, device=device)
dislike_src = torch.tensor(dislike_src, device=device)
dislike_dst = torch.tensor(dislike_dst, device=device)

# 构建图 构建人工异质图数据集，包含两种类型节点 和 六种类型边
hetero_graph = dgl.heterograph({
    ('user', 'follow', 'user'): (follow_src, follow_dst),   # follow_src 和 follow_dst 之间有关系 follow
    ('user', 'followed-by', 'user'): (follow_dst, follow_src),
    ('user', 'click', 'item'): (click_src, click_dst),
    ('item', 'clicked-by', 'user'): (click_dst, click_src),
    ('user', 'dislike', 'item'): (dislike_src, dislike_dst),
    ('item', 'disliked-by', 'user'): (dislike_dst, dislike_src)})
    
hetero_graph = hetero_graph.to(device)

为节点初始化特征，将节点user分为5类，将 边 click 分为 10类

# 使用random方式随机为 节点 初始化特征
hetero_graph.nodes['user'].data['feature'] = torch.randn(n_users, n_hetero_features).to(device)
hetero_graph.nodes['item'].data['feature'] = torch.randn(n_items, n_hetero_features).to(device)

# 标签细分，将节点user分为5种，将 边 click分为10种
hetero_graph.nodes['user'].data['label'] = torch.randint(0, n_user_classes, (n_users,)).to(device)  # 为每个user打一个标签(0-4)。返回一个 [0-n_user_classed)内的int的 形状为(n_users,)的一维张量
hetero_graph.edges['click'].data['label'] = torch.randint(1, n_max_clicks, (n_clicks,)).float().to(device) # 为每个click打一个标签

# 随机在 user节点 和 click边 生成训练mask
hetero_graph.nodes['user'].data['train_mask'] = torch.zeros(n_users, dtype=torch.bool).bernoulli(0.6).to(device)    # 使用伯努利分布选60%的为train
hetero_graph.edges['click'].data['train_mask'] = torch.zeros(n_clicks, dtype=torch.bool).bernoulli(0.6).to(device)  # 使用伯努利分布选60%的为train

损失函数计算损失：

def compute_loss(pos_score, neg_score):
    n_edges = pos_score.shape[0]
    return (1 - pos_score.unsqueeze(1) + neg_score.view(n_edges, -1)).clamp(min=0).mean()

对边进行负采样：

# 负采样
# 对要进行链接预测的边类型构造一个负采样图
def construct_negative_graph(graph, k, etype):
    """

    :param graph:
    :param k: 对于每个正样本，生成k个负样本
    :param etype:  边的类型。生成的图中要包含这种 异质边 的类型
    :return:
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    utype, _, vtype = etype  # 提取异质图的源节点类型和目标节点类型
    src, dst = graph.edges(etype=etype)  # 提取特定类型的所有边的源节点和目标节点索引
    neg_src = src.repeat_interleave(k).to(device)  # 为每个正样本重复 k 次源节点的索引(负采样生成的链接的源节点和正样本链接的源节点一样)
    neg_dst = torch.randint(0, graph.num_nodes(vtype), (len(src) * k,)).to(device)  # 负样本链接的目标节点是随机的。
    return dgl.heterograph(
        {etype: (neg_src, neg_dst)},  # 指定异质边的类型(utype, reltype, vtype). 这里的utype, vtype是负采样得到的节点。边是给定的特定的etype。
        num_nodes_dict={ntype: graph.num_nodes(ntype) for ntype in graph.ntypes})  # 为每种节点类型指定节点数量(保证和原图数量一致)

主函数：

def main():
    # 构建好的图是 hetero_graph 图

    hetero_graph = construct_hete_graph()
    graph_feature_init(hetero_graph)

    model = Model(10, 20, 5, hetero_graph.etypes)   # 输入10， hid20, 输出5，边特征为hetero_graph.etypes
    model = model.to(device)

    # 将节点的特征全部取出
    # 用户特征
    user_feats = hetero_graph.nodes['user'].data['feature']
    # 商品特征
    item_feats = hetero_graph.nodes['item'].data['feature']
    # 将所有节点特征存储到一个字典中
    node_features = {'user': user_feats, 'item': item_feats}

    optimizer = torch.optim.Adam(model.parameters())
    epoches = 10
    k = 5

    for epoch in range(epoches):
        # 生成负图
        negative_graph = construct_negative_graph(hetero_graph, k, ('user', 'click', 'item'))   # 为(user click item) 类型的边 生成负图
        negative_graph.to(device)
        pos_score, neg_score = model(hetero_graph, negative_graph, node_features, ('user', 'click', 'item'))
        loss = compute_loss(pos_score, neg_score)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        print(loss.item())

if __name__ == '__main__':
    main()