wehub-resource-sync
196 lines
6.6 KiB
Python
196 lines
6.6 KiB
Python
# coding:utf-8
|
|
import autograd.numpy as np
|
|
from autograd import elementwise_grad
|
|
|
|
from mla.neuralnet.activations import sigmoid
|
|
from mla.neuralnet.initializations import get_initializer
|
|
from mla.neuralnet.layers import Layer, get_activation, ParamMixin
|
|
from mla.neuralnet.parameters import Parameters
|
|
|
|
"""
|
|
References:
|
|
Understanding LSTM Networks http://colah.github.io/posts/2015-08-Understanding-LSTMs/
|
|
A Critical Review of Recurrent Neural Networks for Sequence Learning http://arxiv.org/pdf/1506.00019v4.pdf
|
|
"""
|
|
|
|
|
|
class LSTM(Layer, ParamMixin):
|
|
def __init__(
|
|
self,
|
|
hidden_dim,
|
|
activation="tanh",
|
|
inner_init="orthogonal",
|
|
parameters=None,
|
|
return_sequences=True,
|
|
):
|
|
self.return_sequences = return_sequences
|
|
self.hidden_dim = hidden_dim
|
|
self.inner_init = get_initializer(inner_init)
|
|
self.activation = get_activation(activation)
|
|
self.activation_d = elementwise_grad(self.activation)
|
|
self.sigmoid_d = elementwise_grad(sigmoid)
|
|
|
|
if parameters is None:
|
|
self._params = Parameters()
|
|
else:
|
|
self._params = parameters
|
|
|
|
self.last_input = None
|
|
self.states = None
|
|
self.outputs = None
|
|
self.gates = None
|
|
self.hprev = None
|
|
self.input_dim = None
|
|
self.W = None
|
|
self.U = None
|
|
|
|
def setup(self, x_shape):
|
|
"""
|
|
Naming convention:
|
|
i : input gate
|
|
f : forget gate
|
|
c : cell
|
|
o : output gate
|
|
|
|
Parameters
|
|
----------
|
|
x_shape : np.array(batch size, time steps, input shape)
|
|
"""
|
|
self.input_dim = x_shape[2]
|
|
# Input -> Hidden
|
|
W_params = ["W_i", "W_f", "W_o", "W_c"]
|
|
# Hidden -> Hidden
|
|
U_params = ["U_i", "U_f", "U_o", "U_c"]
|
|
# Bias terms
|
|
b_params = ["b_i", "b_f", "b_o", "b_c"]
|
|
|
|
# Initialize params
|
|
for param in W_params:
|
|
self._params[param] = self._params.init((self.input_dim, self.hidden_dim))
|
|
|
|
for param in U_params:
|
|
self._params[param] = self.inner_init((self.hidden_dim, self.hidden_dim))
|
|
|
|
for param in b_params:
|
|
self._params[param] = np.full((self.hidden_dim,), self._params.initial_bias)
|
|
|
|
# Combine weights for simplicity
|
|
self.W = [self._params[param] for param in W_params]
|
|
self.U = [self._params[param] for param in U_params]
|
|
|
|
# Init gradient arrays for all weights
|
|
self._params.init_grad()
|
|
|
|
self.hprev = np.zeros((x_shape[0], self.hidden_dim))
|
|
self.oprev = np.zeros((x_shape[0], self.hidden_dim))
|
|
|
|
def forward_pass(self, X):
|
|
n_samples, n_timesteps, input_shape = X.shape
|
|
p = self._params
|
|
self.last_input = X
|
|
|
|
self.states = np.zeros((n_samples, n_timesteps + 1, self.hidden_dim))
|
|
self.outputs = np.zeros((n_samples, n_timesteps + 1, self.hidden_dim))
|
|
self.gates = {
|
|
k: np.zeros((n_samples, n_timesteps, self.hidden_dim))
|
|
for k in ["i", "f", "o", "c"]
|
|
}
|
|
|
|
self.states[:, -1, :] = self.hprev
|
|
self.outputs[:, -1, :] = self.oprev
|
|
|
|
for i in range(n_timesteps):
|
|
t_gates = np.dot(X[:, i, :], self.W) + np.dot(
|
|
self.outputs[:, i - 1, :], self.U
|
|
)
|
|
|
|
# Input
|
|
self.gates["i"][:, i, :] = sigmoid(t_gates[:, 0, :] + p["b_i"])
|
|
# Forget
|
|
self.gates["f"][:, i, :] = sigmoid(t_gates[:, 1, :] + p["b_f"])
|
|
# Output
|
|
self.gates["o"][:, i, :] = sigmoid(t_gates[:, 2, :] + p["b_o"])
|
|
# Cell
|
|
self.gates["c"][:, i, :] = self.activation(t_gates[:, 3, :] + p["b_c"])
|
|
|
|
# (previous state * forget) + input + cell
|
|
self.states[:, i, :] = (
|
|
self.states[:, i - 1, :] * self.gates["f"][:, i, :]
|
|
+ self.gates["i"][:, i, :] * self.gates["c"][:, i, :]
|
|
)
|
|
self.outputs[:, i, :] = self.gates["o"][:, i, :] * self.activation(
|
|
self.states[:, i, :]
|
|
)
|
|
|
|
self.hprev = self.states[:, n_timesteps - 1, :].copy()
|
|
self.oprev = self.outputs[:, n_timesteps - 1, :].copy()
|
|
|
|
if self.return_sequences:
|
|
return self.outputs[:, 0:-1, :]
|
|
else:
|
|
return self.outputs[:, -2, :]
|
|
|
|
def backward_pass(self, delta):
|
|
if len(delta.shape) == 2:
|
|
delta = delta[:, np.newaxis, :]
|
|
|
|
n_samples, n_timesteps, input_shape = delta.shape
|
|
|
|
# Temporal gradient arrays
|
|
grad = {k: np.zeros_like(self._params[k]) for k in self._params.keys()}
|
|
|
|
dh_next = np.zeros((n_samples, input_shape))
|
|
output = np.zeros((n_samples, n_timesteps, self.input_dim))
|
|
|
|
# Backpropagation through time
|
|
for i in reversed(range(n_timesteps)):
|
|
dhi = (
|
|
delta[:, i, :]
|
|
* self.gates["o"][:, i, :]
|
|
* self.activation_d(self.states[:, i, :])
|
|
+ dh_next
|
|
)
|
|
|
|
og = delta[:, i, :] * self.activation(self.states[:, i, :])
|
|
de_o = og * self.sigmoid_d(self.gates["o"][:, i, :])
|
|
|
|
grad["W_o"] += np.dot(self.last_input[:, i, :].T, de_o)
|
|
grad["U_o"] += np.dot(self.outputs[:, i - 1, :].T, de_o)
|
|
grad["b_o"] += de_o.sum(axis=0)
|
|
|
|
de_f = (dhi * self.states[:, i - 1, :]) * self.sigmoid_d(
|
|
self.gates["f"][:, i, :]
|
|
)
|
|
grad["W_f"] += np.dot(self.last_input[:, i, :].T, de_f)
|
|
grad["U_f"] += np.dot(self.outputs[:, i - 1, :].T, de_f)
|
|
grad["b_f"] += de_f.sum(axis=0)
|
|
|
|
de_i = (dhi * self.gates["c"][:, i, :]) * self.sigmoid_d(
|
|
self.gates["i"][:, i, :]
|
|
)
|
|
grad["W_i"] += np.dot(self.last_input[:, i, :].T, de_i)
|
|
grad["U_i"] += np.dot(self.outputs[:, i - 1, :].T, de_i)
|
|
grad["b_i"] += de_i.sum(axis=0)
|
|
|
|
de_c = (dhi * self.gates["i"][:, i, :]) * self.activation_d(
|
|
self.gates["c"][:, i, :]
|
|
)
|
|
grad["W_c"] += np.dot(self.last_input[:, i, :].T, de_c)
|
|
grad["U_c"] += np.dot(self.outputs[:, i - 1, :].T, de_c)
|
|
grad["b_c"] += de_c.sum(axis=0)
|
|
|
|
dh_next = dhi * self.gates["f"][:, i, :]
|
|
|
|
# TODO: propagate error to the next layer
|
|
|
|
# Change actual gradient arrays
|
|
for k in grad.keys():
|
|
self._params.update_grad(k, grad[k])
|
|
return output
|
|
|
|
def shape(self, x_shape):
|
|
if self.return_sequences:
|
|
return x_shape[0], x_shape[1], self.hidden_dim
|
|
else:
|
|
return x_shape[0], self.hidden_dim
|