À propos du projet
Burn est une bibliothèque de tenseurs et un framework d'apprentissage profond écrit en Rust, destiné au calcul numérique, à l'entraînement et à l'inférence. Son objectif affiché est d'éviter la séparation habituelle entre une pile d'entraînement Python et un moteur de déploiement distinct : le même code de modèle peut être entraîné puis exécuté en production, ce que le projet présente comme utile pour la personnalisation sur appareil et l'apprentissage fédéré.
Conception et ergonomie
- API de type PyTorch avec formes et graphes dynamiques, tout en compilant JIT des flux d'opérations tensorielles et en effectuant une fusion automatique de kernels.
- La compilation incrémentale est mise en avant : le README affirme que les modifications du code du modèle se recompilent en moins de 5 secondes, même en mode release.
- Les backends peuvent être échangés sans modifier le code du modèle ; plusieurs backends peuvent coexister dans une même application et le device est choisi à l'exécution via `Device`.
Backends
- GPU : CUDA, ROCm, Metal, Vulkan, WebGPU et LibTorch (déprécié depuis la 0.22.0). La couverture varie selon le fournisseur : Nvidia (CUDA, Vulkan, WebGPU, LibTorch), AMD (ROCm, Vulkan, WebGPU, LibTorch), Apple (Metal, WebGPU, LibTorch), Intel et Qualcomm (Vulkan, WebGPU), Wasm (WebGPU).
- CPU : backend CPU CubeCL, Flex et LibTorch pour x86 et Arm ; Flex couvre aussi Wasm et no_std.
- Décorateurs de backend : Autodiff ajoute la rétropropagation à n'importe quel backend ; Fusion ajoute la fusion de kernels là où c'est pris en charge et est activé par défaut sur les backends accélérés first-party ; Remote (bêta) prend en charge l'exécution client/serveur via Iroh ou WebSocket pour le calcul distribué.
Entraînement et inférence
- Un tableau de bord d'entraînement en terminal construit sur Ratatui affiche les métriques d'entraînement et de validation en temps réel, avec navigation via les touches fléchées et la possibilité d'interrompre la boucle d'entraînement sans plantage.
- Les modèles ONNX peuvent être importés via burn-onnx et convertis en code Rust natif à l'aide des API Burn, afin de s'exécuter sur n'importe quel backend Burn ; le README indique que ce crate est en développement actif avec un ensemble d'opérateurs limité.
- Les poids aux formats PyTorch et Safetensors peuvent être chargés dans des modèles définis avec Burn.
- L'inférence peut s'exécuter dans le navigateur via WebAssembly avec Flex (CPU) ou WGPU (WebGPU), avec des démos MNIST et de classification d'images.
- Les composants principaux prennent en charge no_std pour les environnements embarqués bare-metal ; actuellement, seul le backend Flex fonctionne en no_std.
Écosystème
- CubeCL : langage et compilateur de calcul GPU derrière les backends accélérés, utilisable de manière autonome.
- burn-onnx : import ONNX ; burn-store : sauvegarde/chargement de poids et import PyTorch/Safetensors.
- burn-vision, burn-rl, burn-dataset pour la vision, l'apprentissage par renforcement et les jeux de données.
- Dépôt models avec des modèles pré-entraînés et des exemples ; burn-bench pour comparer les backends dans le temps.
- Une liste de crates communautaires couvre le chargement de données (polars, arrow-rs, image, hf-hub), la tokenisation et le NLP (tokenizers, rust-bert), les bibliothèques numériques (ndarray, nalgebra), le ML classique (linfa, smartcore), les runtimes d'inférence (candle, mistral.rs, ort, tract, wonnx), l'outillage LLM/RAG (rig, langchain-rust), les embeddings et la recherche vectorielle (fastembed, qdrant, lancedb), la vision par ordinateur (kornia-rs), la simulation (rapier) et la visualisation (rerun, plotters).
Pour commencer
- Le Burn Book est la documentation principale, couvrant les tenseurs, les modules, les optimiseurs et les kernels GPU personnalisés.
- Les exemples incluent un workflow MNIST de base, une boucle d'entraînement personnalisée, un kernel WGPU personnalisé, des jeux de données CSV et d'images, la régression, un renderer personnalisé, des démos d'inférence dans le navigateur, l'import de poids PyTorch, la classification et la génération de texte, et un WGAN sur MNIST.
- La licence est MIT/Apache-2.0.
Comments
0 Rating appears after 10 ratings
Sign in to join the discussion.