注意
跳到末尾 下载完整的示例代码。
提升多核 CPU 上的可扩展性
图神经网络 (GNN) 训练在多核 CPU 上存在可扩展性差的问题。具体来说，性能通常在 16 核时达到上限，使用超过 16 核时性能没有提升[1]。ARGO 是一个提供可扩展性能的运行时系统。启用 ARGO 后，我们能够在超过 64 核上进行扩展，ARGO 在 Xeon 8380H 和 Xeon 6430L 上分别将 GNN 训练（按 epoch 时间计）加速高达 4.30 倍和 3.32 倍[2]。本章重点介绍如何设置 ARGO，以释放多核 CPU 的潜力来加速 GNN 训练。
安装
ARGO 利用 scikit-optimize 库进行自动调优。请安装 scikit-optimize 以运行 ARGO：.. code-block:: shell
conda install -c conda-forge “scikit-optimize>=0.9.0”
或 .. code-block:: shell
pip install scikit-optimize>=0.9
在您自己的 GNN 程序中启用 ARGO
在本节中，我们提供了一个逐步教程，介绍如何在 DGL 程序中启用 ARGO。我们使用 ogb_example.py [3] 作为示例。.. note
We also provide the complete example file *ogb_example_ARGO.py* [#f4]_
which followed the steps below to enable ARGO on *ogb_example.py*.
步骤 1
首先，在文件顶部包含所有必需的包。请将您的文件和 argo.py [5] 放在同一目录下。
步骤 2
设置 PyTorch 分布式数据并行 (DDP)
2.1. 在训练程序顶部添加初始化函数，并使用 DDP wrapper 包装 `model` .. code-block:: python
- def train(…)
dist.init_process_group(‘gloo’, rank=rank, world_size=world_size) # 新添加 model = SAGE(…) # 原始代码 model = DistributedDataParallel(model) # 新添加 …
2.2. 在主程序中，在启动训练函数之前添加以下内容 .. code-block:: python
… os.environ[‘MASTER_ADDR’] = ‘127.0.0.1’ os.environ[‘MASTER_PORT’] = ‘29501’ mp.set_start_method(‘fork’, force=True) train(args, device, data) # 启动训练函数的原始代码
步骤 3
通过初始化运行时系统并包装训练函数来启用 ARGO .. code-block:: python
runtime = ARGO(n_search = 15, epoch = args.num_epochs, batch_size = args.batch_size) # 初始化 runtime.run(train, args=(args, device, data)) # 包装训练函数
注意
ARGO 接受三个输入参数：搜索次数 n_search、epoch 数和 mini-batch 大小。增加 n_search 可能找到更好的配置并减少 epoch 时间；然而，搜索本身也会带来额外的开销。我们建议将 n_search 设置在 15 到 45 之间，以获得最佳整体性能。
步骤 4
修改训练函数的输入，在原始输入后直接添加 ARGO 参数。
这是原始函数：.. code-block:: python
def train(args, device, data)
添加以下变量：rank, world_size, comp_core, load_core, counter, b_size, ep .. code-block:: python
def train(args, device, data, rank, world_size, comp_core, load_core, counter, b_size, ep)
步骤 5
修改训练函数中的 dataloader 函数 .. code-block:: python
- dataloader = dgl.dataloading.DataLoader(
g, train_nid, sampler, batch_size=b_size, # 已修改 shuffle=True, drop_last=False, num_workers=len(load_core), # 已修改 use_ddp = True) # 新添加
步骤 6
通过在训练 for 循环前添加 enable_cpu_affinity() 来启用核心绑定，并将 epoch 数更改为变量 ep：.. code-block:: python
- with dataloader.enable_cpu_affinity(loader_cores=load_core, compute_cores=comp_core)
for epoch in range(ep): # 将 num_epochs 更改为 ep
步骤 7
最后一步！在训练前加载模型并在训练后保存。
原始程序：.. code-block:: python
- with dataloader.enable_cpu_affinity(loader_cores=load_core, compute_cores=comp_core)
for epoch in range(ep): … # 训练操作
修改后：.. code-block:: python
PATH = “model.pt” if counter[0] != 0
checkpoint = th.load(PATH) model.load_state_dict(checkpoint[‘model_state_dict’]) optimizer.load_state_dict(checkpoint[‘optimizer_state_dict’]) epoch = checkpoint[‘epoch’] loss = checkpoint[‘loss’]
- with dataloader.enable_cpu_affinity(loader_cores=load_core, compute_cores=comp_core)
for epoch in range(ep): … # 训练操作
dist.barrier() if rank == 0
- th.save({‘epoch’: counter[0],
‘model_state_dict’: model.state_dict(), ‘optimizer_state_dict’: optimizer.state_dict(), ‘loss’: loss, }, PATH)
步骤 8
完成！您现在可以运行启用了 ARGO 的 GNN 程序了。.. code-block:: shell
python <your_code>.py
脚注
脚本总运行时间： (0 minutes 0.000 seconds)