將 PyTorch model 與 CoreML 一起使用時輸入維度重塑

Question

我在 PyTorch 中有一個 seq2seq model 我想用 CoreML 運行。 將 model 導出到 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 model 上運行 predict()。底層異常消息是：編譯 model 時出錯：“編譯器錯誤：找到大小為零的 Blob：

對於 CoreML 中的推理，我希望批處理（第一個）和寬度（最后一個）維度要么是動態的，要么能夠靜態地改變它們。

那可能嗎？

Answer 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"}}
)

現在導出的 model 接受大小為[batch, 1, width]的輸入，其中batch和width是動態的。