uform3 not performing as well as older onnx model (webgpu)

[mattdesl]

I'm currently testing with Chrome v137.0.7151.69 and macOS 13.0. I'm using onnxruntime-web@1.22.0-dev.20250409-89f8206ba4. I can't seem to run the new v3 model with webgpu onnnx EP without a huge performance degradation compared to the older model.

Code:
https://gist.github.com/mattdesl/30bc5de23eb6edfd7362d91d43170922

(change the "provider" and "model" variables in main.js)

I'm testing three models:

{
  // The "new" v3 model
  v3: {
    image_encoder: "uform3-image-text-english-small/image_encoder.onnx",
    text_encoder: "uform3-image-text-english-small/text_encoder.onnx",
  },
  // The "old" models ...
  fp16: {
    text_encoder: "uform-vl-english-small-gpu-fp16/text_encoder.onnx",
    image_encoder: "uform-vl-english-small-gpu-fp16/image_encoder.onnx",
  },
  fp32: {
    text_encoder: "uform-vl-english-small-cpu-fp32/text_encoder.onnx",
    image_encoder: "uform-vl-english-small-cpu-fp32/image_encoder.onnx",
  },
}

Using webgpu backend, testing only image encoding / inference time:
v3 ~7000 ms
fp16 ~800 ms
fp32 ~750 ms

The v3 model seems to produce inaccurate/incorrect cosine similarity in webgpu mode.

Using cpu backend:
v3 ~6500 ms
fp16 N/A
fp32 ~7000 ms

I am hoping it's just something I've done wrong that is causing the v3 webgpu to both fail to infer correctly and perform very slowly?

[mattdesl]

I've investigated this a little more and found that:

1. The conversion to u8 seems problematic performance-wise for onnxruntime-web in webgpu backend.
2. When going from fp32 to fp16, I had to use opset_version=20 to avoid incorrect inference results.

I had to pin these versions to get it running:

pip install onnxruntime==1.21.0 

# I also needed this to get things working for some reason..
pip install numpy==1.26.4 

I used my own conversion script based on your existing conversion notebook:

import onnx
import onnxruntime as rt
import os
import torch
from torch.onnx import export as onnx_export
from onnxconverter_common import float16
from onnx import helper
from onnxruntime.quantization import quantize_dynamic, QuantType

working_directory = 'src/clip_test/public'
model_name = "uform3-image-text-english-small"
model_directory = os.path.join(working_directory, model_name)
model_weights_path = os.path.join(model_directory, "image_encoder.pt")
config_path = os.path.join(model_directory, "config.json")
tokenizer_path = os.path.join(model_directory, "tokenizer.json")

print(model_weights_path)
state_dict = torch.load(model_weights_path, weights_only=True)

from uform.torch_encoders import ImageEncoder, TextEncoder
image_encoder = ImageEncoder.from_pretrained(config_path, state_dict)

image_encoder.eval()

model_out_fp16 = os.path.join(model_directory, 'image_encoder_fp16.onnx')

module = image_encoder
module.eval()
module.return_features = True
module.to(dtype=torch.float32)
model_out_0 = 'image_encoder_fp32.onnx'
torch.onnx.export(
    module,
    torch.randn(1, 3, 224, 224, dtype=torch.float32),  # Dummy input
    os.path.join(model_directory, model_out_0), 
    export_params=True,
    opset_version=20,
    do_constant_folding=True,
    input_names = ['images'], 
    output_names = ['embeddings'],
    dynamic_axes={
        'images' : {0 : 'batch_size'},
        'embeddings' : {0 : 'batch_size'}})

model_in_0_path = os.path.join(model_directory, model_out_0)

# convert to fp16
onnx_model = onnx.load(model_in_0_path)
onnx_model_fp16 = float16.convert_float_to_float16(onnx_model, 
    min_positive_val=1e-7,
    max_finite_val=1e2, 
    keep_io_types=False,     # optional, if you want inputs/outputs at fp32)
)
onnx.save(onnx_model_fp16, model_out_fp16)

Now with these results, I get two webgpu capable models (fp16 and fp32) that give correct inference and also much better performance than what is out of the box here. About ~175-225 ms inference time for the fp16 model.

[ashvardanian]

The model accuracy should have been fine with 8-bit integers, but maybe our data distribution was too different from yours 🤷‍♂️

As for throughput, that's more of a WebGPU issue. I can't remember a piece of hardware that would handle half-precision floats better than small integer, so it must be a software problem on their end.