​基于RotatE模型的知识图谱嵌入技术

pip install torch numpy

import numpy as np

# 定义实体和关系
entities = {'UserA': 0, 'MovieB': 1, 'DirectorC': 2, 'ActorD': 3}
relations = {'likes': 0, 'directed_by': 1, 'acted_in': 2}

# 三元组数据
triples = [
    ('UserA', 'likes', 'MovieB'),
    ('MovieB', 'directed_by', 'DirectorC'),
    ('DirectorC', 'acted_in', 'ActorD')
]

# 将三元组转换为索引
triples_idx = [(entities[h], relations[r], entities[t]) for h, r, t in triples]

import torch
import torch.nn as nn

class RotatE(nn.Module):
    def __init__(self, num_entities, num_relations, embedding_dim):
        super(RotatE, self).__init__()
        self.entity_embeddings = nn.Embedding(num_entities, embedding_dim)
        self.relation_embeddings = nn.Embedding(num_relations, embedding_dim)
        
        # 初始化嵌入
        nn.init.xavier_uniform_(self.entity_embeddings.weight.data)
        nn.init.xavier_uniform_(self.relation_embeddings.weight.data)

    def forward(self, heads, relations, tails):
        # 获取实体和关系嵌入
        head_emb = self.entity_embeddings(heads)
        relation_emb = self.relation_embeddings(relations)
        tail_emb = self.entity_embeddings(tails)

        # 将关系嵌入转换为复数表示
        relation_emb_real = torch.cos(relation_emb)
        relation_emb_imag = torch.sin(relation_emb)

        # 计算旋转后的头实体向量
        rotated_head = head_emb * relation_emb_real + head_emb * relation_emb_imag * 1j
        # 计算尾实体向量的复数形式
        tail_complex = tail_emb[:, 0] + tail_emb[:, 1] * 1j

        # 计算距离
        distance = torch.norm(rotated_head - tail_complex, dim=1)
        return distance

# 定义损失函数
def loss_fn(predictions, labels):
    return torch.mean(torch.clamp(predictions - labels, min=0))

# 创建训练数据
heads = torch.tensor([h for h, r, t in triples_idx])
relations = torch.tensor([r for h, r, t in triples_idx])
tails = torch.tensor([t for h, r, t in triples_idx])

# 定义模型
model = RotatE(num_entities=len(entities), num_relations=len(relations), embedding_dim=50)

# 模拟训练过程
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# 训练轮数
for epoch in range(100):
    optimizer.zero_grad()
    # 获取模型预测
    predictions = model(heads, relations, tails)
    # 使用随机生成的负样本
    labels = torch.zeros_like(predictions)  # 正样本标签
    loss = loss_fn(predictions, labels)
    
    # 反向传播
    loss.backward()
    optimizer.step()
    
    print(f'Epoch {epoch}, Loss: {loss.item()}')