get rid of dynamic rnn... hew

Author: carpedm20

config.py

@@ -47,8 +47,8 @@ def add_argument_group(name):
 
 # Misc
 misc_arg = add_argument_group('Misc')
-misc_arg.add_argument('--log_step', type=int, default=20, help='')
-misc_arg.add_argument('--num_log_samples', type=int, default=2, help='')
+misc_arg.add_argument('--log_step', type=int, default=50, help='')
+misc_arg.add_argument('--num_log_samples', type=int, default=3, help='')
 misc_arg.add_argument('--log_level', type=str, default='INFO', choices=['INFO', 'DEBUG', 'WARN'], help='')
 misc_arg.add_argument('--log_dir', type=str, default='logs')
 misc_arg.add_argument('--data_dir', type=str, default='data')

data_loader.py

@@ -35,6 +35,12 @@ def generate_one_example(n_nodes, rng):
  solutions = solve_tsp_dynamic(nodes)
  return nodes, solutions
 
+def pad(x, max_length):
+ shape = x.shape
+ pad_length = max_length - shape[0]
+ shape[0] = pad_length
+ return np.concatenate([x, x.np.zeros(shape)])
+
 class TSPDataLoader(object):
  def __init__(self, config, rng=None):
  self.config = config
@@ -78,10 +84,10 @@ def _create_input_queue(self, queue_capacity_factor=16):
  min_after_dequeue = 1000
  capacity = min_after_dequeue + 3 * self.batch_size
 
- self.queue_ops[name] = tf.PaddingFIFOQueue(
+ self.queue_ops[name] = tf.FIFOQueue(
  capacity=capacity,
  dtypes=[tf.float32, tf.int32],
- shapes=[[None, 2,], [None]],
+ shapes=[[self.max_length, 2,], [self.max_length]],
  name="fifo_{}".format(name))
  self.enqueue_ops[name] = \
  self.queue_ops[name].enqueue([self.input_ops[name], self.target_ops[name]])
@@ -136,12 +142,13 @@ def _maybe_generate_and_save(self):
  if not os.path.exists(path):
  tf.logging.info("Creating {} for [{}]".format(path, self.task))
 
- x, y = [], []
+ x = np.zeros([num, self.max_length, 2], dtype=np.float32)
+ y = np.zeros([num, self.max_length], dtype=np.int32)
  for i in trange(num, desc="Create {} data".format(name)):
  n_nodes = self.rng.randint(self.min_length, self.max_length+ 1)
  nodes, res = generate_one_example(n_nodes, self.rng)
- x.append(nodes)
- y.append(res)
+ x[i,:len(nodes)] = nodes
+ y[i,:len(res)] = res
 
  np.savez(path, x=x, y=y)
  self.data[name] = TSP(x=x, y=y, name=name)

layers.py

@@ -11,9 +11,10 @@
 
 def decoder_rnn(cell, inputs,
  enc_outputs, enc_final_states,
- seq_length, hidden_dim, num_glimpse,
- max_dec_length, batch_size, is_train,
- end_of_sequence_id=0, initializer=None):
+ seq_length, hidden_dim,
+ num_glimpse, batch_size, is_train,
+ end_of_sequence_id=0, initializer=None,
+ first_decoder_input=None):
  with tf.variable_scope("decoder_rnn") as scope:
  def attention(ref, query, with_softmax, scope="attention"):
  with tf.variable_scope(scope):
@@ -41,40 +42,32 @@ def glimpse(ref, query, scope="glimpse"):
  return tf.reduce_sum(alignments * ref, [1])
 
  def output_fn(ref, query, num_glimpse):
- for idx in range(num_glimpse):
- query = glimpse(ref, query, "glimpse_{}".format(idx))
- return attention(ref, query, with_softmax=False, scope="attention")
-
- maximum_length = tf.convert_to_tensor(max_dec_length, tf.int32)
- def decoder_fn_inference(
- time, cell_state, cell_input, cell_output, context_state):
- if context_state is None:
- context_state = tf.zeros([batch_size,], dtype=tf.int32)
-
- if cell_output is None:
- # time == 0
- cell_state = enc_final_states
- done = tf.zeros([batch_size,], dtype=tf.bool)
+ if query is None:
+ return tf.zeros([11], tf.float32) # only used for shape inference
  else:
- output_logit = output_fn(enc_outputs, cell_output, num_glimpse)
+ for idx in range(num_glimpse):
+ query = glimpse(ref, query, "glimpse_{}".format(idx))
+ return attention(ref, query, with_softmax=False, scope="attention")
 
- sampled_idx = tf.squeeze(
- tf.cast(tf.multinomial(output_logit, 1), tf.int32), -1)
- done = tf.equal(sampled_idx, end_of_sequence_id)
-
- cell_input = tf.stop_gradient(
- tf.gather_nd(enc_outputs, index_matrix_to_pairs(sampled_idx)))
- conext_state = tf.stack([context_state, sampled_idx], 0)
-
- done = tf.cond(tf.greater(time, maximum_length),
- lambda: tf.ones([batch_size,], dtype=tf.bool),
- lambda: done)
- return (done, cell_state, cell_input, cell_output, context_state)
+ def input_fn(sampled_idx):
+ return tf.stop_gradient(
+ tf.gather_nd(enc_outputs, index_matrix_to_pairs(sampled_idx)))
 
  if is_train:
  decoder_fn = simple_decoder_fn_train(enc_final_states)
  else:
- decoder_fn = decoder_fn_inference
+ def decoder_fn(time, cell_state, cell_input, cell_output, context_state):
+ cell_output = output_fn(enc_outputs, cell_output, num_glimpse)
+
+ if cell_state is None:
+ cell_state = enc_final_states
+ next_input = tf.squeeze(first_decoder_input, 1)
+ done = tf.zeros([batch_size,], dtype=tf.bool)
+ else:
+ sampled_idx = tf.cast(tf.argmax(cell_output, 1), tf.int32)
+ next_input = input_fn(sampled_idx)
+ done = tf.equal(sampled_idx, end_of_sequence_id)
+ return (done, cell_state, next_input, cell_output, context_state)
 
  outputs, final_state, final_context_state = \
  dynamic_rnn_decoder(cell, decoder_fn, inputs=inputs,

model.py

@@ -122,20 +122,23 @@ def _build_model(self):
 
  if self.use_terminal_symbol:
  # 0 index indicates terminal
- first_decoder_input = tf.expand_dims(trainable_initial_state(
+ self.first_decoder_input = tf.expand_dims(trainable_initial_state(
  batch_size, self.hidden_dim, name="first_decoder_input"), 1)
  self.enc_outputs = tf.concat_v2(
- [first_decoder_input, self.enc_outputs], axis=1)
+ [self.first_decoder_input, self.enc_outputs], axis=1)
 
  with tf.variable_scope("dencoder"):
+ self.idx_pairs = index_matrix_to_pairs(self.dec_targets)
+ self.embeded_dec_inputs = tf.stop_gradient(
+ tf.gather_nd(self.enc_outputs, self.idx_pairs))
+
  if self.use_terminal_symbol:
  tiled_zero_idxs = tf.tile(tf.zeros(
  [1, 1], dtype=tf.int32), [batch_size, 1], name="tiled_zero_idxs")
  self.dec_targets = tf.concat_v2([self.dec_targets, tiled_zero_idxs], axis=1)
 
- self.idx_pairs = index_matrix_to_pairs(self.dec_targets)
- self.embeded_dec_inputs = tf.stop_gradient(
- tf.gather_nd(self.enc_outputs, self.idx_pairs))
+ self.embeded_dec_inputs = tf.concat_v2(
+ [self.first_decoder_input, self.embeded_dec_inputs], axis=1)
 
  self.dec_cell = LSTMCell(
  self.hidden_dim,
@@ -148,41 +151,51 @@ def _build_model(self):
  self.dec_pred_logits, _, _ = decoder_rnn(
  self.dec_cell, self.embeded_dec_inputs, 
  self.enc_outputs, self.enc_final_states,
- self.dec_seq_length, self.hidden_dim, self.num_glimpse,
- self.max_dec_length, batch_size, is_train=True,
+ self.dec_seq_length, self.hidden_dim,
+ self.num_glimpse, batch_size, is_train=True,
  initializer=self.initializer)
  self.dec_pred_prob = tf.nn.softmax(
  self.dec_pred_logits, 2, name="dec_pred_prob")
  self.dec_pred = tf.argmax(
  self.dec_pred_logits, 2, name="dec_pred")
 
  with tf.variable_scope("dencoder", reuse=True):
- self.dec_inference_outputs, _, self.dec_inference = decoder_rnn(
- self.dec_cell, first_decoder_input,
+ self.dec_inference_logits, _, _ = decoder_rnn(
+ self.dec_cell, None,
  self.enc_outputs, self.enc_final_states,
- self.dec_seq_length, self.hidden_dim, self.num_glimpse,
- self.max_dec_length, batch_size, is_train=False,
- initializer=self.initializer)
+ self.dec_seq_length, self.hidden_dim,
+ self.num_glimpse, batch_size, is_train=False,
+ initializer=self.initializer, first_decoder_input=self.first_decoder_input)
  self.dec_inference_prob = tf.nn.softmax(
- self.dec_inference_outputs, 2, name="dec_inference_prob")
+ self.dec_inference_logits, 2, name="dec_inference_logits")
+ self.dec_inference = tf.argmax(
+ self.dec_inference_logits, 2, name="dec_inference")
 
  def _build_optim(self):
  losses = tf.nn.sparse_softmax_cross_entropy_with_logits(
  labels=self.dec_targets, logits=self.dec_pred_logits)
+ inference_losses = tf.nn.sparse_softmax_cross_entropy_with_logits(
+ labels=self.dec_targets, logits=self.dec_inference_logits)
 
  def apply_mask(op):
  length = tf.cast(op[:1], tf.int32)
  loss = op[1:]
  return tf.multiply(loss, tf.ones(length, dtype=tf.float32))
 
- batch_loss = tf.div(tf.reduce_sum(tf.multiply(losses, self.mask)),
- tf.reduce_sum(self.mask), name="batch_loss")
+ batch_loss = tf.div(
+ tf.reduce_sum(tf.multiply(losses, self.mask)),
+ tf.reduce_sum(self.mask), name="batch_loss")
+
+ batch_inference_loss = tf.div(
+ tf.reduce_sum(tf.multiply(losses, self.mask)),
+ tf.reduce_sum(self.mask), name="batch_inference_loss")
 
  tf.losses.add_loss(batch_loss)
  total_loss = tf.losses.get_total_loss()
 
  self.total_loss = total_loss
  self.target_cross_entropy_losses = losses
+ self.total_inference_loss = batch_inference_loss
 
  self.lr = tf.train.exponential_decay(
  self.lr_start, self.global_step, self.lr_decay_step,

trainer.py

@@ -76,12 +76,13 @@ def train(self):
 
  if result['step'] % self.log_step == 0:
  fetch = {
- 'loss': self.model.total_loss,
- 'pred': self.model.dec_inference,
- 'targets': self.model.dec_targets,
+ 'test l': self.model.total_inference_loss,
+ 'test x': self.model.dec_inference,
+ 'test y': self.model.dec_targets,
  }
  result = self.model.test(self.sess, fetch, self.summary_writer)
 
+ tf.logging.info("")
  tf.logging.info("loss: {}".format(result['loss']))
  for idx in range(self.num_log_samples):
  tf.logging.info("preds: {}".format(result['preds'][idx]))