nato-opt 0.1.0

NATO Optimizer and Spectral Penalties (Rust Port)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

use tch::nn::VarStore;

/// Adjusts learning rates in optimizer's param groups in-place.
/// Note: tch doesn't natively expose per-group parameters with `lr_decay_epoch` in its Optimizer struct.
/// We provide a struct or a simple helper to emulate it, or we expect the user to manually adjust the LR on the optimizer object.
pub fn adjust_learning_rate(
    optimizer: &mut tch::nn::Optimizer,
    epoch: i64,
    initial_lr: f64,
    lr_decay_epoch: i64,
    lr_decay: f64,
) {
    if lr_decay_epoch > 0 {
        let new_lr = initial_lr * lr_decay.powi((epoch / lr_decay_epoch) as i32);
        optimizer.set_lr(new_lr);
    } else {
        optimizer.set_lr(initial_lr);
    }
}