注意
Go to the end to download the full example code
提高多核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”
或 .. 代码块:: shell
pip 安装 scikit-optimize>=0.9
在您自己的GNN程序中启用ARGO
在本节中，我们提供了一个逐步教程，介绍如何在DGL程序上启用ARGO。 我们使用ogb_example.py [3]作为示例。 .. 注意:
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包装器包装`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、训练轮数和 mini-batch 大小。 增加 n_search 可能会在更少的训练轮数内找到更好的配置； 然而，搜索本身也会带来额外的开销。我们建议将 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
通过在训练循环之前添加enable_cpu_affinity()来启用核心绑定，并将训练轮数更改为变量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” 如果 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() 如果 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
.py
脚注
Total running time of the script: (0 minutes 0.000 seconds)