[Web] Demucs model won't run in both WASM and WGPU

[gianlourbano]

Describe the issue

I converted the model from pytorch to onnx as described here, with some issues. The model works in onnx python, but in wasm /webgpu the runtime dies without error. The optimized version of the model runs in wasm, but not webgpu. I don't know if this problem is related to the model conversion or the runtime. I have tested with both @latest and @dev.

To reproduce

Here's a link to a sample repo, instructions in README.

Urgency

Urgent, as this project is related to my thesis

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

1.19.2, 1.20.0-dev.20240907-ad9afbb042

Execution Provider

'wasm'/'cpu' (WebAssembly CPU), 'webgpu' (WebGPU)

[gyagp]

For WebGPU EP, the problem is related to op unsqueeze. According the ONNX spec (https://onnx.ai/onnx/operators/onnx__Unsqueeze.html), axes of unsqueeze is a list of integers, but in your model, it's just a scalar "1".

[gianlourbano]

So the problem is related to the dynamo export of torch?

[fs-eire]

Technically the axes should always be a 1D tensor. However, in reality, the CPU code has loosen the limit:

https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/providers/cpu/tensor/unsqueeze.cc#L60-L62

perhaps webgpu should have same behavior to CPU.

#22054

[gianlourbano]

@gyagp with latest 1.20.0-dev.20240917-afd642a194, that should include both fixes, i still cannot run the model in webgpu, the runtime just aborts after displaying the wgpu experimental warning

[gyagp]

I also hit some issue with the latest code, and I will take a further look.
BTW, I manually modified the model to work around the unsqueeze issue before, and it seems that model can run. I uploaded it to https://huggingface.co/webai-community/models/tree/main (click "download file" after demucs.onnx).

[gianlourbano]

Your model succesfully runs with latest @dev, with timings (60s of audio with 10s chunks):

wasm:
step 0: 12656 ms
step 1: 12864 ms
step 2: 13211 ms
step 3: 13164 ms
step 4: 13643 ms
step 5: 13687 ms

wgpu:
step 0: 10226 ms
step 1: 9612 ms
step 2: 9628 ms
step 3: 9647 ms
step 4: 9600 ms
step 5: 9562 ms

onnx python cpu:
step 0: 4.9 s
step 1: 4.9 s
step 2: 4.6 s
step 3: 4.9 s
step 4: 4.8 s
step 5: 4.6 s

On ryzen 4600H

[gianlourbano]

I have also tried on a macbook m1 pro with an average wgpu step of ~2.8s