wyrm 0.9.1

A low-overhead, define-by-run autodifferentiation library.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

//! Optimization module.
//!
//! Contains a number of optimizers.
use std::cell::Cell;
use {ParameterNode, Variable};
mod adagrad;
mod adam;
mod barrier;
mod sgd;

pub use self::adagrad::Adagrad;
pub use self::adam::Adam;
use self::barrier::SynchronizationBarrier;
pub use self::sgd::SGD;

/// Core trait implemented by all optimizer methods.
pub trait Optimizer {
    /// Perform a single SGD step.
    fn step(&self, parameters: &[Variable<ParameterNode>]);
}

/// Trait implemented by synchronizable optimizers.
///
/// Using a set of synchronized optimizers guarantees that parameter
/// updates will always happen in the same order, guaranteeing reproducible
/// results at the price of some performance relative to asynchronous parallel
/// optimization.
pub trait Synchronizable {
    /// Synchronize this optimizer, producing a set of synchronized optimimzers
    /// to be used by individual fitting threads.
    fn synchronized(&self, num_threads: usize) -> Vec<SynchronizedOptimizer<Self>>
    where
        Self: Sized,
    {
        self.synchronized_with_step(num_threads, 8)
    }
    /// Synchronize this optimizer, producing a set of synchronized optimimzers
    /// to be used by individual fitting threads. The threads will synchonize
    /// their updates every `step_size` steps.
    fn synchronized_with_step(
        &self,
        num_threads: usize,
        step_size: usize,
    ) -> Vec<SynchronizedOptimizer<Self>>
    where
        Self: Sized;
}

/// Synchronized optimizer wrapper.
pub struct SynchronizedOptimizer<'a, T: 'a> {
    step_size: usize,
    num_updates: Cell<usize>,
    optimizer: &'a T,
    barrier_guard: barrier::SynchronizationBarrierGuard,
}

impl<'a, T: 'a> SynchronizedOptimizer<'a, T> {
    fn new(
        optimizer: &'a T,
        barrier_guard: barrier::SynchronizationBarrierGuard,
        step_size: usize,
    ) -> Self {
        SynchronizedOptimizer {
            step_size: step_size,
            num_updates: Cell::new(0),
            optimizer: optimizer,
            barrier_guard: barrier_guard,
        }
    }
}

impl<'a, T> Optimizer for SynchronizedOptimizer<'a, T>
where
    T: Optimizer,
{
    fn step(&self, parameters: &[Variable<ParameterNode>]) {
        self.num_updates.set(self.num_updates.get() + 1);

        if self.num_updates.get() == self.step_size {
            let _barrier = self.barrier_guard.synchronize();
            self.optimizer.step(parameters);

            self.num_updates.set(0);
        }
    }
}

impl<T> Synchronizable for T
where
    T: Optimizer + Sized,
{
    fn synchronized_with_step(
        &self,
        num_threads: usize,
        step_size: usize,
    ) -> Vec<SynchronizedOptimizer<T>> {
        let barrier = SynchronizationBarrier::default();

        (0..num_threads)
            .map(|_| SynchronizedOptimizer::new(self, barrier.register_thread(), step_size))
            .collect()
    }
}

macro_rules! impl_optimizer_enum {
    ($(($tag:ident, $type:ty)),*) => {
        /// Enum containing all optimizers.
        ///
        /// Makes runtime switching between optimizers slightly more ergonomic.
        pub enum Optimizers {
            $(
                #[allow(missing_docs)]
                $tag($type),
            )*
        }

        impl Optimizer for Optimizers {
            fn step(&self, parameters: &[Variable<ParameterNode>]) {
                match self {
                    $(
                        Optimizers::$tag(val) => val.step(parameters),
                    )*
                }
            }
        }
    }
}

impl_optimizer_enum!((SGD, SGD), (Adagrad, Adagrad), (Adam, Adam));