1pub mod adagrad;
8pub mod sgd;
9
10pub use adagrad::{AdaGrad, AdaGradConfig};
11pub use sgd::{Sgd, SgdConfig};
12
13#[derive(Debug, Clone)]
15#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
16pub enum Optimizer {
17 Sgd(Sgd),
19 AdaGrad(AdaGrad),
21}
22
23impl Optimizer {
24 pub fn sgd(feature_count: usize, config: SgdConfig) -> Result<Self, RillError> {
26 Ok(Optimizer::Sgd(Sgd::new(feature_count, config)?))
27 }
28
29 pub fn adagrad(feature_count: usize, config: AdaGradConfig) -> Result<Self, RillError> {
31 Ok(Optimizer::AdaGrad(AdaGrad::new(feature_count, config)?))
32 }
33
34 pub fn param_count(&self) -> usize {
36 match self {
37 Optimizer::Sgd(o) => o.param_count(),
38 Optimizer::AdaGrad(o) => o.param_count(),
39 }
40 }
41
42 pub fn samples_seen(&self) -> u64 {
44 match self {
45 Optimizer::Sgd(o) => o.samples_seen(),
46 Optimizer::AdaGrad(o) => o.samples_seen(),
47 }
48 }
49
50 pub fn step(
55 &mut self,
56 weights: &mut [f64],
57 intercept: &mut f64,
58 grad_weights: &[f64],
59 grad_intercept: f64,
60 ) -> Result<(), RillError> {
61 match self {
62 Optimizer::Sgd(o) => o.step(weights, intercept, grad_weights, grad_intercept),
63 Optimizer::AdaGrad(o) => o.step(weights, intercept, grad_weights, grad_intercept),
64 }
65 }
66
67 pub fn reset(&mut self) {
69 match self {
70 Optimizer::Sgd(o) => o.reset(),
71 Optimizer::AdaGrad(o) => o.reset(),
72 }
73 }
74}
75
76use crate::error::RillError;