Add initial search.py for module search.

For now hardwire to use tfds datasets for image classification.

PiperOrigin-RevId: 275846855

tensorflow_hub/tools/module_search/BUILD

@@ -14,11 +14,24 @@
 # ==============================================================================
 licenses(["notice"])  # Apache 2.0 License
 
+py_binary(
+    name = "search",
+    srcs = ["search.py"],
+    python_version = "PY3",
+    deps = [
+        ":utils",
+        "//tensorflow_hub:expect_numpy_installed",
+        "//tensorflow_hub:expect_tensorflow_installed",
+    ],
+)
+
 py_library(
     name = "utils",
     srcs = ["utils.py"],
     deps = [
         "//tensorflow_hub:expect_numpy_installed",
         "//tensorflow_hub:expect_tensorflow_installed",
+        "//tensorflow_hub:expect_tensorflow_datasets_installed",
+        "//tensorflow_hub",
     ],
 )

tensorflow_hub/tools/module_search/search.py

@@ -0,0 +1,126 @@
+# Copyright 2019 The TensorFlow Hub Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tool to rank modules to use in a downstream classification task."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl import app
+from absl import flags
+
+import pandas as pd
+import numpy as np
+import tensorflow.compat.v2 as tf
+
+from tensorflow_hub.tools.module_search import utils
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("dataset", None,
+                    "Specification of a dataset. E.g. use `cifar10#1000` to "
+                    "perform search using 1000 examples from tfds `cifar10` "
+                    "dataset.")
+
+flags.DEFINE_multi_string("module", None, "Module to consider in the search")
+
+flags.DEFINE_string("module_list", None,
+    "Path to text file with a module per line to be considered in the search."
+    "Empty lines and lines starting with # are ignored")
+
+
+def load_data(data_spec):
+  return utils.load_data(**data_spec)
+
+
+def load_raw_features(data_spec):
+  data = load_data(data_spec=data_spec)
+  return data.map(lambda x: tf.image.resize(x["image"], (224, 224)))
+
+
+def load_labels(data_spec):
+  data = load_data(data_spec=data_spec)
+  return np.array([x for x in data.map(lambda x: x["label"])])
+
+
+def compute_embeddings(module_spec, data_spec):
+  raw_features = load_raw_features(data_spec=data_spec)
+  embedding_fn = utils.load_embedding_fn(
+      module=module_spec)
+  outputs = []
+  for batch in raw_features.batch(10):
+    outputs.extend(embedding_fn(batch))
+  return np.array(outputs)
+
+
+def compute_score(module_spec, data_spec):
+  embeddings = compute_embeddings(module_spec=module_spec,
+                                  data_spec=data_spec)
+  distances = utils.compute_distance_matrix_loo(embeddings)
+  labels = load_labels(data_spec=data_spec)
+  error_rate = utils.knn_errorrate_loo(distances, labels, k=1)
+  return np.array(error_rate)
+
+
+def main(argv):
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+
+  if not FLAGS.dataset:
+    raise app.UsageError("--dataset is a required argument.")
+
+  module_list = []
+  if FLAGS.module:
+    module_list.extend(FLAGS.module)
+
+  if FLAGS.module_list:
+    with tf.io.gfile.GFile(FLAGS.module_list) as f:
+      lines = f.read().split("\n")
+      module_list.extend([l for l in lines if l != "" and not l.startswith("#")])
+
+  ds_sections = FLAGS.dataset.split("#")
+  dataset = ds_sections[0]
+  train_examples = int(ds_sections[1]) if len(ds_sections) != 0 else None
+  data_spec = {
+    "dataset": dataset,
+    "split": "train",
+    "num_examples": train_examples,
+  }
+
+  results = []
+  for module in module_list:
+    results.append((
+        module, data_spec,
+        compute_score(module_spec=module, data_spec=data_spec)))
+
+  df = pd.DataFrame(results, columns=["module", "data", "1nn"])
+  df = df.filter(["module", "1nn"])
+  df.sort_values(["1nn"])
+  df.reset_index(drop=True)
+  df.set_index("module")
+
+  with pd.option_context(
+      "display.max_rows", None,
+      "display.max_columns", None,
+      "display.precision", 3,
+      "max_colwidth", -1,  # Don't truncate columns (e.g. module name).
+      "display.expand_frame_repr", False,  # Don't wrap output.
+  ):
+    print("# Module ranking for %s" % data_spec)
+    print(df)
+
+
+if __name__ == '__main__':
+  app.run(main)