将 PyTorch 模型与 CoreML 结合使用时输入维度重塑

Question

我在 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'},
    }
)

但这会导致RunTimeWarning转换为CoreML-

RuntimeWarning：您将无法在此 Core ML 模型上运行 Predict()。底层异常消息是：编译模型时出错：“编译器错误：发现大小为零的 Blob：

为了在 CoreML 中进行推理，我希望批量（第一个）和宽度（最后一个）维度要么是动态的，要么能够静态地更改它们。

那可能吗？

Answer 1

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

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]的输入，其中批量和宽度是动态的。