update pointer network from other repo

Author: carpedm20

config.py

@@ -14,7 +14,7 @@ def add_argument_group(name):
 
 # Network
 net_arg = add_argument_group('Network')
-net_arg.add_argument('--hidden_dim', type=int, default=128, help='')
+net_arg.add_argument('--hidden_dim', type=int, default=256, help='')
 net_arg.add_argument('--num_layers', type=int, default=1, help='')
 net_arg.add_argument('--input_dim', type=int, default=2, help='')
 net_arg.add_argument('--max_enc_length', type=int, default=None, help='')
@@ -26,7 +26,7 @@ def add_argument_group(name):
 
 # Data
 data_arg = add_argument_group('Data')
-data_arg.add_argument('--task', type=str, default='TSP')
+data_arg.add_argument('--task', type=str, default='tsp')
 data_arg.add_argument('--batch_size', type=int, default=128)
 data_arg.add_argument('--min_data_length', type=int, default=5)
 data_arg.add_argument('--max_data_length', type=int, default=10)
@@ -42,12 +42,13 @@ def add_argument_group(name):
 train_arg.add_argument('--lr_start', type=float, default=0.001, help='')
 train_arg.add_argument('--lr_decay_step', type=int, default=5000, help='')
 train_arg.add_argument('--lr_decay_rate', type=float, default=0.96, help='')
-train_arg.add_argument('--max_grad_norm', type=float, default=1.0, help='')
+train_arg.add_argument('--max_grad_norm', type=float, default=2.0, help='')
 train_arg.add_argument('--checkpoint_secs', type=int, default=300, help='')
 
 # Misc
 misc_arg = add_argument_group('Misc')
-misc_arg.add_argument('--log_step', type=int, default=20, 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

@@ -1,13 +1,26 @@
 # Most of the codes are from 
 # https://github.com/vshallc/PtrNets/blob/master/pointer/misc/tsp.py
 import os
+import re
+import zipfile
 import itertools
 import threading
 import numpy as np
-from tqdm import trange
+from tqdm import trange, tqdm
 from collections import namedtuple
 
 import tensorflow as tf
+from download import download_file_from_google_drive
+
+GOOGLE_DRIVE_IDS = {
+    'tsp5_train.zip': '0B2fg8yPGn2TCSW1pNTJMXzFPYTg',
+    'tsp10_train.zip': '0B2fg8yPGn2TCbHowM0hfOTJCNkU',
+    'tsp5-20_train.zip': '0B2fg8yPGn2TCTWNxX21jTDBGeXc',
+    'tsp50_train.zip': '0B2fg8yPGn2TCaVQxSl9ab29QajA',
+    'tsp20_test.txt': '0B2fg8yPGn2TCdF9TUU5DZVNCNjQ',
+    'tsp40_test.txt': '0B2fg8yPGn2TCcjFrYk85SGFVNlU',
+    'tsp50_test.txt.zip': '0B2fg8yPGn2TCUVlCQmQtelpZTTQ',
+}
 
 TSP = namedtuple('TSP', ['x', 'y', 'name'])
 
@@ -35,22 +48,34 @@ def generate_one_example(n_nodes, rng):
   solutions = solve_tsp_dynamic(nodes)
   return nodes, solutions
 
+def read_paper_dataset(paths, max_length):
+  x, y = [], []
+  for path in paths:
+    tf.logging.info("Read dataset {} which is used in the paper..".format(path))
+    length = max(re.findall('\d+', path))
+    with open(path) as f:
+      for l in tqdm(f):
+        inputs, outputs = l.split(' output ')
+        x.append(np.array(inputs.split(), dtype=np.float32).reshape([-1, 2]))
+        y.append(np.array(outputs.split(), dtype=np.int32)[:-1]) # skip the last one
+  return x, y
+
 class TSPDataLoader(object):
   def __init__(self, config, rng=None):
     self.config = config
     self.rng = rng
 
-    self.task = config.task
+    self.task = config.task.lower()
     self.batch_size = config.batch_size
     self.min_length = config.min_data_length
     self.max_length = config.max_data_length
 
     self.is_train = config.is_train
     self.use_terminal_symbol = config.use_terminal_symbol
+    self.random_seed = config.random_seed
 
     self.data_num = {}
     self.data_num['train'] = config.train_num
-    self.data_num['valid'] = config.valid_num
     self.data_num['test'] = config.test_num
 
     self.data_dir = config.data_dir
@@ -63,7 +88,13 @@ def __init__(self, config, rng=None):
     self.queue_ops, self.enqueue_ops = None, None
     self.x, self.y, self.seq_length, self.mask = None, None, None, None
 
-    self._maybe_generate_and_save()
+    paths = self.download_google_drive_file()
+    if len(paths) != 0:
+      self._maybe_generate_and_save(except_list=paths.keys())
+      for name, path in paths.items():
+        self.read_zip_and_update_data(path, name)
+    else:
+      self._maybe_generate_and_save()
     self._create_input_queue()
 
   def _create_input_queue(self, queue_capacity_factor=16):
@@ -78,11 +109,13 @@ 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.RandomShuffleQueue(
           capacity=capacity,
+          min_after_dequeue=min_after_dequeue,
           dtypes=[tf.float32, tf.int32],
-          shapes=[[None, 2,], [None]],
-          name="fifo_{}".format(name))
+          shapes=[[self.max_length, 2,], [self.max_length]],
+          seed=self.random_seed,
+          name="random_queue_{}".format(name))
       self.enqueue_ops[name] = \
           self.queue_ops[name].enqueue([self.input_ops[name], self.target_ops[name]])
 
@@ -127,21 +160,26 @@ def stop_input_queue(self):
     self.coord.request_stop()
     self.coord.join(threads)
 
-  def _maybe_generate_and_save(self):
+  def _maybe_generate_and_save(self, except_list=[]):
     self.data = {}
 
     for name, num in self.data_num.items():
+      if name in except_list:
+        tf.logging.info("Skip creating {} because of given except_list {}".format(name, except_list))
+        continue
       path = self.get_path(name)
 
       if not os.path.exists(path):
         tf.logging.info("Creating {} for [{}]".format(path, self.task))
 
-        x, y = [], []
-        for i in trange(num, desc="Create {} data".format(name)):
+        x = np.zeros([num, self.max_length, 2], dtype=np.float32)
+        y = np.zeros([num, self.max_length], dtype=np.int32)
+
+        for idx 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[idx,:len(nodes)] = nodes
+          y[idx,:len(res)] = res
 
         np.savez(path, x=x, y=y)
         self.data[name] = TSP(x=x, y=y, name=name)
@@ -154,3 +192,50 @@ def get_path(self, name):
     return os.path.join(
         self.data_dir, "{}_{}={}.npz".format(
             self.task_name, name, self.data_num[name]))
+
+  def download_google_drive_file(self):
+    paths = {}
+    for mode in ['train', 'test']:
+      candidates = []
+      candidates.append(
+          '{}{}_{}'.format(self.task, self.max_length, mode))
+      candidates.append(
+          '{}{}-{}_{}'.format(self.task, self.min_length, self.max_length, mode))
+
+      for key in candidates:
+        for search_key in GOOGLE_DRIVE_IDS.keys():
+          if search_key.startswith(key):
+            path = os.path.join(self.data_dir, search_key)
+            tf.logging.info("Download dataset of the paper to {}".format(path))
+
+            if not os.path.exists(path):
+              download_file_from_google_drive(GOOGLE_DRIVE_IDS[search_key], path)
+              if path.endswith('zip'):
+                with zipfile.ZipFile(path, 'r') as z:
+                  z.extractall(self.data_dir)
+            paths[mode] = path
+
+    tf.logging.info("Can't found dataset from the paper!")
+    return paths
+
+  def read_zip_and_update_data(self, path, name):
+    if path.endswith('zip'):
+      filenames = zipfile.ZipFile(path).namelist()
+      paths = [os.path.join(self.data_dir, filename) for filename in filenames]
+    else:
+      paths = [path]
+
+    x_list, y_list = read_paper_dataset(paths, self.max_length)
+
+    x = np.zeros([len(x_list), self.max_length, 2], dtype=np.float32)
+    y = np.zeros([len(y_list), self.max_length], dtype=np.int32)
+
+    for idx, (nodes, res) in enumerate(tqdm(zip(x_list, y_list))):
+      x[idx,:len(nodes)] = nodes
+      y[idx,:len(res)] = res
+
+    if self.data is None:
+      self.data = {}
+
+    tf.logging.info("Update [{}] data with {} used in the paper".format(name, path))
+    self.data[name] = TSP(x=x, y=y, name=name)

download.py

@@ -0,0 +1,34 @@
+# Code based on
+# http://stackoverflow.com/questions/25010369/wget-curl-large-file-from-google-drive/39225039#39225039
+import requests
+from tqdm import tqdm
+
+def download_file_from_google_drive(id, destination):
+  URL = "https://docs.google.com/uc?export=download"
+
+  session = requests.Session()
+
+  response = session.get(URL, params = { 'id' : id }, stream = True)
+  token = get_confirm_token(response)
+
+  if token:
+    params = { 'id' : id, 'confirm' : token }
+    response = session.get(URL, params = params, stream = True)
+
+  save_response_content(response, destination)  
+  return True
+
+def get_confirm_token(response):
+  for key, value in response.cookies.items():
+    if key.startswith('download_warning'):
+      return value
+
+  return None
+
+def save_response_content(response, destination):
+  CHUNK_SIZE = 32768
+
+  with open(destination, "wb") as f:
+    for chunk in tqdm(response.iter_content(CHUNK_SIZE)):
+      if chunk: # filter out keep-alive new chunks
+        f.write(chunk)

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,
+                max_length=None):
   with tf.variable_scope("decoder_rnn") as scope:
     def attention(ref, query, with_softmax, scope="attention"):
       with tf.variable_scope(scope):
@@ -41,37 +42,37 @@ 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.TensorArray(tf.float32, size=maximum_length)
-
-      if cell_output is None:
-        # invariant tha this is time == 0
-        cell_state = enc_final_states
-        cell_input = inputs[:,0,:]
-        done = tf.zeros([batch_size,], dtype=tf.bool)
+      if query is None:
+        return tf.zeros([max_length], tf.float32) # only used for shape inference
       else:
-        output_logit = output_fn(enc_outputs, cell_output, num_glimpse)
-        sampled_idx = tf.multinomial(output_logit, 1)
+        for idx in range(num_glimpse):
+          query = glimpse(ref, query, "glimpse_{}".format(idx))
+        return attention(ref, query, with_softmax=False, scope="attention")
 
-        context_state.write(time, output_logit)
-        done = tf.squeeze(tf.equal(sampled_idx, end_of_sequence_id), -1)
-
-      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
+      maximum_length = tf.convert_to_tensor(max_length, tf.int32)
+
+      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 = cell_input
+          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)
+
+        done = tf.cond(tf.greater(time, maximum_length),
+          lambda: tf.ones([batch_size,], dtype=tf.bool),
+          lambda: done)
+        return (done, cell_state, next_input, cell_output, context_state)
 
     outputs, final_state, final_context_state = \
         dynamic_rnn_decoder(cell, decoder_fn, inputs=inputs,
@@ -111,8 +112,10 @@ def trainable_initial_state(batch_size, state_size,
 def index_matrix_to_pairs(index_matrix):
   # [[3,1,2], [2,3,1]] -> [[[0, 3], [1, 1], [2, 2]], 
   #                        [[0, 2], [1, 3], [2, 1]]]
-  replicated_first_indices = tf.tile(
-      tf.expand_dims(tf.range(tf.shape(index_matrix)[0]), dim=1), 
-      [1, tf.shape(index_matrix)[1]])
-  return tf.stack([replicated_first_indices, index_matrix], axis=2)
-
+  replicated_first_indices = tf.range(tf.shape(index_matrix)[0])
+  rank = len(index_matrix.get_shape())
+  if rank == 2:
+    replicated_first_indices = tf.tile(
+        tf.expand_dims(replicated_first_indices, dim=1),
+        [1, tf.shape(index_matrix)[1]])
+  return tf.stack([replicated_first_indices, index_matrix], axis=rank)