def fwd_fn(idx, label):
np.random.seed(1)
embedding_shape = (dataset_size, embedding_size)
embedding_initializer = np.random.normal(0, 1, embedding_shape).astype(
np.float32)
weights_shape = (embedding_size, embedding_size)
weights_initializer = np.random.normal(0, 1,
weights_shape).astype(np.float32)
with variable_scope.variable_scope("part1", use_resource=True):
embedding = variable_scope.get_variable(
"c",
dtype=np.float32,
initializer=embedding_initializer,
trainable=True)
weight = variable_scope.get_variable("w0",
dtype=np.float32,
initializer=weights_initializer,
trainable=True)
x = embedding_ops.embedding_lookup(embedding, idx)
x = math_ops.matmul(x, weight)
logits = math_ops.reduce_sum(x, axis=[-1])
loss = math_ops.reduce_mean(
nn.sparse_softmax_cross_entropy_with_logits(logits=logits,
labels=label))
return loss
def stage1(idx, label):
with variable_scope.variable_scope("stage1", use_resource=True):
embedding = variable_scope.get_variable(
"c",
shape=[10, 1216],
dtype=np.float32,
initializer=init_ops.constant_initializer(10.01),
trainable=True)
x = embedding_ops.embedding_lookup(embedding, idx)
return x, label
def _build_decoder_embedding(self, input, seq_len, scope_name):
'''
Args:
input: 2D-tensor, (batch_size, maxlen_tgt), target word index
seq_len: 1D-tensor, (batch_size,), valid target word length
scope_name: scope name
Returns:
x: 3D-tensor, (batch_size, maxlen_tgt, attention_dim)
loss_mask: 3D-tensor, (batch_size, maxlen_tgt-1, 1)
att_mask: 4D-tensor, (1, 1, maxlen_tgt-1, maxlen_tgt-1), constant
'''
with tf.compat.v1.variable_scope(scope_name,
dtype=self.dtype,
use_resource=True) as scope:
loss_mask = tf.sequence_mask(seq_len,
maxlen=self.config['maxlen_tgt'] - 1,
dtype=scope.dtype)
loss_mask = tf.expand_dims(loss_mask, axis=2)
embedding = tf.compat.v1.get_variable(
"embedding_table",
[self.config['vocab_size'], self.config['adim']],
scope.dtype,
initializer=tf.initializers.random_uniform(minval=0,
maxval=1.0,
dtype=scope.dtype),
trainable=self.training,
regularizer=self.kernel_regularizer)
# x = tf.nn.embedding_lookup(embedding, input)
x = embedding_ops.embedding_lookup(embedding, input)
# position embedding
_, max_len, dmodel = x.shape.as_list()
pos_emb = self._build_pos_embedding(max_len, dmodel, reverse=False)
x = math.sqrt(self.config['adim']) * x + pos_emb
if self.training:
if self.config['use_ipu_dropout']:
x = ipu.rand_ops.dropout(x,
rate=self.config['dropout_rate'])
else:
x = tf.nn.dropout(x, rate=self.config['dropout_rate'])
# subsequent_mask
index = tf.range(1, self.config['maxlen_tgt'], 1, dtype=tf.int32)
index = tf.reshape(index, (1, 1, -1))
att_mask = tf.sequence_mask(index, dtype=scope.dtype)
att_mask = (1.0 - att_mask) * self.mask_value
return x, loss_mask, att_mask
def fwd_fn(idx, label):
with variable_scope.variable_scope("part1", use_resource=True):
embedding = variable_scope.get_variable(
"c",
shape=[10, 1216],
dtype=np.float32,
initializer=init_ops.constant_initializer(10.01),
trainable=True)
x = embedding_ops.embedding_lookup(embedding, idx)
logits = math_ops.reduce_sum(x, axis=[-1])
loss = math_ops.reduce_mean(
nn.sparse_softmax_cross_entropy_with_logits(logits=logits,
labels=label))
return loss
def create_policy(*infeed_data):
"""Act according to current policy and generate action probability. """
dis_obs = list(infeed_data[:4])
cont_obs = list(infeed_data[4:8])
state_in = infeed_data[-1]
# Look up embedding for all the discrete obs
emb_lookup = []
with tf.variable_scope("popnn_lookup"):
for index, obs in enumerate(dis_obs):
emb_matrix = tf.get_variable(
f'emb_matrix{index}',
[DIS_OBS_CARDINALITY[index], DIS_OBS_EMB_SIZE[index]], DTYPE)
emb_lookup.append(
embedding_ops.embedding_lookup(emb_matrix,
obs,
name=f'emb_lookup{index}'))
# Clip some continuous observations
cont_obs[-1] = tf.clip_by_value(cont_obs[-1], -5.0, 5.0, name="clip")
# Concat groups of observations
obs_concat = []
for d_obs, c_obs in zip(emb_lookup, cont_obs):
obs_concat.append(tf.concat([d_obs, c_obs], axis=3, name="concat_obs"))
# Fully connected transformations
num_output = 8
obs_concat[-1] = Dense(num_output, dtype=DTYPE)(obs_concat[-1])
# Reduce max
obs_concat = [tf.reduce_max(obs, axis=2) for obs in obs_concat]
# Final concat of all the observations
lstm_input = tf.concat(obs_concat, axis=2, name="concat_all")
# LSTM layer
lstm_input = tf.transpose(
lstm_input, perm=[1, 0, 2],
name="pre_lstm_transpose") # PopnnLSTM uses time-major tensors
lstm_cell = rnn_ops.PopnnLSTM(num_units=LSTM_HIDDEN_SIZE,
dtype=DTYPE,
partials_dtype=DTYPE,
name="lstm")
lstm_output, state_out = lstm_cell(
lstm_input,
training=True,
initial_state=tf.nn.rnn_cell.LSTMStateTuple(state_in[:, 0],
state_in[:, 1]))
lstm_output = tf.transpose(lstm_output,
perm=[1, 0, 2],
name="post_lstm_transpose")
logits = Dense(NUM_ACTIONS, name="logits", dtype=DTYPE)(lstm_output)
log_prob = tf.nn.log_softmax(logits, name="prob")
# make action selection op (outputs int actions, sampled from policy)
actions = tf.random.categorical(logits=tf.reshape(logits,
(-1, NUM_ACTIONS)),
num_samples=1)
actions = tf.reshape(actions, (args.batch_size, args.time_steps))
action_masks = tf.one_hot(actions, NUM_ACTIONS, dtype=DTYPE)
action_prob = tf.reduce_sum(action_masks * log_prob, axis=-1)
return action_prob