(原型) 在 PyTorch 中使用 iOS GPU¶

**作者**：徐涛

简介¶

本教程介绍了在 iOS GPU 上运行模型的步骤。我们将使用 mobilenetv2 模型作为示例。由于移动 GPU 功能目前处于原型阶段，您需要从源代码构建自定义的 PyTorch 二进制文件。目前，仅支持有限数量的运算符，并且某些客户端 API 在未来版本中可能会发生变化。

模型准备¶

由于 GPU 以不同的顺序使用权重，因此我们需要做的第一步是将我们的 TorchScript 模型转换为与 GPU 兼容的模型。此步骤也称为“预打包”。

带有 Metal 的 PyTorch¶

为此，我们将安装包含 Metal 后端的 PyTorch nightly 二进制文件。运行以下命令：

conda install pytorch -c pytorch-nightly
// or
pip3 install --pre torch -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html

此外，您可以从源代码构建自定义的 PyTorch 二进制文件，其中包含 Metal 后端。只需从 github 中检出 PyTorch 源代码并运行以下命令：

cd PYTORCH_ROOT
USE_PYTORCH_METAL_EXPORT=ON python setup.py install --cmake

以上命令将从 master 分支构建自定义的 PyTorch 二进制文件。 install 参数只是告诉 setup.py 覆盖桌面上的现有 PyTorch。构建完成后，打开另一个终端并检查 PyTorch 版本以查看安装是否成功。在编写本食谱时，版本为 1.8.0a0+41237a4。根据您从 master 分支检出代码的时间，您可能会看到不同的数字，但它应该大于 1.7.0。

import torch
torch.__version__ #1.8.0a0+41237a4

Metal 兼容模型¶

下一步将把mobilenetv2 TorchScript模型转换为Metal兼容的模型。我们将利用来自torch.utils模块的optimize_for_mobile API。如下所示

import torch
import torchvision
from torch.utils.mobile_optimizer import optimize_for_mobile

model = torchvision.models.mobilenet_v2(pretrained=True)
scripted_model = torch.jit.script(model)
optimized_model = optimize_for_mobile(scripted_model, backend='metal')
print(torch.jit.export_opnames(optimized_model))
optimized_model._save_for_lite_interpreter('./mobilenetv2_metal.pt')

请注意，torch.jit.export_opnames(optimized_model)将转储optimized_mobile中所有优化的操作符。如果一切正常，您应该能够在控制台中看到以下操作符被打印出来

['aten::adaptive_avg_pool2d',
'aten::add.Tensor',
'aten::addmm',
'aten::reshape',
'aten::size.int',
'metal::copy_to_host',
'metal_prepack::conv2d_run']

这些是在iOS GPU上运行mobilenetv2模型所需的所有操作符。太棒了！现在您已将mobilenetv2_metal.pt保存到磁盘，让我们继续进行iOS部分。

使用支持Metal的PyTorch iOS库¶

支持Metal的PyTorch iOS库LibTorch-Lite-Nightly可在CocoaPods中获取。您可以阅读iOS教程中的在CocoaPods中使用Nightly PyTorch iOS库部分，以详细了解其用法。

我们还有HelloWorld-Metal示例，它展示了如何将所有部分连接在一起。

请注意，如果您运行HelloWorld-Metal示例，您可能会注意到结果与iOS教程中显示的CPU模型的结果略有不同。

- timber wolf, grey wolf, gray wolf, Canis lupus
- malamute, malemute, Alaskan malamute
- Eskimo dog, husky

这是因为默认情况下，Metal使用fp16而不是fp32进行计算。精度损失是预期的。

使用从源代码构建的LibTorch-Lite¶

您还可以从源代码构建自定义LibTorch-Lite，并使用它在iOS Metal上运行GPU模型。在本节中，我们将使用HelloWorld示例来演示此过程。

首先，确保您已从PyTorch根目录中“模型准备”步骤的“build”文件夹中删除了所有内容。然后运行以下命令

IOS_ARCH=arm64 USE_PYTORCH_METAL=1 ./scripts/build_ios.sh

注意IOS_ARCH告诉脚本构建arm64版本的Libtorch-Lite。这是因为在PyTorch中，Metal仅适用于支持Apple A9芯片或更高芯片的iOS设备。构建完成后，请按照iOS教程中的从源代码构建PyTorch iOS库部分正确设置XCode设置。不要忘记将./mobilenetv2_metal.pt复制到您的XCode项目中，并相应地修改模型文件路径。

接下来，我们需要在TorchModule.mm中进行一些更改

...
// #import <Libtorch-Lite/Libtorch-Lite.h>
// If it's built from source with Xcode, comment out the line above
// and use following headers
#include <torch/csrc/jit/mobile/import.h>
#include <torch/csrc/jit/mobile/module.h>
#include <torch/script.h>
...

- (NSArray<NSNumber*>*)predictImage:(void*)imageBuffer {
  c10::InferenceMode mode;
  at::Tensor tensor = torch::from_blob(imageBuffer, {1, 3, 224, 224}, at::kFloat).metal();
  auto outputTensor = _impl.forward({tensor}).toTensor().cpu();
  ...
}
...

如您所见，我们只需调用.metal()将输入张量从CPU移动到GPU，然后调用.cpu()将结果移回。在内部，.metal()会将输入数据从CPU缓冲区复制到具有GPU兼容内存格式的GPU缓冲区。当调用.cpu()时，GPU命令缓冲区将被刷新和同步。在forward完成之后，最终结果将从GPU缓冲区复制回CPU缓冲区。

我们必须做的最后一步是将Accelerate.framework和MetalPerformanceShaders.framework添加到您的xcode项目中（通过XCode打开您的项目，转到项目的“General”选项卡，找到“Frameworks, Libraries and Embedded Content”部分，然后单击“+”按钮）。

如果一切正常，您应该能够在手机上看到推理结果。

结论¶

在本教程中，我们演示了如何将mobilenetv2模型转换为GPU兼容模型。我们通过HelloWorld示例演示了如何使用C++ API在iOS GPU上运行模型。请注意，GPU功能仍在开发中，将继续添加新的操作符。API可能会在将来的版本中发生更改。

感谢阅读！与往常一样，我们欢迎任何反馈，如果您有任何反馈，请在此处创建问题。

了解更多¶

来自Torchvision的Mobilenetv2
要了解有关如何使用optimize_for_mobile的更多信息，请参阅移动性能配方