Skip to main content

Optimizer

nuvec_nn::optim

Trait Optimizer 

Source 
pub trait Optimizer {
    // Required method
    fn apply(
        &mut self,
        parameters: &[Array],
        gradients: &[Array],
    ) -> Result<Vec<Array>>;

    // Provided method
    fn update<M: Module>(
        &mut self,
        model: &mut M,
        gradients: &[Array],
    ) -> Result<()> { ... }
}
Expand description

Updates a model’s trainable parameters from their gradients.

Required Methods§

Source

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

One step: the updated parameters.

Provided Methods§

Source

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.

Examples found in repository?
examples/mnist.rs (line 141)
101fn main() -> Result<()> {
102    let args: Vec<String> = std::env::args().skip(1).collect();
103    let dir = args
104        .iter()
105        .find(|a| !a.starts_with("--"))
106        .cloned()
107        .unwrap_or_else(|| "data/mnist".into());
108    let epochs: usize = args
109        .iter()
110        .position(|a| a == "--epochs")
111        .and_then(|i| args.get(i + 1))
112        .map_or(10, |v| v.parse().expect("--epochs takes a number"));
113    let asynchronous = args.iter().any(|a| a == "--async");
114    let compiled = args.iter().any(|a| a == "--compile");
115    if args.iter().any(|a| a == "--cpu") {
116        set_default_stream(Stream::default_for(Device::cpu()));
117    }
118    let (train_x, train_y, n) = load(&dir, "train").inspect_err(|_| {
119        eprintln!("cannot read MNIST from {dir}: run `just mnist` (it downloads the data)");
120    })?;
121    let (test_x, test_y, n_test) = load(&dir, "t10k")?;
122    let pixels = train_x.len() / n;
123    let test_x = Array::from_slice(&test_x, &[n_test, pixels])?;
124    let test_y = Array::from_slice(&test_y, &[n_test])?;
125
126    let (batch, lr) = (256, 0.1);
127    random::seed(0);
128    let mut model = Mlp::new(&[pixels, 32, 32, 10])?;
129    nuvec::eval(&arrays(&model))?;
130    let mut optimizer = Sgd::new(lr);
131    // The step as a function of the parameters and the batch, returning new parameters.
132    let k = arrays(&model).len();
133    let template = model.clone();
134    let step = nuvec::compile(move |xs: &[Array]| {
135        let mut m = template.clone();
136        m.update_trainable(&xs[..k])?;
137        let (x, y) = (&xs[k], &xs[k + 1]);
138        let (_, grads) = value_and_grad(&m, |m: &Mlp| {
139            cross_entropy(&m.forward(x)?, y, Reduction::Mean)
140        })?;
141        Sgd::new(lr).update(&mut m, &grads)?;
142        Ok(arrays(&m))
143    });
144    let mut shuffle = Shuffle(0);
145    let device = nuvec::default_stream().device;
146    let mut times = Vec::new();
147    let mut accuracy = 0.0;
148    for epoch in 0..epochs {
149        let tic = Instant::now();
150        let perm = shuffle.permutation(n);
151        for ids in perm.chunks(batch) {
152            let mut xb = Vec::with_capacity(ids.len() * pixels);
153            for &i in ids {
154                xb.extend_from_slice(&train_x[i * pixels..(i + 1) * pixels]);
155            }
156            let yb: Vec<u32> = ids.iter().map(|&i| train_y[i]).collect();
157            let x = Array::from_slice(&xb, &[ids.len(), pixels])?;
158            let y = Array::from_slice(&yb, &[ids.len()])?;
159            if compiled {
160                let mut args = arrays(&model);
161                args.extend([x, y]);
162                model.update_trainable(&step(&args)?)?;
163            } else {
164                let (_, grads) = value_and_grad(&model, |m: &Mlp| {
165                    cross_entropy(&m.forward(&x)?, &y, Reduction::Mean)
166                })?;
167                optimizer.update(&mut model, &grads)?;
168            }
169            if asynchronous {
170                nuvec::async_eval(&arrays(&model))?;
171            } else {
172                nuvec::eval(&arrays(&model))?;
173            }
174        }
175        let predictions = argmax(&model.forward(&test_x)?, Some(1), false)?;
176        let hits = astype(&equal(&predictions, &test_y)?, DType::F32)?;
177        accuracy = mean(&hits, &[], false)?.item::<f32>()?;
178        let seconds = tic.elapsed().as_secs_f64();
179        times.push(seconds);
180        println!("Epoch {epoch}: Test accuracy {accuracy:.3}, Time {seconds:.3} (s)");
181    }
182    let mut steady: Vec<f64> = times.iter().skip(1).copied().collect();
183    steady.sort_by(f64::total_cmp);
184    let median = steady.get(steady.len() / 2).copied().unwrap_or(times[0]);
185    println!(
186        "RESULT device={device} async={asynchronous} compiled={compiled} accuracy={accuracy:.4} epoch_seconds={median:.4}"
187    );
188    Ok(())
189}

Dyn Compatibility§

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors§