Burn WGPU 后端

Burn WGPU 后端

此包为 Burn 提供了一个使用 wgpu 的 WGPU 后端。

后端支持 Vulkan、Metal、DirectX11/12、OpenGL、WebGPU。

使用示例

#[cfg(feature = "wgpu")]
mod wgpu {
    use burn_autodiff::Autodiff;
    use burn_wgpu::{AutoGraphicsApi, Wgpu, WgpuDevice};
    use mnist::training;

    pub fn run() {
        let device = WgpuDevice::default();
        training::run::<Autodiff<Wgpu<AutoGraphicsApi, f32, i32>>>(device);
    }
}

配置

您可以将 BURN_WGPU_MAX_TASKS 设置为正整数，以确定向图形 API 提交批处理计算任务的次数。

平台支持