simple_linear_train_script/
simple-linear-train-script.rs

1use hodu::{nn::SGD, prelude::*};
2use std::time::Instant;
3
4fn main() -> Result<(), Box<dyn std::error::Error>> {
5    let input_data: Vec<Vec<f32>> = (0..10000)
6        .map(|i| {
7            vec![
8                (i % 100) as f32 / 100.0,
9                ((i % 100) + 1) as f32 / 100.0,
10                ((i % 100) + 2) as f32 / 100.0,
11            ]
12        })
13        .collect();
14    let target_data: Vec<Vec<f32>> = (0..10000).map(|i| vec![((i % 100) * 10) as f32 / 1000.0]).collect();
15
16    let input_tensor = Tensor::new(input_data)?;
17    let target_tensor = Tensor::new(target_data)?;
18
19    // Build script using new API
20    let builder = Builder::new("linear_training".to_string());
21    builder.start()?;
22
23    let mut linear = Linear::new(3, 1, true, DType::F32)?;
24    let mse_loss = MSELoss::new();
25    let mut optimizer = SGD::new(0.01);
26
27    let input = Tensor::input("input", [10000, 3])?;
28    input.requires_grad()?;
29    let target = Tensor::input("target", [10000, 1])?;
30
31    let epochs = 1000;
32    let mut final_loss = Tensor::scalar(0.0)?;
33
34    for _ in 0..epochs {
35        let pred = linear.forward(&input)?;
36        let loss = mse_loss.forward((&pred, &target))?;
37
38        loss.backward()?;
39
40        optimizer.step(&mut linear.parameters())?;
41        optimizer.zero_grad(&mut linear.parameters())?;
42
43        final_loss = loss;
44    }
45
46    let params = linear.parameters();
47    builder.add_output("loss", final_loss)?;
48    builder.add_output("weight", params[0].clone())?;
49    builder.add_output("bias", params[1].clone())?;
50
51    builder.end()?;
52
53    // Build script
54    let mut script = builder.build()?;
55
56    // Set device
57    #[cfg(feature = "cuda")]
58    script.set_device(Device::CUDA(0));
59    #[cfg(feature = "metal")]
60    script.set_device(Device::Metal);
61
62    // Set compiler type
63    #[cfg(feature = "xla")]
64    script.set_compiler(Compiler::XLA);
65
66    // Set runtime inputs
67    script.set_input("input", input_tensor);
68    script.set_input("target", target_tensor);
69
70    println!("Compiling script...");
71    let compile_start = Instant::now();
72    script.compile()?;
73    let compile_elapsed = compile_start.elapsed();
74    println!("Compilation time: {:?}", compile_elapsed);
75
76    println!("Running script...");
77    let run_start = Instant::now();
78    let output = script.run()?;
79    let run_elapsed = run_start.elapsed();
80
81    println!("Loss: {}", output["loss"]);
82    println!("Weight: {}", output["weight"]);
83    println!("Bias: {}", output["bias"]);
84    println!("Execution time: {:?}", run_elapsed);
85    println!("Total time: {:?}", compile_elapsed + run_elapsed);
86
87    Ok(())
88}
simple_linear_train_script/simple-linear-train-script.rs

simple_linear_train_script/
simple-linear-train-script.rs
