使用 PyTorch-BigGraph 构建和部署大规模图嵌入的完整步骤

conda create -n pytorch-biggraph python=3.8  
conda activate pytorch-biggraph

# 如果您有兼容的GPU，请安装支持CUDA的PyTorch  
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

import pandas as pd  
def format_data(df):  
    """  
    将DataFrame格式化为与PyTorch-BigGraph兼容的TSV格式。  
    参数:  
    df (pd.DataFrame): 至少包含节点和边的两列的DataFrame。  
    返回:  
    无（写入TSV文件）  
    """  
    # 假设DataFrame具有'source'、'target'、'relation'列  
    df[['source', 'target', 'relation']].to_csv("graph_data.tsv", sep='\t', index=False, header=False)  
    print("数据已格式化并保存为'graph_data.tsv'。")  
# 示例用法  
data = {'source': [1, 2, 3], 'target': [4, 5, 6], 'relation': ['follows', 'likes', 'shares']}  
df = pd.DataFrame(data)  
format_data(df)

from torchbiggraph.config import parse_config  
from torchbiggraph.partitionserver import start_partition_server  
config = parse_config("config.json")  
# 启动分区服务器（通常在单独的实例上运行）  
start_partition_server(config)

def prepare_nodes_and_edges(df):  
    """  
    通过为每个实体创建唯一ID来准备节点和边。  
    参数:  
    df (pd.DataFrame): 包含'source'和'target'列的DataFrame。  
    返回:  
    dict: 实体到唯一ID的映射。  
    """  
    nodes = set(df['source']).union(set(df['target']))  
    node_mapping = {node: idx for idx, node in enumerate(nodes)}  
    print("节点和边的准备已完成。")  
    return node_mapping  
# 示例用法  
node_mapping = prepare_nodes_and_edges(df)

config = dict(  
    entity_path="data/entities",  
    relation_path="data/relations",  
    checkpoint_path="models/checkpoints",  
    dimension=200,          # 更高的维度用于更复杂的关系  
    num_epochs=10,           # 根据数据集的实验设置  
    batch_size=10000,        # 根据硬件能力进行调整  
    eval_fraction=0.1,       # 使用10%的数据进行评估  
    num_partitions=4,        # 根据图的大小和可用的GPU数量  
    num_gpus=2,              # 如果使用多GPU设置，指定GPU数量  
    bucket_order="random"    # 控制数据桶的处理顺序  
)

import json  
from torchbiggraph.config import ConfigSchema  
from torchbiggraph.train import train  
# 将配置保存为JSON文件  
with open("config.json", "w") as f:  
    json.dump(config, f)  
# 使用配置初始化并训练模型  
train(config_path="config.json")

def train_model(config_path, num_gpus):  
    if num_gpus > 1:  
        # 对于多GPU设置  
        torch.distributed.init_process_group(backend="nccl")  
    train(config_path=config_path)  
# 运行训练  
train_model("config.json", num_gpus=2)

import torch  
# 监控GPU内存使用情况的示例代码  
print("GPU内存使用情况:", torch.cuda.memory_allocated())

from torchbiggraph.checkpoint import save_checkpoint  
# 示例：在每个周期保存一个检查点  
for epoch in range(config['num_epochs']):  
    # 执行训练步骤  
    save_checkpoint(config['checkpoint_path'], model_state)  
    print(f"检查点已保存，周期 {epoch}")

import torch.distributed as dist  
def distributed_training():  
    dist.init_process_group(backend='nccl')  
    # 现在使用分布式后端进行正常训练  
    train(config_path="config.json")  
# 启动分布式训练  
distributed_training()

import optuna  
def objective(trial):  
    # 采样要调优的超参数  
    batch_size = trial.suggest_int("batch_size", 1000, 10000)  
    dimension = trial.suggest_int("dimension", 50, 300)  
    learning_rate = trial.suggest_loguniform("learning_rate", 1e-5, 1e-1)  
    config['batch_size'] = batch_size  
    config['dimension'] = dimension  
    config['lr'] = learning_rate  
    # 使用新超参数训练模型  
    train_model("config.json", config['num_gpus'])  
    # 返回用于优化的评估指标  
    return evaluate_model(config['checkpoint_path'])  
# 运行超参数优化  
study = optuna.create_study(direction="maximize")  
study.optimize(objective, n_trials=50)

import numpy as np  
def evaluate_embeddings(embeddings, test_data):  
    """  
    使用平均倒数排名（MRR）和Hits@K评估嵌入。  
    参数:  
    embeddings (dict): 节点嵌入的字典。  
    test_data (list): 测试边的列表，形式为（source, target）。  
    返回:  
    dict: MRR和Hits@K分数。  
    """  
    mrr = 0  
    hits_at_k = {1: 0, 3: 0, 10: 0}  
    for source, target in test_data:  
        source_emb = embeddings[source]  
        target_emb = embeddings[target]  
        # 计算与所有其他节点的相似性  
        similarities = {node: np.dot(source_emb, embeddings[node]) for node in embeddings}  
        ranked_nodes = sorted(similarities, key=similarities.get, reverse=True)  
        rank = ranked_nodes.index(target) + 1  
        mrr += 1 / rank  
        for k in hits_at_k.keys():  
            if rank <= k:  
                hits_at_k[k] += 1  
    num_edges = len(test_data)  
    mrr /= num_edges  
    hits_at_k = {k: hits / num_edges for k, hits in hits_at_k.items()}  
    return {"MRR": mrr, "Hits@K": hits_at_k}

from sklearn.manifold import TSNE  
import matplotlib.pyplot as plt  
def visualize_embeddings(embeddings):  
    """  
    使用t-SNE可视化嵌入。  
    参数:  
    embeddings (dict): 节点嵌入的字典。  
    返回:  
    无  
    """  
    nodes = list(embeddings.keys())  
    emb_matrix = np.array([embeddings[node] for node in nodes])  
    tsne = TSNE(n_components=2, perplexity=30)  
    emb_2d = tsne.fit_transform(emb_matrix)  
    plt.figure(figsize=(10, 10))  
    plt.scatter(emb_2d[:, 0], emb_2d[:, 1], alpha=0.7)  
    plt.title("嵌入的t-SNE可视化")  
    plt.show()  
# 示例用法  
visualize_embeddings(embeddings)

import redis  
def deploy_embeddings(embeddings, redis_host="localhost", redis_port=6379):  
    """  
    将嵌入存储在Redis中以便快速检索。  
    参数:  
    embeddings (dict): 节点嵌入的字典。  
    redis_host (str): Redis服务器主机名。  
    redis_port (int): Redis服务器端口。  
    返回:  
    无  
    """  
    r = redis.Redis(host=redis_host, port=redis_port)  
    for node, emb in embeddings.items():  
        r.set(node, np.array2string(emb, separator=','))  
    print("嵌入已部署到Redis。")  
# 将嵌入部署到Redis  
deploy_embeddings(embeddings)

def recommend_items(user_embedding, redis_conn, top_k=5):  
    """  
    根据用户嵌入相似性推荐项目。  
    参数:  
    user_embedding (np.array): 用户或源节点的嵌入。  
    redis_conn (Redis): Redis连接对象。  
    top_k (int): 推荐数量。  
    返回:  
    list: 推荐项目。  
    """  
    items = redis_conn.keys()  
    similarities = {}  
    for item in items:  
        item_emb = np.fromstring(redis_conn.get(item)[1:-1], sep=',')  
        similarity = np.dot(user_embedding, item_emb)  
        similarities[item] = similarity  
    # 按相似性排序项目并返回top_k  
    recommended_items = sorted(similarities, key=similarities.get, reverse=True)[:top_k]  
    return recommended_items  
# 示例用法  
user_embedding = np.random.rand(200)  
recommend_items(user_embedding, redis.Redis())

import torch  
# 检查内存使用情况  
print("分配的GPU内存:", torch.cuda.memory_allocated())  
print("缓存的GPU内存:", torch.cuda.memory_reserved())如果内存使用过高，请尝试减少batch_size，降低dimension，或增加num_partitions以在GPU之间分配数据。

!nvcc --version  
import torch  
print("CUDA版本:", torch.version.cuda)  
print("CUDA是否可用?", torch.cuda.is_available())

# 重新安装具有特定CUDA版本的PyTorch  
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

import pandas as pd  
def validate_data_format(file_path):  
    data = pd.read_csv(file_path, sep='\t', header=None)  
    # 确保存在所需的列  
    assert data.shape[1] >= 2, "数据应至少有2列（source, target）"  
    print("数据格式已验证。")  
# 验证数据  
validate_data_format("graph_data.tsv")

import torch.distributed as dist  
def distributed_training(config_path):  
    dist.init_process_group(backend='nccl')  
    train(config_path=config_path)