burn_optim/lib.rs
1#![cfg_attr(not(feature = "std"), no_std)]
2#![warn(missing_docs)]
3#![cfg_attr(docsrs, feature(doc_cfg))]
4
5//! Optimizers and learning rate schedulers for Burn.
6//!
7//! Applications use these through `burn::optim`, `burn::lr_scheduler` and `burn::grad_clipping`.
8//! An optimizer is built from its config and applied to a module with gradients from a backward pass:
9//!
10//! ```rust,ignore
11//! let mut optimizer = AdamConfig::new().init();
12//! let grads = GradientsParams::from_grads(loss.backward(), &model);
13//! model = optimizer.step(learning_rate, model, grads);
14//! ```
15//!
16//! - Optimizers: SGD, Adam, AdamW, Adagrad, Adafactor, Adan, LAMB, L-BFGS, Lion, Muon and RMSprop,
17//! with momentum, weight decay and gradient accumulation helpers.
18//! - [`lr_scheduler`]: constant, step, exponential, linear, cosine, Noam, and sequential or
19//! composed schedules.
20//! - [`grad_clipping`]: clipping by value or by norm.
21//!
22//! # Feature flags
23//!
24//! - `std` (default): standard library support. Without it the crate is `no_std` with `alloc`.
25//! - `tracing`: instrument operations with the `tracing` crate.
26
27#[macro_use]
28extern crate derive_new;
29
30extern crate alloc;
31
32/// Optimizer module.
33mod optim;
34pub use optim::*;
35
36/// Gradient clipping module.
37pub mod grad_clipping;
38
39/// Learning rate scheduler module.
40#[cfg(feature = "std")]
41pub mod lr_scheduler;
42
43/// Type alias for the learning rate.
44///
45/// LearningRate also implements [learning rate scheduler](crate::lr_scheduler::LrScheduler) so it
46/// can be used for constant learning rate.
47pub type LearningRate = f64; // We could potentially change the type.