minidx_core/
optimizers.rs

1use crate::misc::Decay;
2use crate::{Dtype, Float, Gradients, Unit};
3use num_traits::FromPrimitive;
4use serde::{Deserialize, Serialize};
5
6/// Describes the training parameters at some instant, serialized during recording.
7#[derive(Clone, Serialize, Deserialize, Debug)]
8pub struct TrainInfo {
9    /// The learning rate used this update step.
10    pub lr: f32,
11    /// The L1 regularization applied each update step.
12    pub l1_reg: f32,
13    /// The L2 regularization applied each update step.
14    pub l2_reg: f32,
15    /// The maximum magnitude of any update to any parameter.
16    pub grad_clip: Option<f32>,
17    /// The 0-indexed count of update steps.
18    pub step: usize,
19}
20
21/// An object responsible for tweaking gradient updates, such as to
22/// add the effects of momentum, perform gradient clipping, etc.
23pub trait GradAdjuster<G: Gradients> {
24    fn adjust(&mut self, gradient_updates: G, loss: f32) -> G;
25}
26
27/// An object responsible for applying gradient updates, such as to
28/// add the gradient updates to the weights while performing regularization.
29pub trait GradApplyer {
30    fn apply<G: Gradients>(
31        &mut self,
32        gradient_updates: G,
33        weights: &mut G,
34    ) -> Result<(), crate::Error>;
35
36    fn advance_step(&mut self);
37}
38
39/// Describes the basic set of parameters used in training. Implements
40/// optimizer traits, so it can be passed into training methods.
41#[derive(Clone, Serialize, Deserialize, Debug)]
42pub struct TrainParams {
43    /// The learning rate to use for update steps. See [Decay] for
44    /// how to specify a fixed vs variable learning rate.
45    pub lr: Decay,
46    /// The L1 regularization to apply each update step.
47    pub l1_reg: Option<Decay>,
48    /// The L2 regularization to apply each update step.
49    pub l2_reg: Option<Decay>,
50
51    /// If set, the number of update steps for which the learning rate
52    /// is ramped up at the beginning.
53    pub soft_start_epochs: Option<usize>,
54
55    /// If set, the maximum magnitude of any update to any parameter.
56    pub grad_clip: Option<f32>,
57
58    step: usize,
59}
60
61impl Default for TrainParams {
62    fn default() -> Self {
63        Self {
64            lr: Decay::None(1.0e-6),
65            l1_reg: None,
66            l2_reg: None,
67            soft_start_epochs: None,
68            grad_clip: None,
69            step: 0,
70        }
71    }
72}
73
74impl Into<TrainInfo> for &TrainParams {
75    fn into(self) -> TrainInfo {
76        TrainInfo {
77            lr: self.current_lr(),
78            l1_reg: self
79                .l1_reg
80                .as_ref()
81                .map(|d| d.at_timestep(self.step))
82                .unwrap_or(0.0),
83            l2_reg: self
84                .l2_reg
85                .as_ref()
86                .map(|d| d.at_timestep(self.step))
87                .unwrap_or(0.0),
88            grad_clip: self.grad_clip,
89            step: self.step,
90        }
91    }
92}
93
94impl TrainParams {
95    /// Constructs an object with the given training rate.
96    pub fn with_lr(lr: f32) -> Self {
97        Self {
98            lr: Decay::None(lr),
99            ..Default::default()
100        }
101    }
102
103    /// Sets the l1 regularization weight, keeping all other parameters unaffected.
104    ///
105    /// This method can be chained in a builder-pattern kind of way.
106    pub fn and_l1(mut self, l1: f32) -> Self {
107        self.l1_reg = Some(Decay::None(l1));
108        self
109    }
110
111    /// Sets the l2 regularization weight, keeping all other parameters unaffected.
112    ///
113    /// This method can be chained in a builder-pattern kind of way.
114    pub fn and_l2(mut self, l2: f32) -> Self {
115        self.l2_reg = Some(Decay::None(l2));
116        self
117    }
118
119    /// Sets the learning rate (linear) decay, keeping all other parameters unaffected.
120    ///
121    /// This method can be chained in a builder-pattern kind of way.
122    pub fn and_lr_decay(mut self, lr_decay: f32) -> Self {
123        self.lr = Decay::Linear {
124            start: self.lr.start_value(),
125            decay: lr_decay,
126        };
127        self
128    }
129
130    /// Sets the learning rate (cosine) decay, keeping all other parameters unaffected.
131    ///
132    /// This method can be chained in a builder-pattern kind of way.
133    pub fn and_lr_cosine_decay(mut self, final_lr: f32, over_steps: usize) -> Self {
134        self.lr = Decay::Cosine {
135            start: self.lr.start_value(),
136            end: final_lr,
137            num_steps: over_steps,
138        };
139        self
140    }
141
142    /// Sets the number of epochs the lr should linearly ramp up to the LR at the start of training.
143    ///
144    /// This method can be chained in a builder-pattern kind of way.
145    pub fn and_soft_start(mut self, epochs: usize) -> Self {
146        self.soft_start_epochs = Some(epochs);
147        self
148    }
149
150    /// Sets the maximum magnitude a single gradient can be. Any gradient larger in magnitude is
151    /// limited to this magnitude.
152    pub fn and_gradient_clip(mut self, clip: f32) -> Self {
153        self.grad_clip = Some(clip);
154        self
155    }
156
157    pub fn current_lr(&self) -> f32 {
158        let lr = self.lr.at_timestep(self.step);
159        match self.soft_start_epochs {
160            Some(soft_start_epochs) => {
161                if self.step < soft_start_epochs {
162                    lr / soft_start_epochs as f32 * (1 + self.step) as f32
163                } else {
164                    lr
165                }
166            }
167            None => lr,
168        }
169    }
170
171    fn clip_grad<G: Dtype>(&self, grad: G) -> G {
172        if let Some(clip) = self.grad_clip {
173            let clip = G::from_f32(clip).unwrap();
174            let neg_clip = G::default() - clip;
175            if grad > clip {
176                clip
177            } else if grad < neg_clip {
178                neg_clip
179            } else {
180                grad
181            }
182        } else {
183            grad
184        }
185    }
186
187    /// Returns the [TrainParams] structure. Exists on [TrainParams] to match the signature
188    /// of other updater implementations.
189    pub fn train_params(&self) -> &Self {
190        self
191    }
192}
193
194impl<G: Gradients> GradAdjuster<G> for TrainParams {
195    fn adjust(&mut self, mut gradient_updates: G, loss: f32) -> G {
196        let l = G::Concrete::from_f32(-loss * self.current_lr()).unwrap();
197        gradient_updates
198            .grad_iter_mut()
199            .for_each(|g| *g = self.clip_grad(*g) * l);
200        gradient_updates
201    }
202}
203
204impl GradApplyer for TrainParams {
205    fn apply<G: Gradients>(
206        &mut self,
207        gradient_updates: G,
208        weights: &mut G,
209    ) -> Result<(), crate::Error> {
210        let l1 = self.l1_reg.as_ref().map(|d| d.at_timestep(self.step));
211        let l2 = self.l2_reg.as_ref().map(|d| d.at_timestep(self.step));
212
213        weights
214            .grad_iter_mut_with_class()
215            .zip(gradient_updates.into_grads())
216            .for_each(|((w, c), u)| {
217                let reg_penalty = if c.should_regularize() {
218                    G::Concrete::from_f32(if let Some(l1) = l1 {
219                        if *w > G::Concrete::default() {
220                            l1
221                        } else if *w < G::Concrete::default() {
222                            -l1
223                        } else {
224                            0.0
225                        }
226                    } else {
227                        0.0
228                    })
229                    .unwrap()
230                        + if let Some(l2) = l2 {
231                            (*w) * (G::Concrete::ONE + G::Concrete::ONE)
232                                * G::Concrete::from_f32(l2).unwrap()
233                        } else {
234                            G::Concrete::default()
235                        }
236                } else {
237                    G::Concrete::default()
238                };
239
240                *w += u - reg_penalty;
241            });
242
243        Ok(())
244    }
245
246    fn advance_step(&mut self) {
247        self.step += 1;
248    }
249}
250
251/// Implements momentum computation in addition to the basics provided by [TrainParams].
252///
253/// Implements optimizer traits, so it can be passed into training methods.
254pub struct Momentum<G: Gradients> {
255    params: TrainParams,
256    velocity: G,
257    momentum_coeff: f32,
258
259    similarity_term: Option<f32>,
260}
261
262impl<G: Gradients> Momentum<G> {
263    /// Constructs a new object with the given coefficient for momentum.
264    ///
265    /// Typical values are 0 to 1, with 0.5 being a good starting value.
266    /// 0.9 is used by torch, but leads to oscillations with high learning
267    /// rates or small batch sizes.
268    pub fn new(params: TrainParams, momentum_coeff: f32) -> Momentum<G> {
269        Self {
270            params,
271            momentum_coeff,
272            velocity: G::empty(),
273            similarity_term: None, // Typical value 0.5=>1.5,
274        }
275    }
276
277    /// Sets the penalty term for updates that deviate (in cosine distance) from
278    /// the velocity. Typical values are 0.5 to 1.5.
279    pub fn and_similarity_penalty(mut self, similarity_term: f32) -> Self {
280        self.similarity_term = Some(similarity_term);
281        self
282    }
283
284    fn update(&mut self, gradient_updates: G, loss: f32) -> G {
285        let mc = G::Concrete::from_f32(self.momentum_coeff).unwrap();
286        let loss = G::Concrete::from_f32(-loss * self.params.current_lr()).unwrap();
287
288        let sim_mul = G::Concrete::from_f32(if let Some(term) = self.similarity_term {
289            self.velocity
290                .cosine_similarity(&gradient_updates)
291                .map(|x| (2.0 * (x + 1.0) + term).log2() + 0.25)
292                .unwrap_or(1.0)
293        } else {
294            1.0
295        })
296        .unwrap();
297
298        // v = coeff * last_v + gradient_updates
299        self.velocity
300            .grad_iter_mut()
301            .zip(gradient_updates.into_grads())
302            .for_each(|(v, g)| {
303                *v = (*v * mc) + (self.params.clip_grad(g) * loss * sim_mul);
304            });
305
306        self.velocity.clone()
307    }
308
309    /// Returns the [TrainParams] structure.
310    pub fn train_params(&self) -> &TrainParams {
311        &self.params
312    }
313}
314
315impl<G: Gradients> GradAdjuster<G> for Momentum<G> {
316    fn adjust(&mut self, gradient_updates: G, loss: f32) -> G {
317        self.update(gradient_updates, loss)
318    }
319}
320
321impl<G: Gradients> GradApplyer for Momentum<G> {
322    fn apply<G2: Gradients>(
323        &mut self,
324        gradient_updates: G2,
325        weights: &mut G2,
326    ) -> Result<(), crate::Error> {
327        self.params.apply(gradient_updates, weights)
328    }
329
330    fn advance_step(&mut self) {
331        self.params.advance_step();
332    }
333}
334
335enum RMSPropBase<G: Gradients> {
336    NoMomentum(TrainParams),
337    Momentum(Momentum<G>),
338}
339
340impl<G: Gradients> RMSPropBase<G> {
341    /// Returns the [TrainParams] structure.
342    pub fn train_params(&self) -> &TrainParams {
343        match self {
344            RMSPropBase::NoMomentum(p) => p,
345            RMSPropBase::Momentum(m) => m.train_params(),
346        }
347    }
348}
349
350impl<G: Gradients> GradAdjuster<G> for RMSPropBase<G> {
351    fn adjust(&mut self, gradient_updates: G, loss: f32) -> G {
352        use RMSPropBase::*;
353        match self {
354            NoMomentum(params) => params.adjust(gradient_updates, loss),
355            Momentum(m) => m.adjust(gradient_updates, loss),
356        }
357    }
358}
359
360impl<G: Gradients> GradApplyer for RMSPropBase<G> {
361    fn apply<G2: Gradients>(
362        &mut self,
363        gradient_updates: G2,
364        weights: &mut G2,
365    ) -> Result<(), crate::Error> {
366        use RMSPropBase::*;
367        match self {
368            NoMomentum(params) => params.apply(gradient_updates, weights),
369            Momentum(m) => m.apply(gradient_updates, weights),
370        }
371    }
372
373    fn advance_step(&mut self) {
374        match self {
375            RMSPropBase::NoMomentum(p) => p.advance_step(),
376            RMSPropBase::Momentum(m) => m.advance_step(),
377        }
378    }
379}
380
381/// Implements rmsprop on top of basic training parameters or [Momentum].
382pub struct RMSProp<G: Gradients>
383where
384    G::Concrete: Float,
385{
386    base: RMSPropBase<G>,
387    accumulator: G,
388    beta: f32,
389}
390
391impl<G: Gradients> RMSProp<G>
392where
393    G::Concrete: Float,
394{
395    /// Constructs a new rmsprop optimizer.
396    pub fn new(params: TrainParams, beta: f32) -> RMSProp<G> {
397        Self {
398            beta,
399            base: RMSPropBase::NoMomentum(params),
400            accumulator: G::empty(),
401        }
402    }
403
404    /// Constructs a new rmsprop optimizer with momentum.
405    pub fn new_with_momentum(params: TrainParams, momentum_coeff: f32, beta: f32) -> RMSProp<G> {
406        Self {
407            beta,
408            base: RMSPropBase::Momentum(Momentum::new(params, momentum_coeff)),
409            accumulator: G::empty(),
410        }
411    }
412
413    /// Sets the penalty term for updates that deviate (in cosine distance) from
414    /// the velocity. Typical values are 0.5 to 1.5.
415    pub fn and_similarity_penalty(mut self, similarity_term: f32) -> Self {
416        if let RMSPropBase::Momentum(m) = &mut self.base {
417            m.similarity_term = Some(similarity_term);
418        }
419        self
420    }
421
422    /// Returns the [TrainParams] structure.
423    pub fn train_params(&self) -> &TrainParams {
424        self.base.train_params()
425    }
426}
427
428impl<G: Gradients> GradAdjuster<G> for RMSProp<G>
429where
430    G::Concrete: Float,
431{
432    fn adjust(&mut self, mut gradient_updates: G, loss: f32) -> G {
433        let b = G::Concrete::from_f32(self.beta).unwrap();
434        self.accumulator
435            .grad_iter_mut()
436            .zip(gradient_updates.grad_iter_mut())
437            .for_each(|(a, u)| {
438                let new_a = (*a * b) + (G::Concrete::ONE - b) * (*u) * (*u);
439                *a = new_a;
440
441                // rmsprop divides the learning rate by sqrt(accumulator + epsilon).
442                // the learning rate will be multiplied in next, so we just apply it to
443                // the whole gradient.
444                *u /= (new_a + G::Concrete::SMOL).sqrt();
445            });
446
447        self.base.adjust(gradient_updates, loss)
448    }
449}
450
451impl<G: Gradients> GradApplyer for RMSProp<G>
452where
453    G::Concrete: Float,
454{
455    fn apply<G2: Gradients>(
456        &mut self,
457        gradient_updates: G2,
458        weights: &mut G2,
459    ) -> Result<(), crate::Error> {
460        self.base.apply(gradient_updates, weights)
461    }
462
463    fn advance_step(&mut self) {
464        self.base.advance_step();
465    }
466}
minidx_core/optimizers.rs

minidx_core/
optimizers.rs
