search

tracel-ai / burn

Burn is a new comprehensive dynamic Deep Learning Framework built using Rust with extreme flexibility, compute efficiency and portability as its primary goals.
https://burn.dev
Apache License 2.0
8.95k stars 444 forks source link

Support importing safetensors format #626

Open wdoppenberg opened 1 year ago

wdoppenberg commented 1 year ago

Safetensors (de)serialization

Feature description

Load, at either tensor- or module-level, weights from a .safetensors file.

Feature motivation

To be able to tap into the huggingface ecosystem it would be useful to have the ability to (de)serialize for safetensors. This would obviously have the limitation of these files merely containing a key-value JSON for all weights, and no description of the model's topology/architecture just like ONNX does. Ideally this would help grow the Model Zoo.

This would then mostly be a convenience feature, both for sharing and using existing models.

(Optional) Suggest a Solution

Possibly we could start by deriving the TryFrom trait for their TensorView and our Tensor and take it from there. Huggingface's candle (Rust) framework already has the capability to (de)serialize from .safetensors.

antimora commented 1 year ago

I think natively we should stick to the file format we have currently, primarily because it's difficult to maintain many formats. We already support JSON (for visual debugging), Bincode (for embedding data), and MessagePack (for long-term and cross-language support). However, we should support file conversion from other formats, such as safetensors, numpy, and pickle. The conversion code should be placed in burn-import under a feature flag included in the default feature set.

I also would like for us to build a tool to do basic structure conversion. What I mean by that is we need some ability to remap field names and also change the depth/level of field nodes.

jin-eld commented 6 months ago

Hi, is there any timeline for this feature? Given the large number of .safetensor models the lack of a conversion tool does hurt...

nathanielsimard commented 6 months ago

Not yet, but we may implement this sonner rather than later

antimora commented 6 months ago

@nathanielsimard it would be easy to do now that PyTorchFileRecorder works nicely. We just need to do the same things but using safetensors reader which is available from Candle.