Support importing safetensors format (#2721)

* add support for safetensors in pytorch reader

* update book

* support safetensors format, for it's own crate

* remove some duplicate code and depend on pytorch adapter

* Fix clippy error

* Fix clippy errors

* Update doc strings

* Move common types/functions into common module

* Refactor and add adapter option

* Update doc strings

* Updated example for importing pt and safetensors weights

* Add book section on safetensors

* Mention Safetensors in the main README file

* Remove dead code

* SafeTensors -> Safetensors type

* SafeTensors -> Safetensors in docs/messages

* SafeTensors -> Safetensors in book

* SafeTensors -> Safetensors in example

* SafeTensors -> Safetensors in readme

* Removed TensorFlow

* Removed excess documentation

* Update pytorch-model to look the formatting similar to safetensors md

* Update pytorch-model.md

* Removed pytorch copied tests and replaced with one multilayer test

* Update README.md

Co-authored-by: Guillaume Lagrange <lagrange.guillaume.1@gmail.com>

* Update burn-book/src/import/README.md

Co-authored-by: Guillaume Lagrange <lagrange.guillaume.1@gmail.com>

* Update burn-book/src/import/README.md

Co-authored-by: Guillaume Lagrange <lagrange.guillaume.1@gmail.com>

* Shorter example

* Remove redundant comment

* Removed extra wording

* Replace unwrap_or_else with expect

* Removed imports for pytorch doc consistency

* Removed imports

* Fix formating

---------

Co-authored-by: Dilshod Tadjibaev <939125+antimora@users.noreply.github.com>
Co-authored-by: Guillaume Lagrange <lagrange.guillaume.1@gmail.com>

Author: 3 people

Cargo.lock

@@ -8,8 +8,8 @@ members = [
     "crates/*",
     "crates/burn-import/pytorch-tests",
     "crates/burn-import/onnx-tests",
+    "crates/burn-import/safetensors-tests",
     "examples/*",
-    "examples/pytorch-import/model",
     "xtask",
 ]
 

Cargo.toml

@@ -188,18 +188,18 @@ for more details.
 <div align="left">
 <img align="right" src="https://raw.github.com/tracel-ai/burn/main/assets/backend-chip.png" height="96px"/>
 
-
 Burn strives to be as fast as possible on as many hardwares as possible, with robust implementations.
 We believe this flexibility is crucial for modern needs where you may train your models in the cloud,
 then deploy on customer hardwares, which vary from user to user.
+
 </div>
 
 <br />
 
 **Supported Backends**
 
 | Backend  | Devices                      | Class       |
-| -------  | ---------------------------- | ----------- |
+| -------- | ---------------------------- | ----------- |
 | CUDA     | NVIDIA GPUs                  | First-Party |
 | ROCm     | AMD GPUs                     | First-Party |
 | Metal    | Apple GPUs                   | First-Party |
@@ -426,13 +426,16 @@ Our ONNX support is further described in
 
 <details>
 <summary>
-Importing PyTorch Models 🚚
+Importing PyTorch or Safetensors Models 🚚
 </summary>
 <br />
 
-Support for loading of PyTorch model weights into Burn’s native model architecture, ensuring
-seamless integration. See
-[Burn Book 🔥 section on importing PyTorch](https://burn.dev/burn-book/import/pytorch-model.html)
+You can load weights from PyTorch or Safetensors formats directly into your Burn-defined models. This makes it easy to reuse existing models while benefiting from Burn's performance and deployment features.
+
+Learn more:
+
+- [Import pre-trained PyTorch models into Burn](https://burn.dev/burn-book/import/pytorch-model.html)
+- [Load models from Safetensors format](https://burn.dev/burn-book/import/safetensors-model.html)
 
 </details>
 
@@ -468,24 +471,23 @@ means it can run in bare metal environment such as embedded devices without an o
 
 <br />
 
-
 ### Benchmarks
 
 To evaluate performance across different backends and track improvements over time, we provide a
 dedicated benchmarking suite.
 
 Run and compare benchmarks using [burn-bench](https://github.com/tracel-ai/burn-bench).
 
-
-> ⚠️ **Warning**  
-> When using one of the `wgpu` backends, you may encounter compilation errors related to recursive type evaluation. This is due to complex type nesting within the `wgpu` dependency chain.  
+> ⚠️ **Warning**
+> When using one of the `wgpu` backends, you may encounter compilation errors related to recursive type evaluation. This is due to complex type nesting within the `wgpu` dependency chain.
 > To resolve this issue, add the following line at the top of your `main.rs` or `lib.rs` file:
+>
 > ```rust
 > #![recursion_limit = "256"]
 > ```
+>
 > The default recursion limit (128) is often just below the required depth (typically 130-150) due to deeply nested associated types and trait bounds.
 
-
 ## Getting Started
 
 <div align="left">

README.md

@@ -24,6 +24,7 @@
 - [Import Models](./import/README.md)
   - [ONNX Model](./import/onnx-model.md)
   - [PyTorch Model](./import/pytorch-model.md)
+  - [Safetensors Model](./import/safetensors-model.md)
 - [Models & Pre-Trained Weights](./models-and-pretrained-weights.md)
 - [Quantization (Beta)](./quantization.md)
 - [Advanced](./advanced/README.md)

burn-book/src/SUMMARY.md

@@ -1,10 +1,13 @@
 # Importing Models
 
-The Burn project supports the import of models from various frameworks, emphasizing efficiency and
-compatibility. Currently, it handles two primary model formats:
+Burn supports importing models from other frameworks and file formats, enabling you to use pre-trained weights in your Burn applications.
 
-1. [ONNX](./onnx-model.md): Facilitates direct import, ensuring the model's performance and structure
-   are maintained.
+## Supported Formats
 
-2. [PyTorch](./pytorch-model.md): Enables the loading of PyTorch model weights into Burn’s native model
-   architecture, ensuring seamless integration.
+Burn currently supports three primary model import formats:
+
+| Format | Description | Use Case |
+|--------|-------------|----------|
+| [**ONNX**](./onnx-model.md) | Open Neural Network Exchange format | Direct import of complete model architectures and weights from any framework that supports ONNX export |
+| [**PyTorch**](./pytorch-model.md) | PyTorch weights (.pt, .pth) | Loading weights from PyTorch models into a matching Burn architecture |
+| [**Safetensors**](./safetensors-model.md) | Hugging Face's model serialization format | Loading a model's tensor weights into a matching Burn architecture |