使用 PyTorch 加速图像分割

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

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

我将引导您完成使用名为 SpaceNet5 的卫星图像数据集的步骤，以及我如何通过简单地切换几个关键开关来优化代码，从而在 CPU 上实现深度学习工作负载。

在我们开始之前，有一些准备工作……

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

您可以在 YouTube 上 找到我所做的随附演讲。

我还强烈推荐这些文章，它们详细解释了如何开始使用 SpaceNet5 数据：

• SpaceNet 5 基线 — 第 1 部分：图像和标签准备
• SpaceNet 5 基线 — 第 2 部分：训练道路速度分割模型
• SpaceNet 5 基线 — 第 3 部分：从卫星图像中提取道路速度矢量
• SpaceNet 5 获胜模型发布：道路的尽头

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

• 使用 Intel Sapphire Rapids 加速 PyTorch Transformers，第 1 部分
• 使用 Intel Sapphire Rapids 加速 PyTorch Transformers，第 2 部分

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

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

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

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

图 2. Google Cloud Platform 上的第四代 Intel Xeon 可扩展处理器实例列表（图片作者自绘）

硬件介绍和优化

第四代 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 挑战赛 的卫星图像数据集。可以从 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 脚本生成了分割掩码。

图 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 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 在一个 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）。我们可以看到它似乎将道路紧密地映射到图像上！

图 4. 莫斯科卫星图像和相应的道路预测（图片作者自绘）

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

关于后处理的说明

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

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

后处理脚本

结论

总之，我们

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

您可以在此处找到第四代 Intel Xeon CPU 的详细基准测试。

下一步

通过使用 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 的其他 AI 工具 和 框架 优化，并了解开放的、基于标准的 oneAPI 多架构、多供应商编程模型，它构成了 Intel AI 软件产品组合的基础。

有关第四代 Intel Xeon 可扩展处理器的更多详细信息，请访问 AI 平台，您可以在其中了解 Intel 如何赋能开发者运行高性能、高效的端到端 AI 流水线。

PyTorch 资源

• PyTorch 入门
• 开发者讨论
• 文档