wehub-resource-sync
137 lines
4.2 KiB
Python
137 lines
4.2 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# pylint: disable=invalid-name
|
|
# ruff: noqa: E741
|
|
"""LSTM reference implementation using numpy."""
|
|
|
|
import numpy as np
|
|
|
|
|
|
def lstm_python(
|
|
Xs: np.array,
|
|
Wi: np.array,
|
|
Wh: np.array,
|
|
Bi: np.array = None,
|
|
Bh: np.array = None,
|
|
h_init: np.array = None,
|
|
c_init: np.array = None,
|
|
proj: np.array = None,
|
|
p_i: np.array = None,
|
|
p_f: np.array = None,
|
|
p_o: np.array = None,
|
|
f_act: str = "sigmoid",
|
|
g_act: str = "tanh",
|
|
h_act: str = "tanh",
|
|
reverse: bool = False,
|
|
weight_layout: str = "IFGO",
|
|
):
|
|
"""LSTM reference implementation using numpy
|
|
|
|
Parameters
|
|
----------
|
|
Xs : np.array
|
|
(seq_length, batch_size, in_dim)
|
|
Wi : np.array
|
|
(4 * hidden_dim, in_dim)
|
|
Wh : np.array
|
|
(4 * hidden_dim, out_dim) where out_dim = proj_dim if proj_dim > 0, else hidden_dim
|
|
Bi : np.array, optional
|
|
(4 * hidden_dim,), by default None
|
|
Bh : np.array, optional
|
|
(4 * hidden_dim,), by default None
|
|
h_init : np.array, optional
|
|
(batch_size, out_dim), by default None
|
|
c_init : np.array, optional
|
|
(batch_size, hidden_dim), by default None
|
|
proj : np.array, optional
|
|
(proj_dim, hidden_dim), by default None
|
|
p_i, p_f, p_o: np.array, optional
|
|
(batch_size, hidden_dim), by default None
|
|
f_act, g_act, h_act: str, optional
|
|
activations, by default "sigmoid", "tanh", "tanh"
|
|
reverse : bool, optional
|
|
process Xs in reverse, by default False
|
|
weight_layout : str, optional
|
|
Packed layout for weights and biases, by default "IFGO"
|
|
"""
|
|
i_gate_idx = weight_layout.find("I")
|
|
f_gate_idx = weight_layout.find("F")
|
|
g_gate_idx = weight_layout.find("G")
|
|
o_gate_idx = weight_layout.find("O")
|
|
|
|
str2act = {"sigmoid": lambda x: 1 / (1 + np.exp(-x)), "tanh": np.tanh}
|
|
|
|
f_act = str2act[f_act]
|
|
g_act = str2act[g_act]
|
|
h_act = str2act[h_act]
|
|
|
|
S, B, F = Xs.shape
|
|
H = Wi.shape[0] // 4
|
|
O = Wh.shape[1]
|
|
|
|
# make life a bit easier
|
|
Wi = np.reshape(Wi, (4, H, F))
|
|
Wh = np.reshape(Wh, (4, H, O))
|
|
if Bi is not None:
|
|
Bi = np.reshape(Bi, (4, H))
|
|
if Bh is not None:
|
|
Bh = np.reshape(Bh, (4, H))
|
|
|
|
h0 = h_init if h_init is not None else np.zeros((B, O), "float32")
|
|
c0 = c_init if c_init is not None else np.zeros((B, H), "float32")
|
|
|
|
hs = [h0]
|
|
cs = [c0]
|
|
|
|
for t in range(S):
|
|
x = Xs[S - t - 1 if reverse else t]
|
|
xh = [np.matmul(x, Wi[g].T) for g in range(4)]
|
|
if Bi is not None:
|
|
xh = [xh[g] + Bi[g] for g in range(4)]
|
|
|
|
hh = [np.matmul(hs[t], Wh[g].T) for g in range(4)]
|
|
if Bh is not None:
|
|
hh = [hh[g] + Bh[g] for g in range(4)]
|
|
|
|
sums = [xh[g] + hh[g] for g in range(4)]
|
|
|
|
if p_i is not None and p_f is not None:
|
|
i_gate = f_act(sums[i_gate_idx] + p_i * cs[t])
|
|
f_gate = f_act(sums[f_gate_idx] + p_f * cs[t])
|
|
else:
|
|
i_gate = f_act(sums[i_gate_idx])
|
|
f_gate = f_act(sums[f_gate_idx])
|
|
|
|
g_gate = g_act(sums[g_gate_idx])
|
|
|
|
next_c = f_gate * cs[t] + i_gate * g_gate
|
|
|
|
if p_o is not None:
|
|
o_gate = f_act(sums[o_gate_idx] + p_o * next_c)
|
|
else:
|
|
o_gate = f_act(sums[o_gate_idx])
|
|
|
|
next_h = o_gate * h_act(next_c)
|
|
|
|
if proj is not None:
|
|
next_h = np.matmul(next_h, proj.T)
|
|
|
|
hs.append(next_h)
|
|
cs.append(next_c)
|
|
|
|
return np.stack(hs[1:], axis=0), np.stack(cs[1:], axis=0)
|