burn-import 0.20.1

Library for importing datamodels into the Burn framework
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

#[cfg(feature = "onnx")]
use burn_import::onnx::ModelGen;

#[cfg(feature = "onnx")]
/// Takes an ONNX file and generates a model from it
fn main() {
    let onnx_file = std::env::args().nth(1).expect("No input file provided");
    let output_dir = std::env::args()
        .nth(2)
        .expect("No output directory provided");

    // Generate the model code from the ONNX file.
    // Weights are saved in burnpack format (.bpk file alongside the generated code)
    ModelGen::new()
        .input(onnx_file.as_str())
        .development(true)
        .out_dir(output_dir.as_str())
        .run_from_cli();
}

#[cfg(not(feature = "onnx"))]
fn main() {
    println!("Compiled without Pytorch feature.");
}