Struct Sgd 

pub struct Sgd {
    pub learning_rate: LearningRate,
    pub momentum: f64,
    pub weight_decay: f64,
    pub dampening: f64,
    pub nesterov: bool,
    /* private fields */
}

Stochastic gradient descent with optional momentum, dampening, Nesterov momentum and weight decay, as MLX’s optimizers.SGD.

Fields

learning_rate: LearningRate

The learning rate.

momentum: f64

The momentum strength (0 disables momentum).

weight_decay: f64

L2 weight decay added to the gradient.

dampening: f64

Dampening for momentum.

nesterov: bool

Nesterov momentum.

Implementations

impl Sgd

pub fn new(learning_rate: impl Into<LearningRate>) -> Self

Plain SGD.

Examples found in repository

examples/mnist.rs (line 130)

fn main() -> Result<()> {
    let args: Vec<String> = std::env::args().skip(1).collect();
    let dir = args
        .iter()
        .find(|a| !a.starts_with("--"))
        .cloned()
        .unwrap_or_else(|| "data/mnist".into());
    let epochs: usize = args
        .iter()
        .position(|a| a == "--epochs")
        .and_then(|i| args.get(i + 1))
        .map_or(10, |v| v.parse().expect("--epochs takes a number"));
    let asynchronous = args.iter().any(|a| a == "--async");
    let compiled = args.iter().any(|a| a == "--compile");
    if args.iter().any(|a| a == "--cpu") {
        set_default_stream(Stream::default_for(Device::cpu()));
    }
    let (train_x, train_y, n) = load(&dir, "train").inspect_err(|_| {
        eprintln!("cannot read MNIST from {dir}: run `just mnist` (it downloads the data)");
    })?;
    let (test_x, test_y, n_test) = load(&dir, "t10k")?;
    let pixels = train_x.len() / n;
    let test_x = Array::from_slice(&test_x, &[n_test, pixels])?;
    let test_y = Array::from_slice(&test_y, &[n_test])?;

    let (batch, lr) = (256, 0.1);
    random::seed(0);
    let mut model = Mlp::new(&[pixels, 32, 32, 10])?;
    nuvec::eval(&arrays(&model))?;
    let mut optimizer = Sgd::new(lr);
    // The step as a function of the parameters and the batch, returning new parameters.
    let k = arrays(&model).len();
    let template = model.clone();
    let step = nuvec::compile(move |xs: &[Array]| {
        let mut m = template.clone();
        m.update_trainable(&xs[..k])?;
        let (x, y) = (&xs[k], &xs[k + 1]);
        let (_, grads) = value_and_grad(&m, |m: &Mlp| {
            cross_entropy(&m.forward(x)?, y, Reduction::Mean)
        })?;
        Sgd::new(lr).update(&mut m, &grads)?;
        Ok(arrays(&m))
    });
    let mut shuffle = Shuffle(0);
    let device = nuvec::default_stream().device;
    let mut times = Vec::new();
    let mut accuracy = 0.0;
    for epoch in 0..epochs {
        let tic = Instant::now();
        let perm = shuffle.permutation(n);
        for ids in perm.chunks(batch) {
            let mut xb = Vec::with_capacity(ids.len() * pixels);
            for &i in ids {
                xb.extend_from_slice(&train_x[i * pixels..(i + 1) * pixels]);
            }
            let yb: Vec<u32> = ids.iter().map(|&i| train_y[i]).collect();
            let x = Array::from_slice(&xb, &[ids.len(), pixels])?;
            let y = Array::from_slice(&yb, &[ids.len()])?;
            if compiled {
                let mut args = arrays(&model);
                args.extend([x, y]);
                model.update_trainable(&step(&args)?)?;
            } else {
                let (_, grads) = value_and_grad(&model, |m: &Mlp| {
                    cross_entropy(&m.forward(&x)?, &y, Reduction::Mean)
                })?;
                optimizer.update(&mut model, &grads)?;
            }
            if asynchronous {
                nuvec::async_eval(&arrays(&model))?;
            } else {
                nuvec::eval(&arrays(&model))?;
            }
        }
        let predictions = argmax(&model.forward(&test_x)?, Some(1), false)?;
        let hits = astype(&equal(&predictions, &test_y)?, DType::F32)?;
        accuracy = mean(&hits, &[], false)?.item::<f32>()?;
        let seconds = tic.elapsed().as_secs_f64();
        times.push(seconds);
        println!("Epoch {epoch}: Test accuracy {accuracy:.3}, Time {seconds:.3} (s)");
    }
    let mut steady: Vec<f64> = times.iter().skip(1).copied().collect();
    steady.sort_by(f64::total_cmp);
    let median = steady.get(steady.len() / 2).copied().unwrap_or(times[0]);
    println!(
        "RESULT device={device} async={asynchronous} compiled={compiled} accuracy={accuracy:.4} epoch_seconds={median:.4}"
    );
    Ok(())
}

pub fn momentum(self, momentum: f64) -> Self

Sets the momentum.

pub fn weight_decay(self, weight_decay: f64) -> Self

Sets the weight decay.

pub fn nesterov(self, nesterov: bool) -> Self

Enables Nesterov momentum.

Trait Implementations

impl Clone for Sgd

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for Sgd

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Optimizer for Sgd

fn apply( &mut self, parameters: &[Array], gradients: &[Array], ) -> Result<Vec<Array>>

One step: the updated parameters.

fn update<M: Module>( &mut self, model: &mut M, gradients: &[Array], ) -> Result<()>

Applies one step to model: gradients are in the order of Module::trainable_parameters.