PyTorch Geometric在Intel XPU设备上的环境配置指南

PyTorch Geometric作为图神经网络领域的重要框架，近年来开始支持Intel XPU设备的加速计算。本文将详细介绍如何在Intel GPU设备上配置PyTorch Geometric的运行环境，包括基础环境搭建和Docker容器化部署两种方案。

硬件与软件基础要求

在开始配置前，需要确保系统满足以下基本要求：

• 配备Intel Arc系列或Data Center GPU Max系列的硬件设备
• 操作系统建议使用Ubuntu 22.04 LTS或更新版本
• 已安装Intel GPU驱动程序(version 1.3.0或更高)

基础环境配置

1. 安装Intel oneAPI基础工具包： 首先需要配置Intel的官方软件仓库，然后安装基础运行环境组件：

   wget https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB
   sudo apt-key add GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB
   echo "deb https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list
   sudo apt update
   sudo apt install intel-basekit
   

2. 设置环境变量： 每次使用前需要加载oneAPI的环境配置：

   source /opt/intel/oneapi/setvars.sh
   

3. 安装PyTorch与IPEX扩展： 需要安装专门为Intel GPU优化的PyTorch版本：

   pip install torch==2.1.0a0 torchvision==0.16.0a0 torchaudio==2.1.0a0 \
   --index-url https://pypi.org/simple/
   pip install intel_extension_for_pytorch==2.1.10+xpu -f https://developer.intel.com/ipex-whl-stable-xpu
   

4. 安装PyTorch Geometric： 最后安装PyG框架及其依赖：

   pip install torch_geometric
   pip install torch_scatter torch_sparse -f https://pypi.org/simple/
   

Docker容器化部署方案

对于希望快速部署的用户，可以使用预先配置好的Docker环境：

1. 准备Dockerfile：

   FROM ubuntu:22.04
   RUN apt update && apt install -y wget gnupg
   RUN wget https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB
   RUN apt-key add GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB
   RUN echo "deb https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
   RUN apt update && apt install -y intel-basekit
   RUN apt install -y python3-pip
   RUN pip install torch==2.1.0a0 torchvision==0.16.0a0 torchaudio==2.1.0a0 \
   --index-url https://pypi.org/simple/
   RUN pip install intel_extension_for_pytorch==2.1.10+xpu -f https://developer.intel.com/ipex-whl-stable-xpu
   RUN pip install torch_geometric
   RUN pip install torch_scatter torch_sparse -f https://pypi.org/simple/
   

2. 构建并运行容器：

   docker build -t pyg-xpu .
   docker run -it --device /dev/dri --ipc=host pyg-xpu bash
   

环境验证

配置完成后，可以通过以下Python代码验证环境是否正常工作：

import torch
import torch_geometric

print(f"PyTorch版本: {torch.__version__}")
print(f"PyG版本: {torch_geometric.__version__}")
print(f"XPU可用: {torch.xpu.is_available()}")
print(f"XPU设备数: {torch.xpu.device_count()}")

常见问题解决

1. 驱动程序问题： 如果遇到设备不可用的情况，首先检查驱动程序版本是否符合要求，可以使用以下命令：

   sudo apt install intel-opencl-icd intel-level-zero-gpu level-zero
   sudo reboot
   

2. 内存不足问题： Intel GPU显存管理较为特殊，建议在代码中定期调用：

   torch.xpu.empty_cache()
   

3. 多卡训练配置： 使用多卡时需要显式设置设备可见性：

   os.environ["ZE_AFFINITY_MASK"] = "0,1"  # 使用前两张卡
   

通过以上步骤，开发者可以在Intel XPU设备上充分利用PyTorch Geometric的图神经网络计算能力，获得显著的性能提升。对于大规模图数据，建议结合PyG的分布式采样功能，充分发挥Intel多GPU的并行计算优势。