将 PyTorch 模型与 CoreML 一起使用时输入尺寸重塑

Posted 2023-03-27

【中文标题】将 PyTorch 模型与 CoreML 一起使用时输入尺寸重塑

【英文标题】：Input dimension reshape when using PyTorch model with CoreML

【发布时间】：2020-09-03 02:42:38

【问题描述】：

我在 PyTorch 中有一个 seq2seq 模型，我想用 CoreML 运行它。将模型导出到 ONNX 时，输入尺寸固定为导出期间使用的张量的形状，并再次从 ONNX 转换为 CoreML。

import torch
from onnx_coreml import convert

x = torch.ones((32, 1, 1000))  # N x C x W
model = Model()
torch.onnx.export(model, x, 'example.onnx')

mlmodel = convert(model='example.onnx', minimum_ios_deployment_target='13')
mlmodel.save('example.mlmodel')

对于 ONNX 导出，您可以导出动态维度 -

torch.onnx.export(
    model, x, 'example.onnx',
    input_names = ['input'],
    output_names = ['output'],
    dynamic_axes=
        'input' : 0 : 'batch', 2: 'width',
        'output' : 0 : 'batch', 1: 'owidth',
    
)

但这会在转换为CoreML 时导致RunTimeWarning -

RuntimeWarning：您将无法在此 Core ML 模型上运行 predict()。潜在的异常消息是：错误编译模型：“编译器错误：找到大小为零的 Blob：

对于 CoreML 中的推理，我希望批量 (first) 和宽度 (last) 维度要么是动态的，要么能够静态更改它们。

这可能吗？

【参考方案1】：

通过为torch.onnx.export 指定dynamic_axes，可以在ONNX 中使输入的维度动态化。

torch.onnx.export(
    model,
    x,
    'example.onnx',
    # Assigning names to the inputs to reference in dynamic_axes
    # Your model only has one input: x
    input_names=["input"],
    # Define which dimensions should be dynamic
    # Names of the dimensions are optional, but recommended.
    # Could just be: "input": [0, 2]
    dynamic_axes="input": 0: "batch", 2: "width"
)

现在导出的模型接受大小为 [batch, 1, width] 的输入，其中 batch 和 width 是动态的。

【讨论】：

谢谢，这适用于 ONNX 转换，但 CoreML 对结果图不满意 - 我已经用详细信息更新了问题。

我对 CoreML 一无所知，但它可能不支持 ONNX 的所有内容。但从ONNX dynamic_axes cause compile time errors in final CoreML model #565 来看，动态维度应该可以工作。在使用动态尺寸时，这可能是您的模型的特定问题。