|
- import hetu as ht
- from hetu import init
- import numpy as np
-
-
- def lstm(x, y_):
- '''
- LSTM model, for MNIST dataset.
-
- Parameters:
- x: Variable(hetu.gpu_ops.Node.Node), shape (N, dims)
- y_: Variable(hetu.gpu_ops.Node.Node), shape (N, num_classes)
- Return:
- loss: Variable(hetu.gpu_ops.Node.Node), shape (1,)
- y: Variable(hetu.gpu_ops.Node.Node), shape (N, num_classes)
- '''
- diminput = 28
- dimhidden = 128
- dimoutput = 10
- nsteps = 28
-
- forget_gate_w = init.random_normal(
- shape=(diminput, dimhidden), stddev=0.1, name="lstm_forget_gate_w")
- forget_gate_u = init.random_normal(
- shape=(dimhidden, dimhidden), stddev=0.1, name="lstm_forget_gate_u")
- forget_gate_b = init.random_normal(
- shape=(dimhidden,), stddev=0.1, name="lstm_forget_gate_b")
- input_gate_w = init.random_normal(
- shape=(diminput, dimhidden), stddev=0.1, name="lstm_input_gate_w")
- input_gate_u = init.random_normal(
- shape=(dimhidden, dimhidden), stddev=0.1, name="lstm_input_gate_u")
- input_gate_b = init.random_normal(
- shape=(dimhidden,), stddev=0.1, name="lstm_input_gate_b")
- output_gate_w = init.random_normal(
- shape=(diminput, dimhidden), stddev=0.1, name="lstm_output_gate_w")
- output_gate_u = init.random_normal(
- shape=(dimhidden, dimhidden), stddev=0.1, name="lstm_output_gate_u")
- output_gate_b = init.random_normal(
- shape=(dimhidden,), stddev=0.1, name="lstm_output_gate_b")
- tanh_w = init.random_normal(
- shape=(diminput, dimhidden), stddev=0.1, name="lstm_tanh_w")
- tanh_u = init.random_normal(
- shape=(dimhidden, dimhidden), stddev=0.1, name="lstm_tanh_u")
- tanh_b = init.random_normal(
- shape=(dimhidden,), stddev=0.1, name="lstm_tanh_b")
- out_weights = init.random_normal(
- shape=(dimhidden, dimoutput), stddev=0.1, name="lstm_out_weight")
- out_bias = init.random_normal(
- shape=(dimoutput,), stddev=0.1, name="lstm_out_bias")
- initial_state = ht.Variable(value=np.zeros((1,)).astype(
- np.float32), name='initial_state', trainable=False)
-
- for i in range(nsteps):
- cur_x = ht.slice_op(x, (0, i * diminput), (-1, diminput))
- # forget gate
- if i == 0:
- temp = ht.matmul_op(cur_x, forget_gate_w)
- last_c_state = ht.broadcastto_op(initial_state, temp)
- last_h_state = ht.broadcastto_op(initial_state, temp)
- cur_forget = ht.matmul_op(last_h_state, forget_gate_u) + temp
- else:
- cur_forget = ht.matmul_op(
- last_h_state, forget_gate_u) + ht.matmul_op(cur_x, forget_gate_w)
- cur_forget = cur_forget + ht.broadcastto_op(forget_gate_b, cur_forget)
- cur_forget = ht.sigmoid_op(cur_forget)
- # input gate
- cur_input = ht.matmul_op(
- last_h_state, input_gate_u) + ht.matmul_op(cur_x, input_gate_w)
- cur_input = cur_input + ht.broadcastto_op(input_gate_b, cur_input)
- cur_input = ht.sigmoid_op(cur_input)
- # output gate
- cur_output = ht.matmul_op(
- last_h_state, output_gate_u) + ht.matmul_op(cur_x, output_gate_w)
- cur_output = cur_output + ht.broadcastto_op(output_gate_b, cur_output)
- cur_output = ht.sigmoid_op(cur_output)
- # tanh
- cur_tanh = ht.matmul_op(last_h_state, tanh_u) + \
- ht.matmul_op(cur_x, tanh_w)
- cur_tanh = cur_tanh + ht.broadcastto_op(tanh_b, cur_tanh)
- cur_tanh = ht.tanh_op(cur_tanh)
-
- last_c_state = ht.mul_op(last_c_state, cur_forget) + \
- ht.mul_op(cur_input, cur_tanh)
- last_h_state = ht.tanh_op(last_c_state) * cur_output
-
- x = ht.matmul_op(last_h_state, out_weights)
- y = x + ht.broadcastto_op(out_bias, x)
- loss = ht.softmaxcrossentropy_op(y, y_)
- loss = ht.reduce_mean_op(loss, [0])
- return loss, y
|
