def create_a3c_lstm_network(input_tensor, output_num):
l_hid1 = tflearn.conv_2d(input_tensor, 16, 8, strides=4, activation='relu', scope='conv1', padding='valid')
l_hid2 = tflearn.conv_2d(l_hid1, 32, 4, strides=2, activation='relu', scope='conv2', padding='valid')
l_hid3 = tflearn.fully_connected(l_hid2, 256, activation='relu', scope='dense3')
# reshape l_hid3 to lstm usable shape (1, batch_size, 256)
l_hid3_reshape = tf.reshape(l_hid3, [1, -1, 256])
# have to custom make the lstm output here to use tf.nn.dynamic_rnn
l_lstm = tflearn.BasicLSTMCell(256)
# BasicLSTMCell lists state size as tuple so we need to pass tuple into dynamic_rnn
lstm_state_size = tuple([[1, x] for x in l_lstm.state_size])
# has to specifically be the same type tf.python.ops.rnn_cell.LSTMStateTuple
from tensorflow.python.ops.nn import rnn_cell as _rnn_cell
initial_lstm_state = _rnn_cell.LSTMStateTuple(tf.placeholder(tf.float32, shape=lstm_state_size[0], name='initial_lstm_state1'),
tf.placeholder(tf.float32, shape=lstm_state_size[1], name='initial_lstm_state2'))
# dynamically get the sequence length
sequence_length = tf.reshape(tf.shape(l_hid3)[0], [1])
l_lstm4, new_lstm_state = tf.nn.dynamic_rnn(l_lstm, l_hid3_reshape,
initial_state=initial_lstm_state, sequence_length=sequence_length,
time_major=False, scope='lstm4')
# reshape lstm back to (batch_size, 256)
l_lstm4_reshape = tf.reshape(l_lstm4, [-1, 256])
actor_out = tflearn.fully_connected(l_lstm4_reshape, output_num, activation='softmax', scope='actorout')
critic_out = tflearn.fully_connected(l_lstm4_reshape, 1, activation='linear', scope='criticout')
return actor_out, critic_out, initial_lstm_state, new_lstm_state
def __call__(self,
inputs,
index,
initial_state=None,
recurrent_weights_initializer=None):
"""
:param tf.Tensor inputs: shape (time,batch,n_hidden)
:param tf.Tensor index: shape (time,batch)
:param tf.Tensor|None initial_state: shape (batch,n_hidden)
:param ()->tf.Tensor recurrent_weights_initializer:
:returns: shape (time,batch,n_hidden), shape (batch,n_hidden)
:rtype: (tf.Tensor, tf.Tensor)
"""
W = tf.get_variable(name="W_re",
shape=(self.n_hidden, self.n_hidden * 4),
initializer=recurrent_weights_initializer)
TFUtil.set_param_axes_split_info(
W, [[self.n_hidden], [self.n_hidden] * 4])
if self.rec_weight_dropout:
from TFUtil import dropout
W = dropout(W,
keep_prob=1.0 - self.rec_weight_dropout,
cond_on_train=True,
seed=TFUtil.get_random_seed())
out, _, _, final_cell_state = self.op(*self.map_layer_inputs_to_op(
X=inputs, W=W, i=index, initial_state=initial_state))
from tensorflow.python.ops.nn import rnn_cell
return out, rnn_cell.LSTMStateTuple(h=out[-1], c=final_cell_state)
def __call__(self, inputs, state, scope=None):
"""Long short-term memory cell (LSTM)."""
with tf.variable_scope(scope
or type(self).__name__): # "BasicLSTMCell"
# Parameters of gates are concatenated into one multiply for efficiency.
if self._state_is_tuple:
c, h = state
else:
c, h = array_ops.split(1, 2, state)
concat = _linear([inputs, h], 4 * self._num_units, True, 0.,
self.weights_init, self.trainable, self.restore,
self.reuse)
# i = input_gate, j = new_input, f = forget_gate, o = output_gate
i, j, f, o = array_ops.split(1, 4, concat)
new_c = (c * self._inner_activation(f + self._forget_bias) +
self._inner_activation(i) * self._activation(j))
new_h = self._activation(new_c) * self._inner_activation(o)
if self._state_is_tuple:
new_state = _rnn_cell.LSTMStateTuple(new_c, new_h)
else:
new_state = array_ops.concat(1, [new_c, new_h])
# Retrieve RNN Variables
with tf.variable_scope('Linear', reuse=True):
self.W = tf.get_variable('Matrix')
self.b = tf.get_variable('Bias')
return new_h, new_state
def __call__(self,
inputs,
index,
initial_state=None,
recurrent_weights_initializer=None):
"""
:param tf.Tensor inputs: shape (time,batch,n_hidden)
:param tf.Tensor index: shape (time,batch)
:param tf.Tensor|None initial_state: shape (batch,n_hidden)
:param ()->tf.Tensor recurrent_weights_initializer:
:returns: shape (time,batch,n_hidden), shape (batch,n_hidden)
:rtype: (tf.Tensor, tf.Tensor)
"""
from tensorflow.python.ops.nn import rnn_cell
W = tf.get_variable(name="W_re",
shape=(self.n_hidden, self.n_hidden * 4),
initializer=recurrent_weights_initializer)
TFUtil.set_param_axes_split_info(
W, [[self.n_hidden], [self.n_hidden] * 4])
if self.rec_weight_dropout:
from TFUtil import dropout
W = dropout(W,
keep_prob=1.0 - self.rec_weight_dropout,
cond_on_train=True,
seed=TFUtil.get_random_seed())
inputs.set_shape(tf.TensorShape([None, None, self.n_hidden * 4]))
W.set_shape(tf.TensorShape([self.n_hidden, self.n_hidden * 4]))
index.set_shape(tf.TensorShape([None, None]))
from TFUtil import to_float32
index = to_float32(index)
n_batch = tf.shape(inputs)[1]
if initial_state is None:
c0 = tf.zeros((n_batch, self.n_hidden),
dtype=tf.float32,
name="initial_c")
y0 = tf.zeros((n_batch, self.n_hidden),
dtype=tf.float32,
name="initial_h")
elif isinstance(initial_state, rnn_cell.LSTMStateTuple):
c0 = initial_state.c
y0 = initial_state.h
else:
c0 = initial_state
y0 = tf.zeros((n_batch, self.n_hidden),
dtype=tf.float32,
name="initial_h")
start = tf.constant(0, name="start")
step = tf.constant(self.step or 1, name="step")
out, _, _, final_cell_state = self.op(inputs, W, y0, c0, index, start,
step)
if out.get_shape().as_list()[0] is None or out.get_shape().as_list(
)[0] > 0:
final_output = out[-1]
else:
final_output = y0
return out, rnn_cell.LSTMStateTuple(h=final_output, c=final_cell_state)
def call(self, inputs, state):
assert isinstance(inputs, tf.Tensor)
if not isinstance(state, rnn_cell.LSTMStateTuple):
state = rnn_cell.LSTMStateTuple(*tuple(state))
assert isinstance(state, rnn_cell.LSTMStateTuple)
h = state.h
if self.rec_dropout_h:
h = self.dropout(h, drop_rate=self.rec_dropout_h, inside_loop=True, batch_axis=0)
dim = self.num_units * 4
x = self.linear(h, dim, with_bias=False) + inputs
cell_in, gate_in, gate_forget, gate_out = tf.split(x, 4, axis=-1)
if self.rec_dropout_u:
cell_in = self.dropout(cell_in, drop_rate=self.rec_dropout_u, inside_loop=True, batch_axis=0)
cell_in = tf.tanh(cell_in)
gate_in = tf.sigmoid(gate_in)
gate_forget = tf.sigmoid(gate_forget)
gate_out = tf.sigmoid(gate_out)
cell = state.c * gate_forget + cell_in * gate_in
out = tf.tanh(cell) * gate_out
return out, rnn_cell.LSTMStateTuple(c=cell, h=out)
def __call__(self, inputs, state):
prev_c, prev_h = state
modulated_inputs, mouldated_prev_h = self.mogrify(
inputs, prev_h, self.rank)
inputs_linear = self.linear(modulated_inputs,
4 * self.num_units,
name="modulated_inputs")
prev_h_linear = self.linear(mouldated_prev_h,
4 * self.num_units,
with_bias=False,
name="mouldated_prev_h")
lstm_in = tf.add(inputs_linear, prev_h_linear)
i, g, f, o = tf.split(lstm_in, num_or_size_splits=4, axis=-1)
new_c = tf.sigmoid(f) * prev_c + tf.sigmoid(i) * self.activation(g)
new_h = tf.sigmoid(o) * self.activation(new_c)
return new_h, rnn_cell.LSTMStateTuple(new_c, new_h)
def __call__(self,
inputs,
index,
initial_state=None,
recurrent_weights_initializer=None):
"""
:param tf.Tensor inputs: shape (time,batch,n_hidden)
:param tf.Tensor index: shape (time,batch)
:param tf.Tensor|None initial_state: shape (batch,n_hidden)
:param ()->tf.Tensor recurrent_weights_initializer:
:returns: shape (time,batch,n_hidden), shape (batch,n_hidden)
:rtype: (tf.Tensor, tf.Tensor)
"""
W = tf.get_variable(name="W_re",
shape=(self.n_hidden, self.n_hidden * 4),
initializer=recurrent_weights_initializer)
out, _, _, final_cell_state = self.op(*self.map_layer_inputs_to_op(
X=inputs, W=W, i=index, initial_state=initial_state))
from tensorflow.python.ops.nn import rnn_cell
return out, rnn_cell.LSTMStateTuple(h=out[-1], c=final_cell_state)
def create_basic_lstm_layer(self,
input_data,
input_size,
num_of_units,
scope=None):
batch_size = tf.shape(input_data)[0]
input_data_reshape = tf.reshape(input_data, [1, -1, input_size])
lstm_layer = tflearn.BasicLSTMCell(num_units=num_of_units,
state_is_tuple=True)
lstm_state_size = tuple([[1, x] for x in lstm_layer.state_size])
initial_lstm_state = _rnn_cell.LSTMStateTuple(
tf.placeholder(tf.float32,
shape=lstm_state_size[0],
name='initial_lstm_state1'),
tf.placeholder(tf.float32,
shape=lstm_state_size[1],
name='initial_lstm_state2'))
sequence_length = tf.reshape(batch_size, [1])
lstm_output, new_lstm_state = tf.nn.dynamic_rnn(
lstm_layer,
input_data_reshape,
initial_state=initial_lstm_state,
sequence_length=sequence_length,
time_major=False,
scope=scope)
lstm_output_reshape = tf.reshape(lstm_output, [-1, num_of_units])
self.hiddens.append(lstm_output_reshape)
self.num_of_hidden_layers += 1
return lstm_output_reshape, initial_lstm_state, new_lstm_state
def state_size(self):
return (_rnn_cell.LSTMStateTuple(self._num_units, self._num_units)
if self._state_is_tuple else 2 * self._num_units)
def state_size(self):
from tensorflow.python.ops.nn import rnn_cell
return rnn_cell.LSTMStateTuple(c=self.n_hidden, h=self.n_hidden)
def state_size(self):
return rnn_cell.LSTMStateTuple(self.num_units, self.num_units)
def call(self, inputs, state):
"""
:param (tf.Tensor,tf.Tensor) inputs: each (batch, num_units)
:param rnn_cell.LSTMStateTuple state: (batch, num_units)
:return: output, new_state
"""
_x1, _x2 = inputs
assert isinstance(_x1, tf.Tensor) and isinstance(_x2, tf.Tensor), "inputs %r unexpected" % (inputs,)
if not isinstance(state, rnn_cell.LSTMStateTuple):
state = rnn_cell.LSTMStateTuple(*tuple(state))
assert isinstance(state, rnn_cell.LSTMStateTuple)
with tf.name_scope("prepare"):
# All internal steps performed with moved feature axis, should be faster.
x1 = self.linear.move_feature_axis(_x1, old_axis=-1, new_axis=self.linear.mul_feature_axis)
x2 = self.linear.move_feature_axis(_x2, old_axis=-1, new_axis=self.linear.mul_feature_axis)
h = self.linear.move_feature_axis(state.h, old_axis=-1, new_axis=self.linear.mul_feature_axis)
c = self.linear.move_feature_axis(state.c, old_axis=-1, new_axis=self.linear.mul_feature_axis)
batch_axis = self.linear.mul_feature_axis - 1
if self.rec_dropout_h:
h = self.dropout(h, drop_rate=self.rec_dropout_h, inside_loop=True, batch_axis=batch_axis)
with tf.variable_scope("step0"):
dim = self.num_units
h = self.linear(h, dim, with_bias=False, feature_axis=self.linear.mul_feature_axis)
h *= x1
with tf.variable_scope("step1"):
h = self.linear(h, dim, with_bias=False, feature_axis=self.linear.mul_feature_axis)
h += x2
h = tf.nn.relu(h)
for step in range(2, self.depth):
with tf.variable_scope('step%i' % step):
dim = self.num_units
h = self.linear(h, dim, feature_axis=self.linear.mul_feature_axis)
h = tf.nn.relu(h)
h.set_shape((dim if self.linear.mul_feature_axis == 0 else None, None))
with tf.variable_scope("gating"):
with tf.variable_scope("cell_in"):
cell_in = self.linear(h, self.num_units, feature_axis=self.linear.mul_feature_axis)
if self.rec_dropout_u:
cell_in = self.dropout(cell_in, drop_rate=self.rec_dropout_u, inside_loop=True, batch_axis=batch_axis)
cell_in = tf.tanh(cell_in)
with tf.variable_scope("gate_in"):
gate_in = self.linear(h, self.num_units, feature_axis=self.linear.mul_feature_axis)
gate_in = tf.sigmoid(gate_in)
with tf.variable_scope("gate_forget"):
gate_forget = self.linear(
h, self.num_units, bias_init=self.forget_gate_bias_init, feature_axis=self.linear.mul_feature_axis)
gate_forget = tf.sigmoid(gate_forget)
with tf.variable_scope("gate_out"):
gate_out = self.linear(h, self.num_units, feature_axis=self.linear.mul_feature_axis)
gate_out = tf.sigmoid(gate_out)
cell = c * gate_forget + cell_in * gate_in
out = tf.tanh(cell) * gate_out
# Move feature axis back to where it is expected.
cell = self.linear.move_feature_axis(cell, old_axis=self.linear.mul_feature_axis, new_axis=-1)
out = self.linear.move_feature_axis(out, old_axis=self.linear.mul_feature_axis, new_axis=-1)
cell.set_shape((None, self.num_units))
out.set_shape((None, self.num_units))
assert out.get_shape().dims[0].value == cell.get_shape().dims[0].value == _x1.get_shape().dims[0].value, 'b.dim'
return out, rnn_cell.LSTMStateTuple(c=cell, h=out)
