[Web] WebGPU incorrect output for Segment Anything Model (SAM)

[xenova]

Describe the issue

Attempting to run the original Segment Anything Models (SAM) results in warnings and incorrect output. cc @guschmue @fs-eire

Output

• Incorrect (WebGPU) output:
  

• Correct (WASM) output:
  

Warnings

An uncaught WebGPU validation error was raised: Tint WGSL reader failure: :47:20 error: index 3 out of bounds [0..2]
      inputIndices[3] = max(0, min(col, undefined - 1));
                   ^

 - While validating [ShaderModuleDescriptor]
 - While calling [Device].CreateShaderModule([ShaderModuleDescriptor]).

An uncaught WebGPU validation error was raised: [Invalid ShaderModule] is invalid.
 - While validating compute stage ([Invalid ShaderModule], entryPoint: main).
 - While calling [Device].CreateComputePipeline([ComputePipelineDescriptor]).

[Invalid ComputePipeline] is invalid.
 - While Validating GetBindGroupLayout (0) on [Invalid ComputePipeline]

To reproduce

Run the following code in a browser that supports WebGPU:

import 'https://cdn.jsdelivr.net/npm/[email protected]/dist/ort.webgpu.min.js'

ort.env.wasm.wasmPaths='https://cdn.jsdelivr.net/npm/[email protected]/dist/'

const options = {
    executionProviders: ['webgpu'],
    graphOptimizationLevel: 'disabled',
}


const path = './sam-vit-base/onnx/vision_encoder.onnx';
const session = await ort.InferenceSession.create(path, options);
console.log(session);

// Dummy inputs
const inputs = {
    pixel_values: new ort.Tensor('float32', new Float32Array(1 * 3 * 1024 * 1024).fill(2.248908281326294), [1, 3, 1024, 1024])
}

const output = await session.run(inputs);
console.log(output);

Urgency

high - Blocks WebGPU support in Transformers.js

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

1.16.1

Execution Provider

'webgpu' (WebGPU)

[guschmue]

We have this https://github.com/microsoft/onnxruntime-inference-examples/tree/main/js/segment-anything example and I did not run into it but I used a newer dev build and used samexporter.
Looking into it today.

[guschmue]

ah, the model inputs/outputs are quite a bit different than the models I use.
The fix from @qjia7 looks good, values for wasm and webgpu are close:

webgpu
[-0.032305262982845306, -0.03442259877920151, -0.04370711371302605, -0.044772855937480927, -0.04492305591702461,
wasm
[-0.032305341213941574, -0.03442276269197464, -0.043707218021154404, -0.04477302357554436, -0.04492315277457237,

waiting for CI to finish.

[xenova]

I've just tried 1.16.3, but the error still seems to be persist. Was the latest release meant to include this fix? cc @guschmue