CoreMLTools转换DeepLabv3模型时输出格式问题的解决方案

问题背景

在使用CoreMLTools将PyTorch的DeepLabv3模型转换为mlpackage格式时，开发者遇到了一个关于输出格式的编译错误。当尝试将模型输出指定为ImageType格式时，系统提示"main output 'output' must be a tensor of type Float32 or Float16"的错误信息。

问题分析

DeepLabv3模型默认的输出是一个四维张量(1×num_classes×height×width)，包含了每个像素点对所有类别的预测概率。当开发者尝试将这个输出转换为ImageType格式时，需要先通过torch.argmax操作获取每个像素点概率最高的类别索引。

然而，torch.argmax操作返回的是整数类型的张量，而CoreML的ImageType输出要求必须是Float32或Float16类型的张量。这就是导致编译错误的根本原因。

解决方案

解决这个问题的关键步骤是在模型的前向传播过程中，将argmax的结果显式转换为浮点类型。具体实现如下：

class WrappedDeeplabv3Resnet101(nn.Module):
    def __init__(self):
        super(WrappedDeeplabv3Resnet101, self).__init__()
        self.model = torch.hub.load(
            "pytorch/vision:v0.6.0",
            "deeplabv3_resnet101",
            pretrained=True
        )
        self.model.eval()

    def forward(self, input):
        res = self.model(input)["out"]
        return torch.argmax(res, dim=1, keepdim=True).float()  # 关键修改：添加.float()

完整转换代码

以下是完整的模型转换代码示例：

import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
import coremltools as ct

# 定义包装模型类
class WrappedDeeplabv3Resnet101(nn.Module):
    def __init__(self):
        super(WrappedDeeplabv3Resnet101, self).__init__()
        self.model = torch.hub.load(
            "pytorch/vision:v0.6.0",
            "deeplabv3_resnet101",
            pretrained=True
        )
        self.model.eval()

    def forward(self, input):
        res = self.model(input)["out"]
        return torch.argmax(res, dim=1, keepdim=True).float()

# 准备输入数据
input_image = Image.open("images/cat_dog.jpg")
preprocess = transforms.Compose([
    transforms.ToTensor()
])
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0)

# 转换模型
torch_model = WrappedDeeplabv3Resnet101().eval()
traced_model = torch.jit.trace(torch_model, input_batch)

# 设置输入输出参数
scale = 1/(0.226*255.0)
bias = [-0.485/(0.229), -0.456/(0.224), -0.406/(0.225)]

image_input = ct.ImageType(
    name="input",
    shape=input_batch.shape,
    scale=scale, 
    bias=bias
)

image_output = ct.ImageType(
    name="output",
    color_layout=ct.colorlayout.GRAYSCALE
)

# 执行转换
mlmodel_from_trace = ct.convert(
    traced_model,
    inputs=[image_input],
    outputs=[image_output]
)

# 保存模型
mlmodel_from_trace.save("SegmentationModel.mlpackage")

技术要点

1. 数据类型一致性：CoreML对模型输入输出的数据类型有严格要求，特别是当使用ImageType格式时，必须确保数据是浮点类型。

2. 图像预处理：在定义ImageType输入时，需要正确设置scale和bias参数，这些值应与模型训练时使用的归一化参数一致。

3. 输出格式：对于分割任务，输出通常是一个单通道的灰度图像，每个像素值代表预测的类别索引。

4. 模型追踪：使用torch.jit.trace时，需要提供具有代表性的输入数据，以确保模型的所有分支都能被正确追踪。

常见问题

1. 为什么其他实现不需要.float()转换？ 这可能是因为其他实现可能在模型内部已经处理了数据类型转换，或者使用了不同的CoreMLTools版本，其类型检查机制可能有所不同。

2. 如何验证转换后的模型？ 可以使用CoreMLTools提供的预览功能，或者直接在Xcode中导入模型进行测试。

3. 性能考虑：

   • 使用ImageType输入输出通常比MLMultiArray更高效
   • 对于移动端部署，可以考虑量化模型以减少大小和提高推理速度

通过遵循上述解决方案和技术要点，开发者可以成功将DeepLabv3模型转换为CoreML格式，并在iOS/macOS应用中实现高效的图像分割功能。