update logging and data generation

.gitignore

@@ -1,7 +1,7 @@
 # Data
 data/hand
 data/gaze
-data/*.png
+data/*
 samples
 outputs
 

config.py

@@ -26,10 +26,13 @@ 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)
+data_arg.add_argument('--train_num', type=int, default=1000000)
+data_arg.add_argument('--valid_num', type=int, default=1000)
+data_arg.add_argument('--test_num', type=int, default=1000)
 
 # Training / test parameters
 train_arg = add_argument_group('Training')
@@ -45,6 +48,7 @@ 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('--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')
 misc_arg.add_argument('--output_dir', type=str, default='outputs')

data_loader.py

@@ -1,8 +1,15 @@
-# Most of the codes are from https://github.com/vshallc/PtrNets/blob/master/pointer/misc/tsp.py
+# Most of the codes are from 
+# https://github.com/vshallc/PtrNets/blob/master/pointer/misc/tsp.py
 import os
 import itertools
+import threading
 import numpy as np
 from tqdm import trange
+from collections import namedtuple
+
+import tensorflow as tf
+
+TSP = namedtuple('TSP', ['x', 'y', 'name'])
 
 def length(x, y):
   return np.linalg.norm(np.asarray(x) - np.asarray(y))
@@ -23,47 +30,116 @@ def solve_tsp_dynamic(points):
   res = min([(A[d][0] + all_distances[0][d[1]], A[d][1]) for d in iter(A)])
   return np.asarray(res[1]) + 1
 
-def generate_one_example(n_nodes):
-  nodes = np.random.rand(n_nodes, 2)
-  res = solve_tsp_dynamic(nodes)
-  return nodes, res
-
-def generate_examples(num, n_min, n_max, desc=""):
-  examples = []
-  for i in trange(num, desc=desc):
-    n_nodes = np.random.randint(n_min, n_max + 1)
-    nodes, res = generate_one_example(n_nodes)
-    examples.append((nodes, res))
-  return examples
+def generate_one_example(n_nodes, rng):
+  nodes = rng.rand(n_nodes, 2).astype(np.float32)
+  solutions = solve_tsp_dynamic(nodes)
+  return nodes, solutions
 
 class TSPDataLoader(object):
   def __init__(self, config, rng=None):
     self.config = config
     self.rng = rng
 
     self.task = config.task
+    self.batch_size = config.batch_size
     self.min_length = config.min_data_length
     self.max_length = config.max_data_length
 
-    self.task_name = "{}_{}_{}".format(self.task, self.min_length, self.max_length)
-    self.npz_path = os.path.join(config.data_dir, "{}.npz".format(self.task_name))
+    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
+    self.task_name = "{}_({},{})".format(
+        self.task, self.min_length, self.max_length)
+
+    self.data = None
+    self.coord = None
+    self.input_ops, self.target_ops = None, None
+    self.queue_ops, self.enqueue_ops = None, None
 
     self._maybe_generate_and_save()
+    self._create_input_queue()
+
+  def _create_input_queue(self, queue_capacity_factor=16):
+    self.input_ops, self.target_ops = {}, {}
+    self.queue_ops, self.enqueue_ops = {}, {}
+
+    for name in self.data_num.keys():
+      self.input_ops[name] = tf.placeholder(tf.float32, shape=[None, None])
+      self.target_ops[name] = tf.placeholder(tf.int32, shape=[None])
+
+      min_after_dequeue = 5000
+      capacity = min_after_dequeue + 3 * self.batch_size
+
+      if self.is_training:
+        self.queue_ops[name] = tf.RandomShuffleQueue(
+            capacity=capacity,
+            min_after_dequeue=min_after_dequeue,
+            dtypes=[tf.float32, tf.int32],
+            name="random_{}".format(name))
+      else:
+        self.queue_ops[name] = tf.FIFOQueue(
+            capacity=capacity,
+            dtypes=[tf.float32, tf.int32],
+            name="fifo_{}".format(name))
+
+      self.enqueue_ops[name] = \
+          self.queue_ops[name].enqueue([self.input_ops[name], self.target_ops[name]])
+
+    tf.train.queue_runner.add_queue_runner(tf.train.queue_runner.QueueRunner(
+          values_queue, enqueue_ops))
+
+  def run_input_queue(self, sess):
+    threads = []
+    self.coord = tf.train.Coordinator()
+
+    for name in self.data_num.keys():
+      def load_and_enqueue(sess, name, input_ops, enqueue_ops, coord):
+        idx = 0
+        while not coord.should_stop():
+          feed_dict = {
+              input_ops[name]: self.data[name].x[idx],
+              target_ops[name]: self.data[name].y[idx],
+          }
+          sess.run(self.enqueue_ops[name], feed_dict=feed_dict)
+          idx += 1
+
+      args = (sess, name, self.input_ops, self.enqueue_ops, self.coord)
+      t = threading.Thread(target=load_and_enqueue, args=args)
+      t.start()
+      threads.append(t)
+      tf.logging.info("Thread start for [{}]".format(name))
+
+  def stop_input_queue(self):
+    self.coord.request_stop()
+    self.coord.join(threads)
 
   def _maybe_generate_and_save(self):
-    if not os.path.exists(self.npz_path):
-      print("[*] Creating dataset for {}".format(self.task))
-
-      train = generate_examples(
-          1000000, self.min_length, self.max_length, "Train data..")
-      valid = generate_examples(
-          1000, self.min_length, self.max_length, "Valid data..")
-      test = generate_examples(
-          1000, self.max_length, self.max_length, "Test data..")
-
-      np.savez(self.npz_path, train=train, test=test, valid=valid)
-    else:
-      print("[*] Loading dataset for {}".format(self.task))
-      data = np.load(self.npz_path)
-      self.train, self.test, self.valid = \
-          data['train'], data['test'], data['valid']
+    self.data = {}
+
+    for name, num in self.data_num.items():
+      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)):
+          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)
+
+        np.savez(path, x=x, y=y)
+        self.data[name] = TSP(x=x, y=y, name=name)
+      else:
+        tf.logging.info("Skip creating {} for [{}]".format(path, self.task))
+        tmp = np.load(path)
+        self.data[name] = TSP(x=tmp['x'], y=tmp['y'], name=name)
+
+  def get_path(self, name):
+    return os.path.join(
+        self.data_dir, "{}_{}={}.npz".format(
+            self.task_name, name, self.data_num[name]))

main.py

@@ -4,15 +4,15 @@
 
 from trainer import Trainer
 from config import get_config
-from utils import prepare_dirs, save_config
+from utils import prepare_dirs_and_logger, save_config
 
 config = None
 
 def main(_):
-  prepare_dirs(config)
+  prepare_dirs_and_logger(config)
 
-  if not config.task.startswith('tsp'):
-    raise Exception("[!] Task should starts with tsp")
+  if not config.task.lower().startswith('tsp'):
+    raise Exception("[!] Task should starts with TSP")
 
   rng = np.random.RandomState(config.random_seed)
   tf.set_random_seed(config.random_seed)

model.py

@@ -34,12 +34,13 @@ def __init__(self, config, data_loader, is_critic=False):
 
     self.layer_dict = {}
 
+    #self._build_input_ops()
     self._build_model()
     if is_critic:
       self._build_critic_model()
 
-    self._build_optim()
-    self._build_summary()
+    #self._build_optim()
+    #self._build_summary()
 
     show_all_variables()
 
@@ -49,6 +50,15 @@ def _build_summary(self):
   def _build_critic_model(self):
     pass
 
+  def _build_input_ops(self):
+    min_queue_examples = values_per_shard * input_queue_capacity_factor
+    capacity = min_queue_examples + 100 * batch_size
+    values_queue = tf.RandomShuffleQueue(
+        capacity=capacity,
+        min_after_dequeue=min_queue_examples,
+        dtypes=[tf.string],
+        name="random_" + value_queue_name)
+
   def _build_model(self):
     self.global_step = tf.Variable(0, trainable=False)
 
@@ -146,8 +156,19 @@ def _build_model(self):
           initializer=self.initializer)
 
   def _build_optim(self):
-    self.loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
-        logits=self.dec_output_logits, labels=self.dec_targets)
+    losses = tf.nn.sparse_softmax_cross_entropy_with_logits(
+        labels=self.dec_targets, logits=self.dec_output_logits)
+
+    weights = tf.ones(input_length, dtype=tf.int32)
+    batch_loss = tf.div(tf.reduce_sum(tf.multiply(losses, weights)),
+                        tf.reduce_sum(weights),
+                        name="batch_loss")
+
+    tf.losses.add_loss(batch_loss)
+    total_loss = tf.losses.get_total_loss()
+
+    tf.summary.scalar("losses/batch_loss", batch_loss)
+    tf.summary.scalar("losses/total_loss", total_loss)
 
     # TODO: length masking
     #mask = tf.sign(tf.to_float(targets_flat))

trainer.py

@@ -52,11 +52,16 @@ def _build_session(self):
     self.sess = sv.prepare_or_wait_for_session(config=sess_config)
 
   def train(self):
-    print("[*] Training starts...")
-    pass
+    tf.logging.info("Training starts...")
+
+    self.data_loader.run_input_queue(self.sess)
+    import ipdb; ipdb.set_trace() 
+    x = 123
+
+    self.data_loader.run_input_queue()
 
   def test(self):
-    pass
+    tf.logging.info("Testing starts...")
 
   def _inject_summary(self, tag, feed_dict, step):
     summaries = self.sess.run(self.summary_ops[tag], feed_dict)

utils.py

@@ -1,11 +1,26 @@
 import os
 import json
+import logging
 import numpy as np
 from datetime import datetime
 
 import tensorflow as tf
+import tensorflow.contrib.slim as slim
+
+def prepare_dirs_and_logger(config):
+  formatter = logging.Formatter(
+      "%(asctime)s:%(levelname)s:%(message)s")
+  logger = logging.getLogger('tensorflow')
+
+  for hdlr in logger.handlers:
+    logger.removeHandler(hdlr)
+
+  handler = logging.StreamHandler()
+  handler.setFormatter(formatter)
+
+  logger.addHandler(handler)
+  logger.setLevel(tf.logging.INFO)
 
-def prepare_dirs(config):
   if config.load_path:
     config.model_name = "{}_{}".format(config.task, config.load_path)
   else:
@@ -21,23 +36,14 @@ def get_time():
   return datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
 
 def show_all_variables():
-  print("")
-  total_count = 0
-  for idx, op in enumerate(tf.trainable_variables()):
-    shape = op.get_shape()
-    count = np.prod(shape)
-    print("[%2d] %s %s = %s" % (idx, op.name, shape, "{:,}".format(int(count))))
-    total_count += int(count)
-  print("=" * 40)
-  print("[Total] variable size: %s" % "{:,}".format(total_count))
-  print("=" * 40)
-  print("")
+  model_vars = tf.trainable_variables()
+  slim.model_analyzer.analyze_vars(model_vars, print_info=True)
 
 def save_config(model_dir, config):
   param_path = os.path.join(model_dir, "params.json")
 
-  print("[*] MODEL dir: %s" % model_dir)
-  print("[*] PARAM path: %s" % param_path)
+  tf.logging.info("MODEL dir: %s" % model_dir)
+  tf.logging.info("PARAM path: %s" % param_path)
 
   with open(param_path, 'w') as fp:
     json.dump(config.__dict__, fp,  indent=4, sort_keys=True)