Clean up to start add modern models (#24)

EBMs/inception.py

@@ -0,0 +1,105 @@
+# Code derived from tensorflow/tensorflow/models/image/imagenet/classify_image.py
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os.path
+import sys
+import tarfile
+
+import numpy as np
+from six.moves import urllib
+import tensorflow as tf
+import glob
+import scipy.misc
+import math
+import sys
+
+import horovod.tensorflow as hvd
+
+MODEL_DIR = '/tmp/imagenet'
+DATA_URL = 'http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz'
+softmax = None
+
+config = tf.ConfigProto()
+config.gpu_options.visible_device_list = str(hvd.local_rank())
+sess = tf.Session(config=config)
+
+# Call this function with list of images. Each of elements should be a 
+# numpy array with values ranging from 0 to 255.
+def get_inception_score(images, splits=10):
+  # For convenience
+  if len(images[0].shape) != 3:
+    return 0, 0
+
+  # Bypassing all the assertions so that we don't end prematuraly'
+  # assert(type(images) == list)
+  # assert(type(images[0]) == np.ndarray)
+  # assert(len(images[0].shape) == 3)
+  # assert(np.max(images[0]) > 10)
+  # assert(np.min(images[0]) >= 0.0)
+  inps = []
+  for img in images:
+    img = img.astype(np.float32)
+    inps.append(np.expand_dims(img, 0))
+  bs = 1
+  preds = []
+  n_batches = int(math.ceil(float(len(inps)) / float(bs)))
+  for i in range(n_batches):
+      sys.stdout.write(".")
+      sys.stdout.flush()
+      inp = inps[(i * bs):min((i + 1) * bs, len(inps))]
+      inp = np.concatenate(inp, 0)
+      pred = sess.run(softmax, {'ExpandDims:0': inp})
+      preds.append(pred)
+  preds = np.concatenate(preds, 0)
+  scores = []
+  for i in range(splits):
+    part = preds[(i * preds.shape[0] // splits):((i + 1) * preds.shape[0] // splits), :]
+    kl = part * (np.log(part) - np.log(np.expand_dims(np.mean(part, 0), 0)))
+    kl = np.mean(np.sum(kl, 1))
+    scores.append(np.exp(kl))
+  return np.mean(scores), np.std(scores)
+
+# This function is called automatically.
+def _init_inception():
+  global softmax
+  if not os.path.exists(MODEL_DIR):
+    os.makedirs(MODEL_DIR)
+  filename = DATA_URL.split('/')[-1]
+  filepath = os.path.join(MODEL_DIR, filename)
+  if not os.path.exists(filepath):
+    def _progress(count, block_size, total_size):
+      sys.stdout.write('\r>> Downloading %s %.1f%%' % (
+          filename, float(count * block_size) / float(total_size) * 100.0))
+      sys.stdout.flush()
+    filepath, _ = urllib.request.urlretrieve(DATA_URL, filepath, _progress)
+    print()
+    statinfo = os.stat(filepath)
+    print('Succesfully downloaded', filename, statinfo.st_size, 'bytes.')
+  tarfile.open(filepath, 'r:gz').extractall(MODEL_DIR)
+  with tf.gfile.FastGFile(os.path.join(
+      MODEL_DIR, 'classify_image_graph_def.pb'), 'rb') as f:
+    graph_def = tf.GraphDef()
+    graph_def.ParseFromString(f.read())
+    _ = tf.import_graph_def(graph_def, name='')
+  # Works with an arbitrary minibatch size.
+  pool3 = sess.graph.get_tensor_by_name('pool_3:0')
+  ops = pool3.graph.get_operations()
+  for op_idx, op in enumerate(ops):
+      for o in op.outputs:
+          shape = o.get_shape()
+          shape = [s.value for s in shape]
+          new_shape = []
+          for j, s in enumerate(shape):
+              if s == 1 and j == 0:
+                  new_shape.append(None)
+              else:
+                  new_shape.append(s)
+          o.set_shape(tf.TensorShape(new_shape))
+  w = sess.graph.get_operation_by_name("softmax/logits/MatMul").inputs[1]
+  logits = tf.matmul(tf.squeeze(pool3, [1, 2]), w)
+  softmax = tf.nn.softmax(logits)
+
+if softmax is None:
+  _init_inception()