pub trait Module: Clone {
// Required method
fn visit(&mut self, p: &mut Parameters<'_>);
// Provided methods
fn parameters(&self) -> Vec<(String, Array)> { ... }
fn trainable_parameters(&self) -> Vec<(String, Array)> { ... }
fn update_trainable(&mut self, values: &[Array]) -> Result<()> { ... }
fn load_weights(
&mut self,
weights: impl IntoIterator<Item = (String, Array)>,
strict: bool,
) -> Result<()> { ... }
}Expand description
A model: a tree of named parameters, as MLX’s nn.Module.
A module implements Module::visit, reporting each parameter and submodule by name
in a fixed order; everything else (named parameters, trainable parameters, updates,
loading weights) is built on it. Models are cheap to clone: arrays are shared.
Required Methods§
Sourcefn visit(&mut self, p: &mut Parameters<'_>)
fn visit(&mut self, p: &mut Parameters<'_>)
Reports every parameter and submodule to p, frozen ones included.
Provided Methods§
Sourcefn parameters(&self) -> Vec<(String, Array)>
fn parameters(&self) -> Vec<(String, Array)>
Every parameter with its path (layers.0.weight), in visiting order.
Sourcefn trainable_parameters(&self) -> Vec<(String, Array)>
fn trainable_parameters(&self) -> Vec<(String, Array)>
The parameters that are trained (not frozen), with their paths.
Sourcefn update_trainable(&mut self, values: &[Array]) -> Result<()>
fn update_trainable(&mut self, values: &[Array]) -> Result<()>
Replaces the trainable parameters, given in the order of
Module::trainable_parameters (as optimizers return them).
Examples found in repository?
examples/mnist.rs (line 136)
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}Sourcefn load_weights(
&mut self,
weights: impl IntoIterator<Item = (String, Array)>,
strict: bool,
) -> Result<()>
fn load_weights( &mut self, weights: impl IntoIterator<Item = (String, Array)>, strict: bool, ) -> Result<()>
Replaces parameters by path, as MLX’s load_weights: each weight must have its
parameter’s shape (the dtype is kept as loaded). With strict, every parameter
must be given and no unknown names are allowed.
Dyn Compatibility§
This trait is not dyn compatible.
In older versions of Rust, dyn compatibility was called "object safety".
Implementations on Foreign Types§
Source§impl<M: Module> Module for Vec<M>
A sequence of modules is a module; its parts are named 0, 1, … (MLX’s lists).
impl<M: Module> Module for Vec<M>
A sequence of modules is a module; its parts are named 0, 1, … (MLX’s lists).