在CoreML中使用PyTorch模型时,输入尺寸会改变形状

我在PyTorch中有一个想与CoreML一起运行的seq2seq模型。将模型导出到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'},
    }
)

But this leads to a RunTimeWarning when converting to CoreML -

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

为了在CoreML中进行推断,我希望批处理(第一个)和宽度(最后一个)维是动态的或具有静态更改它们的能力。

那可能吗?

评论
  • alber1986
    alber1986 回复

    The dimensions of the input can be made dynamic in ONNX by specifying dynamic_axes for 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]的输入,其中batch和width是动态的。