Burn Deep Learning Framework

repository·main·Indexed 12 days ago

https://github.com/tracel-ai/burn

A high-performance deep learning framework and tensor library in Rust that unifies training and inference through a single, multi-platform API. It includes support for automatic differentiation via Burn Autodiff, multiple backends including CUDA, Wgpu, and CubeCL, and tools for kernel fusion and distributed tensor sharing.

Tokens
159.7K
Snippets
434
Records
668
Agent score
94%

What's inside Burn

  1. Overview of the Burn Standard Library (burn-std)

    main
    The burn-std crate provides the foundational core types and utilities used throughout the Burn ecosystem. It is designed to be lightweight and highly portable, offering essential definitions for shapes, indexing, and various data types required for deep learning operations.
  2. Overview of the Burn Backend crate

    main
    The burn-backend crate provides the core interfaces and data structures required to execute tensor operations within the Burn ecosystem. It defines the abstraction layer that allows Burn's high-level tensor API to interact with various hardware-specific implementations (backends).
  3. Overview of Burn Dataset

    main
    The Burn Dataset library is a specialized tool for streamlining machine learning (ML) data pipeline creation. It provides various dataset implementations, transformation functions, and multiple data sources to facilitate efficient data handling during training and inference.
  4. Overview of burn-flex capabilities and performance

    main

    burn-flex is a complete replacement for burn-ndarray designed to overcome limitations such as the 6-dimension limit and lack of native f16/bf16 support.

    Key Features

    • Trait Implementation: Implements all required Backend traits: FloatTensorOps, IntTensorOps, BoolTensorOps, QTensorOps, ModuleOps, ActivationOps, and TransactionOps.
    • Memory Efficiency: Uses Arc-based Copy-On-Write (COW) and in-place mutation to reduce memory usage.
    • Performance: Achieves compute improvements (1.1-9.7x) via better algorithms and massive structural improvements (up to 166,000x) by representing operations as zero-copy views instead of eagerly materializing data.

    Limitations

    • Does not support optional BLAS acceleration (uses the gemm crate instead).
    • Does not include the export_tests reference implementation feature.
  5. Overview of Burn Store

    main

    Burn Store is an advanced model storage and serialization library for the Burn deep learning framework. It provides tools for efficient model serialization, cross-framework interoperability (such as PyTorch support), and advanced tensor management.

    Key features include:

    • Burnpack Format: Native Burn format using CBOR metadata, supporting memory-mapped loading, ParamId persistence for stateful training, and no-std environments.
    • SafeTensors Format: Industry-standard secure and efficient tensor serialization.
    • PyTorch Support: Direct loading of .pth or .pt files with automatic weight transformation.
    • Zero-Copy Loading: Utilizes memory-mapped files and lazy tensor materialization.
    • Flexible Filtering & Remapping: Ability to load/save specific model subsets using regex or custom predicates, and rename tensors during the process.
    • Half-Precision Storage: Automatic conversion (e.g., F32 to F16) to reduce file sizes.
    • No-std Support: Available for embedded and WASM environments via Burnpack and SafeTensors.
  6. Overview of the Burn CubeCL Backend

    main
    The burn-cubecl backend is a generic backend for the Burn framework that supports Just-In-Time (JIT) compilation to any shader language target. This allows for flexible execution across different hardware by compiling kernels to the appropriate shader language at runtime.
  7. What is Burn?

    main

    Burn is a unified deep learning framework and tensor library written in Rust, optimized for numerical computing, training, and inference.

    Unlike traditional workflows where models are trained in Python and exported to separate production engines (like ONNX or TensorRT), Burn allows you to use the exact same codebase for both training and production. This eliminates the brittle and lossy export step, making it ideal for complex architectures and deployment scenarios like on-device personalization or federated learning.

    Key features include:

    • Unified API: Execute multi-platform tensor operations with a single API.
    • PyTorch-like Ergonomics: Supports dynamic shapes and dynamic graphs.
    • High Performance: Uses JIT-compilation for streams of tensor operations to perform automatic kernel fusion.
    • Fast Research Loop: Designed for incremental compilation, allowing model code modifications to recompile in under 5 seconds (even in release mode).
  8. Overview of burn-flex

    main

    burn-flex is a high-performance, memory-efficient CPU backend for the Burn framework. It is designed to support a wide range of environments, including std, no_std (embedded), and WebAssembly (WASM).

    Key characteristics include:

    • Performance: Optimized via SIMD (via macerator), multi-threading (via rayon), and highly optimized matrix multiplication (via gemm).
    • Memory Efficiency: Extensive use of zero-copy operations (e.g., transpose, slice, expand) and Arc-based Copy-on-Write (COW) semantics for $O(1)$ tensor cloning.
    • Versatility: Supports various data types (f32, f64, f16, bf16, and integers i8-i64, u8-u64) and specialized operations like quantized arithmetic, FFT, and attention mechanisms.