burn_import/
lib.rs

1#![warn(missing_docs)]
2#![cfg_attr(docsrs, feature(doc_cfg))]
3#![allow(clippy::ptr_arg)]
4#![allow(clippy::single_match)]
5#![allow(clippy::upper_case_acronyms)]
6#![allow(clippy::approx_constant)]
7
8//! `burn-import` is a crate designed to simplify the process of importing models trained in other
9//! machine learning frameworks into the Burn framework. This tool generates a Rust source file that
10//! aligns the imported model with Burn's model and converts tensor data into a format compatible with
11//! Burn.
12
13#[cfg(any(feature = "pytorch", feature = "onnx", feature = "safetensors"))]
14#[macro_use]
15extern crate derive_new;
16
17// Enabled when the `pytorch` or `onnx` feature is enabled.
18#[cfg(any(feature = "pytorch", feature = "onnx"))]
19mod logger;
20
21/// The onnx module.
22#[cfg(feature = "onnx")]
23pub mod onnx;
24
25/// The module for generating the burn code.
26#[cfg(feature = "onnx")]
27pub mod burn;
28
29/// The PyTorch module for recorder.
30#[cfg(feature = "pytorch")]
31pub mod pytorch;
32
33/// The Safetensors module for recorder.
34#[cfg(feature = "safetensors")]
35pub mod safetensors;
36
37// Enabled when the `pytorch` or `safetensors` feature is enabled.
38#[cfg(any(feature = "pytorch", feature = "safetensors"))]
39mod common;
40
41mod formatter;
42pub use formatter::*;