Burn WGPU backend
This crate provides a WGPU backend for Burn using the wgpu.
The backend supports 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);
}
}
You can set BURN_WGPU_MAX_TASKS
to a positive integer that determines how many computing tasks are submitted in batches to the graphics API.
Option | CPU | GPU | Linux | MacOS | Windows | Android | iOS | WASM |
---|---|---|---|---|---|---|---|---|
Metal | No | Yes | No | Yes | No | No | Yes | No |
Vulkan | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No |
OpenGL | No | Yes | Yes | Yes | Yes | Yes | Yes | No |
WebGpu | No | Yes | No | No | No | No | No | Yes |
Dx11/Dx12 | No | Yes | No | No | Yes | No | No | No |