作者：Intel

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

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

我将引导您完成使用 SpaceNet5 卫星图像数据集的步骤，以及我是如何通过几个关键开关优化代码，使得深度学习工作负载在 CPU 上变得可行。

在我们开始之前，做一些前期准备…

本文附带的代码位于 Intel Extension for PyTorch 仓库的 examples 文件夹中。我大量借鉴了 City-Scale Road Extraction from Satellite Imagery (CRESI) 仓库的代码。我针对第 4 代 Intel Xeon 处理器进行了适配，并加入了 PyTorch 优化和 Intel Extension for PyTorch 优化。特别是，我能够利用此处的 notebook 拼凑出一个工作流程。

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

我也强烈推荐以下文章，它们详细解释了如何开始使用 SpaceNet5 数据

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

在主要云服务提供商 (CSP) 上使用 CPU 实例而非 GPU 实例可以节省大量成本。最新的处理器仍在向 CSP 推广，因此我使用了托管在 Intel® Developer Cloud（您可以在此处注册 Beta：cloud.intel.com）上的第 4 代 Intel Xeon 处理器。

在 AWS 上，您可以选择 r7iz.* EC2 实例，前提是您在此处注册预览（图 1）。在撰写本文时，新的 AI 加速引擎 Intel® Advanced Matrix Extensions (Intel® AMX) 仅在裸机上可用，但很快将在虚拟机上启用。

List of 4th Gen Xeon instances on AWS EC2

图 1。AWS EC2 上的第 4 代 Xeon 实例列表（图片由作者提供）

在 Google Cloud* Platform 上，您可以选择第 4 代 Xeon Scalable 处理器 C3 VM（图 2）。

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

图 2。Google Cloud Platform 上的第 4 代 Intel Xeon Scalable 处理器实例列表（图片由作者提供）

硬件介绍和优化

第 4 代 Intel Xeon 处理器于 2023 年 1 月发布，我使用的裸机实例有两个插槽（每个有 56 个物理核心）、504 GB 内存和 Intel 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 Challenge 的卫星图像数据集。可以从 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 x 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 的训练/验证集划分，并确保指向正确的卫星图像和相应训练掩模文件夹。配置参数在 Intel Extension for PyTorch 在 GitHub 示例下的 notebook 中有更详细的解释。

训练 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()。

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

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 在一个插槽上使用 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 语句中添加 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 训练，包括使用 Intel Extension for PyTorch 和 BF16
  • 在第 4 代 Intel Xeon CPU 上训练了一个高性能的 ResNet34 + UNet 模型
  • 运行了初步推理以查看限速掩模的预测结果

您可以在此处找到第 4 代 Intel Xeon CPU 的详细基准测试结果：edc.intel.com/content/www/us/en/products/performance/benchmarks/4th-generation-intel-xeon-scalable-processors/。

下一步

使用 Intel Extension for PyTorch 扩展在 Intel CPU 上的优化

pip install intel-extension-for-pytorch

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

如果您有更多问题，请在 LinkedIn 上与我联系！

有关 Intel Extension for PyTorch 的更多信息，请参见此处。

获取软件

我鼓励您了解 Intel 的其他人工智能工具和框架优化，并了解构成 Intel 人工智能软件产品组合基础的开放式、基于标准的oneAPI 多架构、多供应商编程模型。

有关第 4 代 Intel Xeon Scalable 处理器的更多详细信息，请访问人工智能平台，在那里您可以了解 Intel 如何赋能开发者运行高性能、高效的端到端人工智能流水线。

PyTorch 资源