iOS 上的图像分割 DeepLabV3¶
作者: Jeff Tang
审阅者: Jeremiah Chung
简介¶
语义图像分割是一种计算机视觉任务，它使用语义标签来标记输入图像的特定区域。PyTorch 语义图像分割 DeepLabV3 模型 可用于使用 20 个语义类别 对图像区域进行标记，例如自行车、公共汽车、汽车、狗和人。图像分割模型在自动驾驶和场景理解等应用中非常有用。
在本教程中，我们将提供一个逐步指南，介绍如何准备和运行 PyTorch DeepLabV3 模型在 iOS 上，带您从开始使用您可能想在 iOS 上使用的模型到最终拥有使用该模型的完整 iOS 应用程序。我们还将涵盖有关如何检查您下一个最喜欢的预训练 PyTorch 模型是否可以在 iOS 上运行以及如何避免陷阱的实用且通用的提示。
注意
在学习本教程之前，您应该查看 适用于 iOS 的 PyTorch Mobile 并尝试一下 PyTorch iOS HelloWorld 示例应用程序。本教程将超越图像分类模型，该模型通常是第一个部署在移动设备上的模型。本教程的完整代码可在此处获得 这里.
学习目标¶
在本教程中，您将学习如何
将 DeepLabV3 模型转换为适用于 iOS 部署的模型。
获取模型对示例输入图像的输出，并在 Python 中将其与 iOS 应用程序的输出进行比较。
构建一个新的 iOS 应用或重用 iOS 示例应用来加载转换后的模型。
将输入准备成模型期望的格式并处理模型输出。
完成 UI、重构、构建和运行应用程序，以查看图像分割的实际效果。
先决条件¶
PyTorch 1.6 或 1.7
torchvision 0.7 或 0.8
Xcode 11 或 12
步骤¶
1. 将 DeepLabV3 模型转换为 iOS 部署¶
在 iOS 上部署模型的第一步是将模型转换为TorchScript格式。
注意
并非所有 PyTorch 模型都可以在此时转换为 TorchScript，因为模型定义可能使用 TorchScript 中没有的语言特性，而 TorchScript 是 Python 的一个子集。有关更多详细信息，请参见脚本和优化食谱。
只需运行以下脚本即可生成脚本化的模型deeplabv3_scripted.pt
import torch
# use deeplabv3_resnet50 instead of deeplabv3_resnet101 to reduce the model size
model = torch.hub.load('pytorch/vision:v0.8.0', 'deeplabv3_resnet50', pretrained=True)
model.eval()
scriptedm = torch.jit.script(model)
torch.jit.save(scriptedm, "deeplabv3_scripted.pt")
生成的deeplabv3_scripted.pt模型文件的大小应约为 168MB。理想情况下，模型还应在部署到 iOS 应用程序之前进行量化，以实现显著的尺寸缩减和更快的推理速度。要对量化有一个基本的了解，请参见量化食谱以及其中的资源链接。我们将详细介绍如何在未来的教程或食谱中，正确地将称为训练后静态量化的量化工作流程应用于 DeepLabV3 模型。
2. 在 Python 中获取模型的示例输入和输出¶
现在我们有了脚本化的 PyTorch 模型，让我们用一些示例输入进行测试，以确保模型在 iOS 上正常工作。首先，让我们编写一个 Python 脚本，使用该模型进行推理并检查输入和输出。对于 DeepLabV3 模型的这个示例，我们可以重用步骤 1 中的代码以及DeepLabV3 模型集线器站点中的代码。在上面的代码中添加以下代码片段
from PIL import Image
from torchvision import transforms
input_image = Image.open("deeplab.jpg")
preprocess = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0)
with torch.no_grad():
output = model(input_batch)['out'][0]
print(input_batch.shape)
print(output.shape)
从这里下载deeplab.jpg并运行上面的脚本，查看模型的输入和输出形状
torch.Size([1, 3, 400, 400])
torch.Size([21, 400, 400])
因此，如果您向模型提供大小为 400x400 的相同图像输入deeplab.jpg，则模型的输出大小应为 [21, 400, 400]。您还应打印出输入和输出的实际数据的至少开头部分，以便在下面的步骤 4 中用于与模型在 iOS 应用程序中运行时的实际输入和输出进行比较。
3. 构建一个新的 iOS 应用程序或重用一个示例应用程序并加载模型¶
首先，按照为 iOS 准备模型食谱中的步骤 3 来在启用了 PyTorch Mobile 的 Xcode 项目中使用我们的模型。因为本教程中使用的 DeepLabV3 模型和 PyTorch Hello World iOS 示例中使用的 MobileNet v2 模型都是计算机视觉模型，您可以选择从HelloWorld 示例库开始，将其用作模板来重用加载模型和处理输入和输出的代码。
现在让我们将步骤 2 中使用的deeplabv3_scripted.pt和deeplab.jpg添加到 Xcode 项目中，并修改ViewController.swift以类似于
class ViewController: UIViewController {
var image = UIImage(named: "deeplab.jpg")!
override func viewDidLoad() {
super.viewDidLoad()
}
private lazy var module: TorchModule = {
if let filePath = Bundle.main.path(forResource: "deeplabv3_scripted",
ofType: "pt"),
let module = TorchModule(fileAtPath: filePath) {
return module
} else {
fatalError("Can't load the model file!")
}
}()
}
然后在return module行设置断点，构建并运行应用程序。应用程序应在断点处停止，这意味着步骤 1 中的脚本化模型已成功加载到 iOS 上。
4. 处理模型输入和输出以进行模型推理¶
在上一步骤中模型加载后，让我们验证它是否按预期输入工作并能生成预期输出。由于 DeepLabV3 模型的模型输入是图像，与 Hello World 示例中的 MobileNet v2 相同，我们将重用 Hello World 中TorchModule.mm文件中的部分输入处理代码。将TorchModule.mm中的predictImage方法实现替换为以下代码
- (unsigned char*)predictImage:(void*)imageBuffer {
// 1. the example deeplab.jpg size is size 400x400 and there are 21 semantic classes
const int WIDTH = 400;
const int HEIGHT = 400;
const int CLASSNUM = 21;
at::Tensor tensor = torch::from_blob(imageBuffer, {1, 3, WIDTH, HEIGHT}, at::kFloat);
torch::autograd::AutoGradMode guard(false);
at::AutoNonVariableTypeMode non_var_type_mode(true);
// 2. convert the input tensor to an NSMutableArray for debugging
float* floatInput = tensor.data_ptr<float>();
if (!floatInput) {
return nil;
}
NSMutableArray* inputs = [[NSMutableArray alloc] init];
for (int i = 0; i < 3 * WIDTH * HEIGHT; i++) {
[inputs addObject:@(floatInput[i])];
}
// 3. the output of the model is a dictionary of string and tensor, as
// specified at https://pytorch.ac.cn/hub/pytorch_vision_deeplabv3_resnet101
auto outputDict = _impl.forward({tensor}).toGenericDict();
// 4. convert the output to another NSMutableArray for easy debugging
auto outputTensor = outputDict.at("out").toTensor();
float* floatBuffer = outputTensor.data_ptr<float>();
if (!floatBuffer) {
return nil;
}
NSMutableArray* results = [[NSMutableArray alloc] init];
for (int i = 0; i < CLASSNUM * WIDTH * HEIGHT; i++) {
[results addObject:@(floatBuffer[i])];
}
return nil;
}
注意
DeepLabV3 模型的模型输出是一个字典，因此我们使用toGenericDict来正确提取结果。对于其他模型，模型输出也可能是一个单一张量或一个张量元组，等等。
通过上述代码更改，您可以在填充inputs和results的两个循环后设置断点，并将它们与您在步骤 2 中看到的模型输入和输出数据进行比较，以查看它们是否匹配。对于在 iOS 和 Python 上运行的模型的相同输入，您应该获得相同的输出。
到目前为止，我们所做的只是确认我们感兴趣的模型可以在我们的 iOS 应用程序中像在 Python 中一样被脚本化并正确运行。我们迄今为止为在 iOS 应用程序中使用模型所执行的步骤，占据了应用程序开发时间的大部分，如果不是大部分的话，这与数据预处理是典型机器学习项目中最繁重的任务类似。
5. 完成 UI、重构、构建和运行应用程序¶
现在我们已经准备好完成应用程序和 UI，以实际查看处理后的结果作为新图像。输出处理代码应如下所示，添加到步骤 4 中TorchModule.mm中的代码片段末尾 - 请记住，首先删除return nil;行，该行是临时添加的，以使代码构建和运行
// see the 20 semantic classes link in Introduction
const int DOG = 12;
const int PERSON = 15;
const int SHEEP = 17;
NSMutableData* data = [NSMutableData dataWithLength:
sizeof(unsigned char) * 3 * WIDTH * HEIGHT];
unsigned char* buffer = (unsigned char*)[data mutableBytes];
// go through each element in the output of size [WIDTH, HEIGHT] and
// set different color for different classnum
for (int j = 0; j < WIDTH; j++) {
for (int k = 0; k < HEIGHT; k++) {
// maxi: the index of the 21 CLASSNUM with the max probability
int maxi = 0, maxj = 0, maxk = 0;
float maxnum = -100000.0;
for (int i = 0; i < CLASSNUM; i++) {
if ([results[i * (WIDTH * HEIGHT) + j * WIDTH + k] floatValue] > maxnum) {
maxnum = [results[i * (WIDTH * HEIGHT) + j * WIDTH + k] floatValue];
maxi = i; maxj = j; maxk = k;
}
}
int n = 3 * (maxj * width + maxk);
// color coding for person (red), dog (green), sheep (blue)
// black color for background and other classes
buffer[n] = 0; buffer[n+1] = 0; buffer[n+2] = 0;
if (maxi == PERSON) buffer[n] = 255;
else if (maxi == DOG) buffer[n+1] = 255;
else if (maxi == SHEEP) buffer[n+2] = 255;
}
}
return buffer;
此处的实现基于对 DeepLabV3 模型的理解，该模型为宽度*高度的输入图像输出大小为 [21, 宽度, 高度] 的张量。宽度*高度输出数组中的每个元素的值介于 0 到 20 之间（总共 21 个在简介中描述的语义标签），该值用于设置特定颜色。此处的分割颜色编码基于具有最高概率的类，您可以扩展您自己数据集中的所有类的颜色编码。
在输出处理之后，您还需要调用一个帮助函数来将 RGBbuffer转换为UIImage实例，以便在UIImageView上显示。您可以参考代码库中UIImageHelper.mm中定义的示例代码convertRGBBufferToUIImage。
此应用程序的 UI 也类似于 Hello World 的 UI，只是您不需要UITextView来显示图像分类结果。您还可以添加两个按钮Segment和Restart，如代码库所示，以运行模型推理并在显示分割结果后显示回原始图像。
在我们可以运行应用程序之前，最后一步是将所有部分连接在一起。修改ViewController.swift文件以使用predictImage，该函数已在库中重构并更改为segmentImage，以及您构建的帮助函数，如库中ViewController.swift中的示例代码所示。将按钮连接到操作，您应该就可以开始了。
现在，当您在 iOS 模拟器或实际 iOS 设备上运行应用程序时，您将看到以下屏幕
