『1024 | 码项目』在低资源环境下训练知识图谱嵌入的详细指南

import torch
import torch.nn as nn
import numpy as np

# 模拟一个简单的知识图谱数据集，包括头实体、关系和尾实体
triples = [
    ("Alice", "knows", "Bob"),
    ("Bob", "likes", "Pizza"),
    ("Alice", "likes", "IceCream"),
    ("Pizza", "is_a", "Food"),
    ("IceCream", "is_a", "Food")
]

# 为了进行嵌入训练，需要对实体和关系进行编码
entity2id = {"Alice": 0, "Bob": 1, "Pizza": 2, "IceCream": 3, "Food": 4}
relation2id = {"knows": 0, "likes": 1, "is_a": 2}

class TransE(nn.Module):
    def __init__(self, entity_count, relation_count, embedding_dim):
        super(TransE, self).__init__()
        # 为实体和关系定义嵌入矩阵
        self.entity_embedding = nn.Embedding(entity_count, embedding_dim)
        self.relation_embedding = nn.Embedding(relation_count, embedding_dim)

    def forward(self, head, relation, tail):
        # 获取头实体、关系、尾实体的嵌入向量
        h = self.entity_embedding(head)
        r = self.relation_embedding(relation)
        t = self.entity_embedding(tail)
        
        # 使用L2距离度量h + r 与 t 的距离
        score = torch.norm(h + r - t, p=2, dim=1)
        return score

# 定义超参数
embedding_dim = 50  # 嵌入向量维度
entity_count = len(entity2id)
relation_count = len(relation2id)

# 实例化模型
model = TransE(entity_count, relation_count, embedding_dim)

# 定义损失函数
def loss_function(pos_score, neg_score, margin=1.0):
    return torch.max(pos_score - neg_score + margin, torch.tensor(0.0)).mean()

# 定义优化器
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 定义负采样函数
def negative_sample(triples, entity_count):
    neg_triples = []
    for h, r, t in triples:
        # 随机替换头实体或尾实体，生成负样本
        if np.random.rand() > 0.5:
            h = np.random.randint(0, entity_count)
        else:
            t = np.random.randint(0, entity_count)
        neg_triples.append((h, r, t))
    return neg_triples

# 训练模型
num_epochs = 100
for epoch in range(num_epochs):
    total_loss = 0
    for triple in triples:
        head, relation, tail = entity2id[triple[0]], relation2id[triple[1]], entity2id[triple[2]]
        neg_triple = negative_sample([triple], entity_count)
        neg_head, neg_relation, neg_tail = neg_triple[0]

        # 转换为tensor
        head = torch.tensor([head], dtype=torch.long)
        relation = torch.tensor([relation], dtype=torch.long)
        tail = torch.tensor([tail], dtype=torch.long)

        neg_head = torch.tensor([neg_head], dtype=torch.long)
        neg_relation = torch.tensor([neg_relation], dtype=torch.long)
        neg_tail = torch.tensor([neg_tail], dtype=torch.long)

        # 前向传播正负样本
        pos_score = model(head, relation, tail)
        neg_score = model(neg_head, neg_relation, neg_tail)

        # 计算损失
        loss = loss_function(pos_score, neg_score)

        # 反向传播与优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    if epoch % 10 == 0:
        print(f"Epoch {epoch}, Loss: {total_loss:.4f}")

import random

# 示例三元组 (药物A, 治疗, 疾病B)
triples = [("药物A", "治疗", "疾病B")]

# 生成反向三元组
augmented_triples = []
for h, r, t in triples:
    augmented_triples.append((t, "被治疗", h))

# 添加噪声数据
noise_triple = ("药物C", "治疗", "疾病D")
if random.random() > 0.5:
    augmented_triples.append(noise_triple)

print(augmented_triples)

from openke.module.model import TransE

# 加载预训练模型
pretrained_model = TransE(
    ent_tot=1000, rel_tot=500, dim=100, p_norm=1, norm_flag=True
)
pretrained_model.load_checkpoint('./checkpoint/pretrained_transe.ckpt')

# 在新数据上进行微调
new_train_dataloader = TrainDataLoader(
    in_path="./new_data/", nbatches=50, threads=8, sampling_mode="normal"
)

# 使用预训练模型进行微调
trainer = Trainer(model=pretrained_model, data_loader=new_train_dataloader, train_times=500, alpha=0.5, use_gpu=False)
trainer.run()

from openke.module.model import TransE
from openke.module.loss import MarginLoss
from openke.config import Trainer

# 定义教师模型
teacher_model = TransE(ent_tot=1000, rel_tot=500, dim=200, p_norm=1, norm_flag=True)

# 定义学生模型（轻量化）
student_model = TransE(ent_tot=1000, rel_tot=500, dim=50, p_norm=1, norm_flag=True)

# 使用知识蒸馏进行训练
class DistillationTrainer(Trainer):
    def distillation_loss(self, student_preds, teacher_preds):
        # 定义知识蒸馏损失
        return ((teacher_preds - student_preds) ** 2).mean()

# 初始化训练器
trainer = DistillationTrainer(model=student_model, data_loader=new_train_dataloader, train_times=500, alpha=0.5, use_gpu=False)
trainer.run()

import networkx as nx
import random

# 创建一个小型的图来模拟知识图谱
G = nx.Graph()

# 添加实体（节点）和关系（边）
G.add_edges_from([
    ('EntityA', 'EntityB', {'relation': 'friend'}),
    ('EntityB', 'EntityC', {'relation': 'colleague'}),
    ('EntityC', 'EntityD', {'relation': 'neighbor'}),
    ('EntityD', 'EntityE', {'relation': 'friend'}),
    ('EntityE', 'EntityF', {'relation': 'colleague'}),
])

# 随机游走采样函数
def random_walk_sample(G, start_node, walk_length):
    walk = [start_node]
    for _ in range(walk_length - 1):
        neighbors = list(G.neighbors(walk[-1]))
        if len(neighbors) > 0:
            walk.append(random.choice(neighbors))
        else:
            break
    return walk

# 从'EntityA'开始随机游走，路径长度为4
sampled_walk = random_walk_sample(G, 'EntityA', walk_length=4)
print("采样的路径:", sampled_walk)

# 邻域采样函数，采样每个节点的邻居节点
def neighborhood_sampling(G, node, num_samples):
    neighbors = list(G.neighbors(node))
    if len(neighbors) > num_samples:
        sampled_neighbors = random.sample(neighbors, num_samples)
    else:
        sampled_neighbors = neighbors
    return sampled_neighbors

# 对'EntityA'进行邻域采样，采样数量为2
sampled_neighbors = neighborhood_sampling(G, 'EntityA', num_samples=2)
print("邻域采样的节点:", sampled_neighbors)

import torch
import torch.nn as nn

# 定义一个简单的线性层模型
class SimpleModel(nn.Module):
    def __init__(self, input_size, output_size):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(input_size, output_size)
        self.fc2 = nn.Linear(output_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        return self.fc2(x)

# 创建模型
model = SimpleModel(input_size=100, output_size=50)

# 模型剪枝操作，移除小于一定阈值的权重
def prune_model(model, pruning_percentage):
    for name, param in model.named_parameters():
        if 'weight' in name:
            threshold = torch.quantile(param.abs(), pruning_percentage)
            mask = (param.abs() > threshold).float()
            param.data.mul_(mask)

# 剪枝10%的权重
prune_model(model, pruning_percentage=0.1)

import torch
import torch.nn as nn
import torch.optim as optim

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self, input_size, output_size):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, x):
        return self.fc(x)

# 生成模型
model = SimpleModel(input_size=100, output_size=50)

# 定义MAML更新步骤
def maml_update(model, loss_fn, x_train, y_train, x_val, y_val, inner_lr=0.01, outer_lr=0.001):
    # 克隆模型用于内环更新
    fast_model = SimpleModel(input_size=100, output_size=50)
    fast_model.load_state_dict(model.state_dict())
    
    # 内环更新（task-specific update）
    optimizer = optim.SGD(fast_model.parameters(), lr=inner_lr)
    loss = loss_fn(fast_model(x_train), y_train)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 计算外环更新
    val_loss = loss_fn(fast_model(x_val), y_val)
    model_params = list(model.parameters())
    fast_params = list(fast_model.parameters())
    
    # 使用验证集的损失来更新初始模型参数
    for i in range(len(model_params)):
        model_params[i].data -= outer_lr * (model_params[i].data - fast_params[i].data)

# 定义损失函数
loss_fn = nn.MSELoss()

# 训练过程示例
x_train = torch.randn(10, 100)  # 训练数据
y_train = torch.randn(10, 50)   # 训练标签
x_val = torch.randn(5, 100)     # 验证数据
y_val = torch.randn(5, 50)      # 验证标签

# 执行MAML更新
maml_update(model, loss_fn, x_train, y_train, x_val, y_val)