IrError(Type(ParametrizedDimensionUnsupported("batch")))

[Ryul0rd]

Describe the bug Exporting a HuggingFace model using the recommended method results in the following error: thread 'main' panicked at 'called 'Result::unwrap()' on an 'Err' value: IrError(Type(ParametrizedDimensionUnsupported("batch")))' The inclusion of the batch dimension is not only what HuggingFace tool does but also what the official PyTorch docs recommend for exporting to onnx.

To Reproduce

1. pip install transformers[onnx]
2. python -m transformers.onnx --model=bert-base-uncased --feature=default onnx/

fn main() {
    #[cfg(not(target_arch = "wasm32"))]
    {
        pollster::block_on(run());
    }
}

async fn run () {
    let model_path = Path::new("onnx/model.onnx");
    let _session = wonnx::Session::from_path(model_path).await.unwrap();
}

Expected behavior The unwrap call should not encounter an error.

Desktop Linux PopOS 20.04

[haixuanTao]

Unfortunately, We don't support parametrized dimension yet. You can replace those dimensions using onnx-simplifier. There should be an example in the documentation.

[Ryul0rd]

I've also tried using onnx-simplifier as described in the documentation and get the same error. Is there something I need to set to have onnx-simplifier remove a batch dimension?

[haixuanTao]

Ok, maybe onnx-simplifier has changed since last time I have used it. Basically in the past you could tell onnx-simplifier to replace batch with a static number like 1.

To make WebGPU work we need to have fixed sized buffers. Therefore we need to know in advance the size of each tensors, that's why we are not currently managing parametrized dimensions.

Later on we could prob make this dynamic.

[pixelspark]

@Ryul0rd the command should be something along the lines of:

python -m onnxsim ./original.onnx ./original-inferred.onnx --input-shape "batch:1"

If you still receive an error afterwards, check if it is for another parameter (other than batch which after running the above command should not appear anymore). You might also want to use a tool such as Netron (https://netron.app) to see what's going on (it will show parametrized dimensions).

[Ryul0rd]

@pixelspark Thanks for pointing me in the right direction. The input shape argument is deprecated but you can currently do python -m onnxsim onnx/model.onnx onnx/simp-model.onnx --overwrite-input-shape "input_ids:1,512" "attention_mask:1,512" "token_type_ids:1,512" In the case of Huggingface's BERT. This fixes the input but unfortunately, the output remains dynamically sized and there doesn't seem to be a way for onnx simplifier to handle this at the moment. I suppose the workaround is to just DIY an export script using PyTorch since their export gives control over I/O shapes without needing to use onnx simplifier.

[superbattery]

@pixelspark Thanks for pointing me in the right direction. The input shape argument is deprecated but you can currently do python -m onnxsim onnx/model.onnx onnx/simp-model.onnx --overwrite-input-shape "input_ids:1,512" "attention_mask:1,512" "token_type_ids:1,512" In the case of Huggingface's BERT. This fixes the input but unfortunately, the output remains dynamically sized and there doesn't seem to be a way for onnx simplifier to handle this at the moment. I suppose the workaround is to just DIY an export script using PyTorch since their export gives control over I/O shapes without needing to use onnx simplifier.

@Ryul0rd onnxruntime supports both the input and output shape. try this: python -m onnxruntime.tools.make_dynamic_shape_fixed --dim_param batch --dim_value 1 orig.onnx result.onnx

[pixelspark]

You can now also use nnx prepare model.onnx model-prepared.onnx --set some_dim=123 (if you have a fully inferred model that is)