def __init__(self, layer_idx, is_training, tau, prev_neurons=None):
rnn_config = get_rnn_config()
self.train_config = get_train_config()
self.layer_config = rnn_config['layer_configs'][layer_idx]
if prev_neurons is None:
self.w_shape = (rnn_config['layout'][layer_idx - 1],
rnn_config['layout'][layer_idx])
else:
self.w_shape = (prev_neurons, rnn_config['layout'][layer_idx])
self.b_shape = (1, self.w_shape[1])
self.is_training = is_training
# Activation summaries and specific neurons to gather individual histograms
self.acts = dict()
self.act_neurons = np.random.choice(
range(self.b_shape[1]),
size=(get_info_config()['tensorboard']['single_acts'], ),
replace=False)
if self.train_config['batchnorm'][
'type'] == 'batch' and 'fc' in self.train_config['batchnorm'][
'modes']:
self.bn_s_x = get_batchnormalizer()
self.bn_b_x = get_batchnormalizer()
with tf.variable_scope(self.layer_config['var_scope']):
var_keys = ['w', 'b']
self.weights = Weights(var_keys, self.layer_config, self.w_shape,
self.b_shape, tau)
def __init__(self, rnn_config, train_config, layer_idx, is_training):
self.rnn_config = rnn_config
self.train_config = train_config
self.is_training = is_training
self.layer_config = rnn_config['layer_configs'][layer_idx]
self.w_shape = (rnn_config['layout'][layer_idx-1], rnn_config['layout'][layer_idx])
self.b_shape = (1, self.w_shape[1])
if 'fc' in self.train_config['batchnorm']['modes']:
self.bn_x = get_batchnormalizer()
with tf.variable_scope(self.layer_config['var_scope']):
w_init_vals, b_init_vals = generate_init_values(self.layer_config['init_config'], self.w_shape, self.b_shape)
w_initializer = tf.constant_initializer(w_init_vals)
b_initializer = tf.constant_initializer(b_init_vals)
self.w = tf.get_variable(name='w', shape=self.w_shape, initializer=w_initializer)
self.b = tf.get_variable(name='b', shape=self.b_shape, initializer=b_initializer)
if self.layer_config['regularization']['mode'] == 'l2':
self.layer_loss = self.layer_config['regularization']['strength'] * tf.nn.l2_loss(self.w)
else:
self.layer_loss = 0
def create_var_fp(self, x, time_idx, init, init_cell=None):
if init:
cell_shape = (tf.shape(x)[0], self.b_shape[1])
if init_cell is not None:
self.weights.tensor_dict['cs'] = init_cell
else:
self.weights.tensor_dict['cs'] = tf.zeros(cell_shape)
self.weights.tensor_dict['co'] = tf.zeros(cell_shape)
co = self.weights.tensor_dict['co']
if self.train_config['batchnorm']['type'] == 'batch':
bn_idx = min(time_idx, self.train_config['batchnorm']['tau'] - 1)
if 'x' in self.train_config['batchnorm']['modes']:
if len(self.bn_s_x) == bn_idx:
self.bn_s_x.append(get_batchnormalizer())
x = self.bn_s_x[bn_idx](x, self.is_training)
if 'h' in self.train_config['batchnorm']['modes'] and bn_idx > 0:
if len(self.bn_s_h) == bn_idx - 1:
self.bn_s_h.append(get_batchnormalizer())
co = self.bn_s_h[bn_idx - 1](co, self.is_training)
x = tf.concat([x, co], axis=1)
i_act = tf.matmul(
x, self.weights.var_dict['wi']) + self.weights.var_dict['bi']
c_act = tf.matmul(
x, self.weights.var_dict['wc']) + self.weights.var_dict['bc']
o_act = tf.matmul(
x, self.weights.var_dict['wo']) + self.weights.var_dict['bo']
if self.train_config['batchnorm']['type'] == 'layer':
i_act = tf.contrib.layers.layer_norm(i_act)
c_act = tf.contrib.layers.layer_norm(c_act)
o_act = tf.contrib.layers.layer_norm(o_act)
if init:
for act_type, act in zip(['i', 'c', 'o'], [i_act, c_act, o_act]):
self.acts[act_type] = act
for neuron_idc in range(len(self.act_neurons)):
self.acts[act_type + '_' + str(neuron_idc)] = tf.slice(
act, begin=(0, neuron_idc), size=(-1, 1))
else:
for act_type, act in zip(['i', 'c', 'o'], [i_act, c_act, o_act]):
self.acts[act_type] = tf.concat([act, self.acts[act_type]],
axis=0)
for neuron_idc in range(len(self.act_neurons)):
self.acts[act_type + '_' + str(neuron_idc)] = \
tf.concat([tf.slice(act, begin=(0, neuron_idc), size=(-1, 1)),
self.acts[act_type + '_' + str(neuron_idc)]], axis=0)
if 'i' in self.rnn_config['act_disc']:
i = tf.cast(tf.cast(
tf.sigmoid(i_act) * self.rnn_config['act_bins'],
dtype=tf.int32),
dtype=tf.float32) / self.rnn_config['act_bins']
if get_info_config()['cell_access']:
self.cell_access_mat.append(i)
else:
i = tf.sigmoid(i_act)
f = 1. - i
if 'c' in self.rnn_config['act_disc']:
c = tf.cast(tf.cast(
tf.sigmoid(c_act) * self.rnn_config['act_bins'],
dtype=tf.int32),
dtype=tf.float32) * 2 / self.rnn_config['act_bins'] - 1
else:
c = tf.tanh(c_act)
if 'o' in self.rnn_config['act_disc']:
o = tf.cast(tf.cast(
tf.sigmoid(o_act) * self.rnn_config['act_bins'],
dtype=tf.int32),
dtype=tf.float32) / self.rnn_config['act_bins']
else:
o = tf.sigmoid(o_act)
self.weights.tensor_dict['cs'] = tf.multiply(
f, self.weights.tensor_dict['cs']) + tf.multiply(i, c)
if 'i' in self.rnn_config['act_disc']:
self.weights.tensor_dict['co'] = tf.multiply(
o, self.weights.tensor_dict['cs'])
else:
self.weights.tensor_dict['co'] = tf.multiply(
o, tf.tanh(self.weights.tensor_dict['cs']))
return self.weights.tensor_dict['co'], self.weights.tensor_dict['cs']
def create_g_sampling_pass(self,
x,
mod_layer_config,
time_idx,
init,
second_arm_pass=False,
data_key=None,
init_cell=None):
#if len(self.rnn_config['act_disc']) != 0:
#raise Exception('classic reparametrization does not work with discrete activations')
if init:
self.weights.create_tensor_samples(second_arm_pass=second_arm_pass,
data_key=data_key)
cell_shape = (tf.shape(x)[0], self.b_shape[1])
if init_cell is not None:
self.weights.tensor_dict['cs'] = init_cell
else:
self.weights.tensor_dict['cs'] = tf.zeros(cell_shape)
self.weights.tensor_dict['co'] = tf.zeros(cell_shape)
co = self.weights.tensor_dict['co']
if self.train_config['batchnorm']['type'] == 'batch':
bn_idx = min(time_idx, self.train_config['batchnorm']['tau'] - 1)
if 'x' in self.train_config['batchnorm']['modes']:
if len(self.bn_b_x) == bn_idx:
self.bn_b_x.append(get_batchnormalizer())
x = self.bn_b_x[bn_idx](x, self.is_training)
if 'h' in self.train_config['batchnorm']['modes'] and bn_idx > 0:
if len(self.bn_b_h) == bn_idx - 1:
self.bn_b_h.append(get_batchnormalizer())
co = self.bn_b_h[bn_idx - 1](co, self.is_training)
x = tf.concat([x, co], axis=1)
i_act = tf.matmul(
x, self.weights.tensor_dict['wi']) + self.weights.tensor_dict['bi']
c_act = tf.matmul(
x, self.weights.tensor_dict['wc']) + self.weights.tensor_dict['bc']
o_act = tf.matmul(
x, self.weights.tensor_dict['wo']) + self.weights.tensor_dict['bo']
if self.train_config['batchnorm']['type'] == 'layer':
i_act = tf.contrib.layers.layer_norm(i_act)
c_act = tf.contrib.layers.layer_norm(c_act)
o_act = tf.contrib.layers.layer_norm(o_act)
if 'i' in self.rnn_config['act_disc']:
print(self.rnn_config['act_bins'])
i = disc_sigmoid(i_act, self.rnn_config['act_bins'])
else:
i = tf.sigmoid(i_act)
f = 1. - i
if 'c' in self.rnn_config['act_disc']:
c = disc_tanh(c_act, self.rnn_config['act_bins'])
else:
c = tf.tanh(c_act)
if 'o' in self.rnn_config['act_disc']:
o = disc_sigmoid(o_act, self.rnn_config['act_bins'])
else:
o = tf.sigmoid(o_act)
self.weights.tensor_dict['cs'] = tf.multiply(
f, self.weights.tensor_dict['cs']) + tf.multiply(i, c)
self.weights.tensor_dict['co'] = tf.multiply(
o, tf.tanh(self.weights.tensor_dict['cs']))
return self.weights.tensor_dict['co'], self.weights.tensor_dict['cs']
def create_l_sampling_pass(self,
x,
mod_layer_config,
time_idx,
init,
init_cell=None):
if init:
cell_shape = (tf.shape(x)[0], self.b_shape[1])
if init_cell is not None:
self.weights.tensor_dict['cs'] = init_cell
else:
self.weights.tensor_dict['cs'] = tf.zeros(cell_shape)
self.weights.tensor_dict['co'] = tf.zeros(cell_shape)
co = self.weights.tensor_dict['co']
if self.train_config['batchnorm']['type'] == 'batch':
bn_idx = min(time_idx, self.train_config['batchnorm']['tau'] - 1)
if 'x' in self.train_config['batchnorm']['modes']:
if len(self.bn_b_x) == bn_idx:
self.bn_b_x.append(get_batchnormalizer())
x = self.bn_b_x[bn_idx](x, self.is_training)
if 'h' in self.train_config['batchnorm']['modes'] and bn_idx > 0:
if len(self.bn_b_h) == bn_idx - 1:
self.bn_b_h.append(get_batchnormalizer())
co = self.bn_b_h[bn_idx - 1](co, self.is_training)
x = tf.concat([x, co], axis=1)
if 'i' in self.rnn_config['act_disc']:
i = self.weights.sample_activation(
'wi', 'bi', x, 'sig', init,
self.train_config['batchnorm']['type'] == 'layer')
else:
a_i = self.weights.sample_activation(
'wi', 'bi', x, None, init,
self.train_config['batchnorm']['type'] == 'layer')
i = tf.sigmoid(a_i)
f = 1. - i
if 'c' in self.rnn_config['act_disc']:
c = self.weights.sample_activation(
'wc', 'bc', x, 'tanh', init,
self.train_config['batchnorm']['type'] == 'layer')
else:
a_c = self.weights.sample_activation(
'wc', 'bc', x, None, init,
self.train_config['batchnorm']['type'] == 'layer')
c = tf.tanh(a_c)
if 'o' in self.rnn_config['act_disc']:
o = self.weights.sample_activation(
'wo', 'bo', x, 'sig', init,
self.train_config['batchnorm']['type'] == 'layer')
else:
a_o = self.weights.sample_activation(
'wo', 'bo', x, None, init,
self.train_config['batchnorm']['type'] == 'layer')
o = tf.sigmoid(a_o)
self.weights.tensor_dict['cs'] = tf.multiply(
f, self.weights.tensor_dict['cs']) + tf.multiply(i, c)
if 'i' in self.rnn_config['act_disc']:
self.weights.tensor_dict['co'] = tf.multiply(
self.weights.tensor_dict['cs'], o)
else:
self.weights.tensor_dict['co'] = tf.multiply(
tf.tanh(self.weights.tensor_dict['cs']), o)
return self.weights.tensor_dict['co'], self.weights.tensor_dict['cs']
def create_forward_pass(self, layer_input, mod_layer_config, init, time_idx):
if init:
cell_shape = (tf.shape(layer_input)[0], self.b_shape[1])
self.cell_state = tf.zeros(cell_shape)
self.cell_output = tf.zeros(cell_shape)
co = self.cell_output
if self.train_config['batchnorm']['type'] == 'batch':
bn_idx = min(time_idx, self.train_config['batchnorm']['tau'] - 1)
if 'x' in self.train_config['batchnorm']['modes']:
if len(self.bn_x) == bn_idx:
self.bn_x.append(get_batchnormalizer())
layer_input = self.bn_x[bn_idx](layer_input, self.is_training)
if 'h' in self.train_config['batchnorm']['modes'] and bn_idx > 0:
if len(self.bn_h) == bn_idx - 1:
self.bn_h.append(get_batchnormalizer())
co = self.bn_h[bn_idx - 1](self.cell_output, self.is_training)
x = tf.concat([layer_input, co], axis=1)
if self.rnn_config['weight_type'] == 'continuous':
i_act = self.bi + tf.matmul(x, self.wi, name='i')
o_act = self.bo + tf.matmul(x, self.wo, name='o')
c_act = self.bc + tf.matmul(x, self.wc, name='c')
elif self.rnn_config['weight_type'] == 'ternary':
i_act = self.bi + tf.matmul(x, ternarize_weight(self.wi), name='i')
o_act = self.bo + tf.matmul(x, ternarize_weight(self.wo), name='o')
c_act = self.bc + tf.matmul(x, ternarize_weight(self.wc), name='c')
elif self.rnn_config['weight_type'] == 'binary':
i_act = self.bi + tf.matmul(x, binarize_weight(self.wi), name='i')
o_act = self.bo + tf.matmul(x, binarize_weight(self.wo), name='o')
c_act = self.bc + tf.matmul(x, binarize_weight(self.wc), name='c')
else:
raise Exception('weight type not understood')
if self.train_config['batchnorm']['type'] == 'layer':
i_act = tf.contrib.layers.layer_norm(i_act)
o_act = tf.contrib.layers.layer_norm(o_act)
c_act = tf.contrib.layers.layer_norm(c_act)
i = tf.sigmoid(i_act)
o = tf.sigmoid(o_act)
c = tf.tanh(c_act)
if 'i' in self.rnn_config['discrete_acts']:
i = disc_sigmoid(i_act, self.rnn_config['n_bins'])
if 'c' in self.rnn_config['discrete_acts']:
c = disc_tanh(c_act, self.rnn_config['n_bins'])
if 'o' in self.rnn_config['discrete_acts']:
o = disc_sigmoid(o_act, self.rnn_config['n_bins'])
f = 1. - i
updated_state = tf.multiply(f, self.cell_state, name='f_cs') + tf.multiply(i, c, name='i_cs')
updated_output = tf.multiply(o, tf.tanh(updated_state))
if mod_layer_config['regularization']['mode'] == 'zoneout':
state_mask = self.state_zoneout_dist.sample(tf.shape(self.cell_state))
output_mask = self.output_zoneout_dist.sample(tf.shape(self.cell_output))
self.cell_state = tf.multiply(state_mask, self.cell_state) + tf.multiply(1 - state_mask, updated_state)
self.cell_output = tf.multiply(output_mask, self.cell_output) + tf.multiply(1 - output_mask, updated_output)
else:
self.cell_state = updated_state
self.cell_output = updated_output
return self.cell_output