CoreML转换PyTorch模型后无法输出指定长度数组的求助

问题背景 #

我是CoreML新手，遇到一个棘手的问题：我有一个用于多标签分类的PyTorch模型，不需要softmax层，输出是长度为100的数组（对应每个分类的独立概率）。想用coremltools转成MLModel在Swift应用里用，但转换后Xcode显示输出是单个Float32，虽然能当作数组访问，但偶尔会触发内存错误EXC_BAD_ACCESS (code=1, address=0x13fcffff0)，开启AddressSanitizer后每次都崩溃，推测是模型元数据里预期输出单个值，实际处理时当作数组访问导致内存越界。

测试模型及转换代码如下：

import coremltools as ct
import torch
from torchinfo import summary

# From pytorch: https://pytorch.org/tutorials/beginner/introyt/modelsyt_tutorial.html 
class TinyModel(torch.nn.Module):

    def __init__(self):
        super(TinyModel, self).__init__()

        self.linear1 = torch.nn.Linear(100, 200)
        self.activation = torch.nn.ReLU()
        self.linear2 = torch.nn.Linear(200, 100)

    def forward(self, x):
        x = self.linear1(x)
        x = self.activation(x)
        x = self.linear2(x)
        return x

def main():
    guitar_net = TinyModel()
    guitar_net.eval()

    # Trace the model with random data.
    example_input = torch.rand(100)

    summary(guitar_net, (100,))

    traced_model = torch.jit.trace(guitar_net, example_input)
    out = traced_model(example_input)

    model = ct.convert(
        model=traced_model,
        inputs=[ct.TensorType(shape=example_input.shape)],
    )

    model.save("TorchModel.mlmodel")


if __name__ == '__main__':
    main()

解决方案 #

问题根源是CoreML对PyTorch一维张量的形状推断出现偏差，误将输出识别为标量。可以通过以下两种方式修复：

方法1：显式指定输出形状 #

在ct.convert时，通过outputs参数明确告诉CoreML输出是长度为100的张量，强制修正元数据中的输出形状：

model = ct.convert(
    model=traced_model,
    inputs=[ct.TensorType(shape=example_input.shape)],
    # 显式定义输出的形状为(100,)
    outputs=[ct.TensorType(shape=(100,))]
)

方法2：使用带批量维度的输入 #

CoreML对带批量维度的二维张量形状推断更稳定，将输入改为(1, 100)的二维张量，模型输出会对应为(1, 100)，转换后Xcode能正确识别为数组：

def main():
    guitar_net = TinyModel()
    guitar_net.eval()

    # 生成带批量维度的输入：shape=(1, 100)
    example_input = torch.rand(1, 100)

    summary(guitar_net, (1, 100))

    traced_model = torch.jit.trace(guitar_net, example_input)
    out = traced_model(example_input)

    model = ct.convert(
        model=traced_model,
        inputs=[ct.TensorType(shape=example_input.shape)],
    )

    model.save("TorchModel.mlmodel")

转换后在Swift中，输出会是形状为[1, 100]的MLMultiArray，只需取第一个维度的元素即可获取长度为100的概率数组，不会再触发内存错误。

验证转换结果 #

转换完成后，可通过以下代码确认输入输出形状是否正确：

print("输入信息:", model.input_description)
print("输出信息:", model.output_description)

内容的提问来源于stack exchange，提问作者Connor Hargus