由英特尔提供

使用 Intel® Extension for PyTorch 提升图像处理性能

PyTorch 提供了出色的 CPU 性能，并且可以通过 Intel® Extension for PyTorch 进一步加速。我使用 PyTorch 1.13.1（带有 ResNet34 + UNet 架构）训练了一个 AI 图像分割模型，从卫星图像中识别道路和限速标志，所有这些都在第 4 代 Intel® Xeon®可扩展处理器上完成。

我将向您介绍如何使用名为 SpaceNet5 的卫星图像数据集，以及我是如何通过切换几个关键开关来优化代码，使深度学习工作负载在 CPU 上变得可行。

在我们开始之前，做一些准备工作...

本文所附带的代码可在 Intel Extension for PyTorch 仓库的示例文件夹中找到。我从 City-Scale Road Extraction from Satellite Imagery (CRESI)仓库中借鉴了很多。我将其适配到第 4 代 Intel Xeon 处理器，并进行了 PyTorch 优化和 Intel Extension for PyTorch 优化。特别是，我能够使用这里的工作簿构建一个工作流程。

您可以在 YouTube 上找到我做的相关演讲。

我还强烈推荐以下文章，以详细了解如何开始使用 SpaceNet5 数据：

我参考了 Julien Simon 在 Hugging Face 上的两篇博客；他在 AWS 实例 r7iz.metal-16xl 上进行了测试：

在主要云服务提供商（CSP）上使用 CPU 实例而不是 GPU 实例可能带来的潜在成本节约是显著的。最新的处理器仍在向 CSPs 推出，因此我使用的是托管在英特尔®开发者云上的第 4 代英特尔 Xeon 处理器（您可以在以下链接注册 Beta 测试：cloud.intel.com）。

在 AWS 上，您可以在注册此预览后选择 r7iz.* EC2 实例（图 1）。在撰写本文时，新的 AI 加速引擎英特尔®高级矩阵扩展（Intel® AMX）仅在裸机可用，但很快也将在虚拟机上启用。

List of 4th Gen Xeon instances on AWS EC2

图 1. AWS EC2 上的第 4 代 Xeon 实例列表（作者图片）

在 Google Cloud*平台上，您可以选择第 4 代 Xeon 可扩展处理器 C3 虚拟机（图 2）。

List of 4th Gen Intel Xeon Scalable processor instances on Google Cloud Platform

图 2. Google Cloud 平台上的第 4 代英特尔 Xeon 可扩展处理器实例列表（作者图片）

硬件介绍与优化

第四代英特尔至强处理器于 2023 年 1 月发布，我使用的裸机实例具有两个插槽（每个插槽有 56 个物理核心），504GB 的内存，以及英特尔 AMX 加速。我在后端安装了一些关键库来控制并监控我使用的 CPU 上的插槽、内存和核心：

numactl （含 sudo apt-get install numactl ）

libjemalloc-dev （含 sudo apt-get install libjemalloc ）

intel-openmp （含 conda install intel-openmp ）

gperftools （带有 conda install gperftools -c conda-forge ）

PyTorch 和 Intel Extension for PyTorch 都提供了辅助脚本，因此无需显式使用 intel-openmp 和 numactl ，但它们需要在后端安装。如果您想为其他工作设置它们，以下是我用于 OpenMP*的设置...

export OMP_NUM_THREADS=36
export KMP_AFFINITY=granularity=fine,compact,1,0
export KMP_BLOCKTIME=1

…其中 OMP_NUM_THREADS 是分配给作业的线程数， KMP_AFFINITY 影响线程亲和力设置（包括将线程紧密打包，线程固定状态）， KMP_BLOCKTIME 设置空闲线程在进入睡眠之前应等待的毫秒数。

这里是我用于 numactl 的设置...

numactl -C 0-35 --membind=0 train.py

…其中 -C 指定使用哪些核心， --membind 指示程序仅使用一个插槽（在这种情况下为插槽 0）。

SpaceNet 数据

我正在使用来自 SpaceNet 5 挑战赛的卫星图像数据集。不同的城市可以从 AWS S3 存储桶免费下载：

aws s3 ls s3://spacenet-dataset/spacenet/SN5_roads/tarballs/ --human-readable

2019-09-03 20:59:32    5.8 GiB SN5_roads_test_public_AOI_7_Moscow.tar.gz
2019-09-24 08:43:02    3.2 GiB SN5_roads_test_public_AOI_8_Mumbai.tar.gz
2019-09-24 08:43:47    4.9 GiB SN5_roads_test_public_AOI_9_San_Juan.tar.gz
2019-09-14 13:13:26   35.0 GiB SN5_roads_train_AOI_7_Moscow.tar.gz
2019-09-14 13:13:34   18.5 GiB SN5_roads_train_AOI_8_Mumbai.tar.gz

您可以使用以下命令下载和解压文件：

aws s3 cp s3://spacenet-dataset/spacenet/SN5_roads/tarballs/SN5_roads_train_AOI_7_Moscow.tar.gz .
tar -xvzf ~/spacenet5data/moscow/SN5_roads_train_AOI_7_Moscow.tar.gz

数据集准备

我使用了莫斯科卫星图像数据集，该数据集包含 1,352 张 1,300×1,300 像素的图像，以及相应的街道标签文本文件。数据集包含 8 波段多光谱图像和 3 波段 RGB 图像。图 3 展示了四个样本 RGB 卫星图像及其对应的生成掩码。我使用了 CRESI 仓库中的 speed_masks.py 脚本来生成分割掩码。

Satellite image 3-channel RGB chips from Moscow (top row) and corresponding pixel segmentation masks with varying speed limits

图 3. 莫斯科卫星图像 3 通道 RGB 芯片（顶部行）及其对应不同速度限制的像素分割掩码（底部行）（图片由作者提供）

存在一个 JSON 配置文件，必须更新所有剩余组件：训练和验证分割、训练和推理。示例配置在此处。我执行了 80:20 的训练/验证分割，确保指向正确的卫星图像文件夹及其对应掩码的训练文件夹。配置参数在 GitHub 上 Intel Extension for PyTorch 的示例笔记本中有更详细的解释。

训练 ResNet34 + UNet 模型

我对下面描述的 cresi 代码进行了一些修改，以便在 CPU 上运行并优化训练。要在 CPU 上原生运行，请在 train.py 脚本中将 self.model = nn.DataParallel(model).cuda() 替换为 self.model = nn.DataParallel(model) 。在 01_train.py 脚本中，删除 torch.randn(10).cuda() 。

为了优化训练，将 import intel_extension_for_pytorch as ipex 添加到 train.py 脚本的导入语句中。在定义模型和优化器之后：

self.model = nn.DataParallel(model)
self.optimizer = optimizer(self.model.parameters(), lr=config.lr)

添加 ipex.optimize 行以使用 BF16 精度，而不是 FP32：

self.model, self.optimizer = ipex.optimize(self.model, 
    optimizer=self.optimizer,dtype=torch.bfloat16)

在运行前向传递和计算损失函数之前添加一行以进行混合精度训练：

with torch.cpu.amp.autocast():
    if verbose:
        print("input.shape, target.shape:", input.shape, target.shape)
    output = self.model(input)
    meter = self.calculate_loss_single_channel(output, target, meter, training, iter_size)

现在我们已经优化了训练代码，我们可以开始训练我们的模型了。

就像 SpaceNet 5 竞赛的获胜者一样，我训练了一个 ResNet34 编码器+UNet 解码器模型。该模型从 ImageNet 权重中预训练，训练过程中主干网络完全未冻结。可以使用 01_train.py 脚本运行训练，但为了控制硬件的使用，我使用了一个辅助脚本。实际上有两个辅助脚本：一个是随 PyTorch 标准版提供的，另一个是随 Intel Extension for PyTorch 提供的。它们都完成了相同的功能，但来自标准 PyTorch 的第一个是 torch.backends.xeon.run_cpu ，来自 Intel Extension for PyTorch 的第二个是 ipexrun 。

下面是我命令行中运行的内容：

python -m torch.backends.xeon.run_cpu --ninstances 1 \
  --ncores_per_instance 32 \
  --log_path /home/devcloud/spacenet5data/moscow/v10_xeon4_devcloud22.04/logs/run_cpu_logs \
  /home/devcloud/cresi/cresi/01_train.py \
  /home/devcloud/cresi/cresi/configs/ben/v10_xeon4_baseline_ben.json --fold=0

ipexrun --ninstances 1 \
--ncore_per_instance 32 \
/home/devcloud/cresi/cresi/01_train.py \
/home/devcloud/cresi/cresi/configs/ben/v10_xeon4_baseline_ben.json --fold=0

在这两种情况下，我要求 PyTorch 在一个 socket 上使用 32 个核心进行训练。运行后，我会打印出后端设置的环境变量，以了解 PyTorch 如何使用硬件：

INFO - Use TCMalloc memory allocator
INFO - OMP_NUM_THREADS=32
INFO - Using Intel OpenMP
INFO - KMP_AFFINITY=granularity=fine,compact,1,0
INFO - KMP_BLOCKTIME=1
INFO - LD_PRELOAD=/home/devcloud/.conda/envs/py39/lib/libiomp5.so:/home/devcloud/.conda/envs/py39/lib/libtcmalloc.so
INFO - numactl -C 0-31 -m 0 /home/devcloud/.conda/envs/py39/bin/python -u 01_train.py configs/ben/v10_xeon4_baseline_ben.json --fold=0

在训练过程中，我确保我的总损失函数在减小（即模型正在收敛到解决方案）。

推理

训练好模型后，我们可以开始仅从卫星图像中进行预测。在 eval.py 推理脚本中，将 import intel_extension_for_pytorch as ipex 添加到导入语句中。在加载 PyTorch 模型后，使用 Intel Extension for PyTorch 优化模型以进行 BF16 推理：

model = torch.load(os.path.join(path_model_weights, 
    'fold{}_best.pth'.format(fold)), 
    map_location = lambda storage, 
    loc: storage)
model.eval()
model = ipex.optimize(model, dtype = torch.bfloat16)

在运行预测之前，请添加两行用于混合精度：

with torch.no_grad():
    with torch.cpu.amp.autocast():
        for data in pbar:
            samples = torch.autograd.Variable(data['image'], volatile=True)
            predicted = predict(model, samples, flips=self.flips)

要运行推理，我们可以使用 02_eval.py 脚本。现在我们有了训练好的模型，我们可以在卫星图像上（图 4）进行预测。我们可以看到它似乎将道路与图像紧密对应！

Moscow satellite image and accompanying prediction of roads

图 4. 莫斯科卫星图像及其伴随的道路预测（图片由作者提供）

我意识到我训练的模型过度拟合了莫斯科图像数据，可能无法很好地推广到其他城市。然而，这个挑战的获胜方案使用了六个城市（拉斯维加斯、巴黎、上海、喀土穆、莫斯科、孟买）的数据，并在新城市上表现良好。未来，值得测试的一件事是在所有六个城市进行训练，并在另一个城市进行推理以重现他们的结果。

后处理笔记

还可以执行进一步的后处理步骤，将掩码作为图特征添加到地图中。您可以在此处了解更多关于后处理步骤的信息：

SpaceNet 5 基线 — 第 3 部分：从卫星影像中提取道路速度矢量

后处理脚本

结论

总结来说，我们：

  • 创建了 1,352 个图像训练掩码（含限速），以对应我们的训练卫星图像数据（来自.geojson 文本文件标签）
  • 定义了我们的训练和推理配置文件
  • 将我们的数据分为训练集和验证集
  • 优化了我们的代码以进行 CPU 训练，包括使用 PyTorch 的 Intel 扩展和 BF16
  • 在第 4 代 Intel Xeon CPU 上训练了一个性能良好的 ResNet34 + UNet 模型
  • 运行初始推理以查看速度限制掩码的预测结果

您可以在此处找到关于第 4 代英特尔至强 CPU 的详细基准测试。

下一步

通过使用英特尔 PyTorch 扩展来扩展英特尔 CPU 上的优化：

pip install intel-extension-for-pytorch

git clone https://github.com/intel/intel-extension-for-pytorch

如果您有任何更多问题，请通过 LinkedIn 联系我！

更多关于英特尔 PyTorch 扩展的信息，请在此处查看。

获取软件

我鼓励您查看英特尔的其他 AI 工具和框架优化，并了解作为英特尔 AI 软件组合基础的开放、基于标准的 oneAPI 多架构、多厂商编程模型。

想了解更多关于第 4 代英特尔至强可扩展处理器的详细信息，请访问 AI 平台，在那里您可以了解英特尔如何赋能开发者运行高性能、高效的端到端 AI 管道。

PyTorch 资源