haskell-vae/GradUtils.hs at master · declanoller/haskell-vae · GitHub

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
module GradUtils
( get_init_grads
, get_init_grads_vae
, add_grads
, mult_grads
, grads_momentum
, get_init_adam_optim
, step_vae_adam_optim
) where


import Numeric.LinearAlgebra
import MatNNGradTypes
import MatrixUtils


get_init_grads :: NN -> Grads
get_init_grads (WeightMatList nn) = grads_zeroed
  where mat_list = map (\(Layer x) -> x) nn
        g = GradMatList mat_list
        grads_zeroed = mult_grads 0 g


get_init_grads_vae :: VAE -> (Grads, Grads)
get_init_grads_vae (VAE (front_half, back_half)) = (front_half_zeroed, back_half_zeroed)
  where front_half_zeroed = get_init_grads front_half
        back_half_zeroed = get_init_grads back_half


-- Set up initial Adam optimizer. Uses a passed VAE to determine the sizes of the
-- matrices. Params used are same as PyTorch.
get_init_adam_optim :: VAE -> Double -> VAEAdamOptim
get_init_adam_optim (VAE (front_half, back_half)) alpha = VAEAdamOptimParams (front_half_nn_optim, back_half_nn_optim, alpha)
  where GradMatList front_half_zeroed = get_init_grads front_half
        GradMatList back_half_zeroed = get_init_grads back_half
        front_half_nn_optim = NNAdamOptimParams $ map (\x -> AdamOptimParams (x, x, 1, beta_1, beta_2)) front_half_zeroed
        back_half_nn_optim = NNAdamOptimParams $ map (\x -> AdamOptimParams (x, x, 1, beta_1, beta_2)) back_half_zeroed
        beta_1 = 0.9
        beta_2 = 0.999


add_grads :: Grads -> Grads -> Grads
add_grads (GradMatList g1) (GradMatList g2) = g3
  where g3 = GradMatList $ zipWith (+) g1 g2

mult_grads :: Double -> Grads -> Grads
mult_grads x (GradMatList g) = g_out
  where g_out = GradMatList $ map (scale x) g


grads_momentum :: (Grads, Grads) -> (Grads, Grads) -> Double -> (Grads, Grads)
grads_momentum (last_vel_front, last_vel_back) (cur_grads_front, cur_grads_back) alpha = (new_vel_front, new_vel_back)
  where new_vel_front = add_grads (mult_grads momentum last_vel_front) (mult_grads (alpha) cur_grads_front)
        new_vel_back = add_grads (mult_grads momentum last_vel_back) (mult_grads (alpha) cur_grads_back)
        momentum = 0.9


{-
Steps the Adam optimizer. Note: the epsilon term, 10**(-6), is typically
smaller, 10**(-8), but I found that that was numerically unstable (a known
issue with Adam). I found that decreasing it to 10**(-6) prevented that, without
a noticeable drop in performance. However, since it's typically 10**(-8),
in other implementations, it's possible something is wrong somewhere. Also note:
it's much slower to multiply by an elementwise inverted matrix than to divide by
one.
-}
step_adam_optim :: AdamOptim -> Matrix R -> (AdamOptim, Matrix R)
step_adam_optim adam_optim g = (new_adam_optim, adam_grads)
  where AdamOptimParams (m, v, t, beta_1, beta_2) = adam_optim
        new_m = (scale beta_1 m) + (scale (1 - beta_1) g)
        new_v = (scale beta_2 v) + (scale (1 - beta_2) (g**2))
        m_hat = scale ((1 - beta_1^t)**(-1.0)) new_m
        v_hat = scale ((1 - beta_2^t)**(-1.0)) new_v
        adam_grads = m_hat*((v_hat**0.5 + 10**(-6.0))**(-1.0))
        new_adam_optim = AdamOptimParams (new_m, new_v, t+1, beta_1, beta_2)


step_nn_adam_optim :: NNAdamOptim -> Grads -> (NNAdamOptim, Grads)
step_nn_adam_optim nn_adam_optim cur_grads = (NNAdamOptimParams new_layer_adam_optims, GradMatList new_layer_grads)
  where (NNAdamOptimParams layer_adam_optims) = nn_adam_optim
        (GradMatList layer_grads) = cur_grads
        (new_layer_adam_optims, new_layer_grads) = unzip $ zipWith step_adam_optim layer_adam_optims layer_grads


-- Unpack and step the Adam optimizer. Essentially steps each of the NNs, and each
-- of their layers, separately.
step_vae_adam_optim :: VAEAdamOptim -> (Grads, Grads) -> (VAEAdamOptim, (Grads, Grads))
step_vae_adam_optim vae_adam_optim (cur_grads_front, cur_grads_back) = (new_vae_adam_optim, new_grads_tup)
  where VAEAdamOptimParams (front_adam_optim, back_adam_optim, alpha) = vae_adam_optim
        (new_nn_front_optim, new_grads_front) = step_nn_adam_optim front_adam_optim cur_grads_front
        (new_nn_back_optim, new_grads_back) = step_nn_adam_optim back_adam_optim cur_grads_back
        new_vae_adam_optim = VAEAdamOptimParams (new_nn_front_optim, new_nn_back_optim, alpha)
        new_grads_tup = (new_grads_front, new_grads_back)