使用Burn/Wgpu在WASM运行MobileNetV2时触发Condvar Wait异常求助

问题描述 #

使用Rust、WASM和Burn框架开发浏览器端图像分类模型，基于Burn-ONNX的图像分类示例项目开发。使用自带的Squeezenet1模型时WebGPU后端运行正常，但切换为MobileNetV2模型后，NdArray后端可正常运行，WebGPU后端触发panic。

报错信息 #

/rustc/4a4ef493e3a1488c6e321570238084b38948f6db/library/std/src/sys/sync/condvar/no_threads.rs:20:9:

condvar wait not supported
gnu/lib/rustlib/src/rust/library/std/src/sys/sync/mutex/no_threads.rs:19:9:

assertion left == right failed: cannot recursively acquire mutex

left: true

right: false

触发panic的代码 #

pub async fn set_backend_wgpu(&mut self) -> Result<(), JsValue> {
        log::info!("Loading the model to the Wgpu backend");
        let start = Instant::now();
        let device = WgpuDevice::default();

        if let Ok(_) =
            WGPU_INITIALIZED.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
        {
            init_setup_async::<AutoGraphicsApi>(&device, Default::default()).await;
        }
        self.model = ModelType::WithWgpuBackend(Model::new(&device));
        let duration = start.elapsed();
        log::debug!("Model is loaded to the Wgpu backend in {duration:?}");

        log::debug!("Warming up the model");
        let start = Instant::now();
        let _ = self.inference(&[0.0; HEIGHT * WIDTH * CHANNELS]).await;
        let duration = start.elapsed();
        log::debug!("Warming up is completed in {duration:?}");
        Ok(())
    }

核心结论与分析 #

不是MobileNetV2的架构本身不支持WASM+WebGPU环境，问题根源在于WASM单线程环境对Rust同步原语的限制，以及MobileNetV2复杂结构触发了Burn WebGPU后端的相关逻辑：

• 浏览器默认的WASM环境是单线程的，Rust标准库中的condvar（条件变量）和递归互斥锁在单线程目标下并不被支持，因为这类同步原语依赖多线程调度机制。
• MobileNetV2的模型结构比Squeezenet更复杂，在初始化或热身推理阶段，Burn的WebGPU后端内部触发了需要多线程同步的逻辑，而Squeezenet的推理路径未触及这部分代码。
• 当前代码中，init_setup_async与inference的连续调用，在处理MobileNetV2时可能导致后端内部尝试递归获取互斥锁，或调用了不支持的条件变量等待操作。

解决方向 #

1. 验证Burn版本兼容性：检查使用的Burn版本是否已修复WASM单线程环境下的同步原语问题，优先使用最新稳定版。
2. 配置单线程安全路径：查看Burn WebGPU后端的配置选项，是否存在禁用多线程相关逻辑的开关，强制使用单线程安全的代码路径。
3. 调整模型加载逻辑：拆分模型初始化与热身推理的步骤，避免在同一个异步块中连续触发可能引发锁竞争的操作，例如调整init_setup_async与Model::new的调用顺序，或增加异步等待间隔。
4. 启用WASM多线程：若浏览器支持（需开启跨源隔离和SharedArrayBuffer），可在Cargo配置中启用多线程相关特性（如features = ["parallel"]），让Rust标准库的同步原语正常工作。

内容的提问来源于stack exchange，提问作者ddhello