1use flodl::*;
9use flodl::monitor::Monitor;
10
11fn main() -> Result<()> {
12 let model = FlowBuilder::from(Linear::new(2, 16)?)
14 .through(GELU)
15 .through(LayerNorm::new(16)?)
16 .also(Linear::new(16, 16)?)
17 .through(Linear::new(16, 2)?)
18 .build()?;
19
20 let params = model.parameters();
22 let mut optimizer = Adam::new(¶ms, 0.01);
23 model.train();
24
25 let opts = TensorOptions::default();
27 let batches: Vec<(Tensor, Tensor)> = (0..32)
28 .map(|_| {
29 let x = Tensor::randn(&[16, 2], opts).unwrap();
30 let y = Tensor::randn(&[16, 2], opts).unwrap();
31 (x, y)
32 })
33 .collect();
34
35 let num_epochs = 50;
37 let mut monitor = Monitor::new(num_epochs);
38 for epoch in 0..num_epochs {
42 let t = std::time::Instant::now();
43
44 for (input_t, target_t) in &batches {
45 let input = Variable::new(input_t.clone(), true);
46 let target = Variable::new(target_t.clone(), false);
47
48 let pred = model.forward(&input)?;
49 let loss = mse_loss(&pred, &target)?;
50
51 optimizer.zero_grad();
52 loss.backward()?;
53 clip_grad_norm(¶ms, 1.0)?;
54 optimizer.step()?;
55
56 model.record_scalar("loss", loss.item()?);
57 }
58
59 model.flush(&[]);
60 monitor.log(epoch, t.elapsed(), &model);
61 }
62
63 monitor.finish();
64 Ok(())
65}