Burn WGPU backend
This crate provides a WGPU backend for Burn utilizing the wgpu.
The backend supports Vulkan, Metal, DirectX11/12, OpenGL, WebGPU.
#[cfg(feature = "wgpu")]
mod wgpu {
use burn_autodiff::ADBackendDecorator;
use burn_wgpu::{AutoGraphicsApi, WgpuBackend, WgpuDevice};
use mnist::training;
pub fn run() {
let device = WgpuDevice::default();
training::run::<ADBackendDecorator<WgpuBackend<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.
The best value should be the smallest one that allows 100% GPU usage.
A high value might increase GPU memory usage with no benefit.