今年年初，伴随着 tensorflow 更新到 1.7.0 版本，Google 发布了 tensorflow hub。tensorflow hub 的主要目标是为模型提供一种简便的封装方式，同时可以简便地复用已封装的模型，可以说 tf hub 是为迁移学习而生的。
熟悉自然语言处理的同学都知道大部分 nlp 模型的底层都是 word2vec 词向量作为一个词的特征，当然近几年越来越多的模型会构建于语言模型之上，例如用 ELMo 代替词向量。其实无论哪种方式，downstream 的任务都建立在这些预训练好的向量之上，downstream 的任务与底层 embedding 的训练是高度解耦的。所以完全可以有专门的团队负责底层 embedding 的优化与开发，让后将它们用 tf hub 封装成 module 供下游应用团队使用，这些 module 对于使用人员就是黑盒子，他们无需关心 module 的实现细节。
在图像领域也是一样的，通常一些 downstream 的任务都会建立在一些经典的模型（vgg, resnet, mobilenet 等）之上，它们会利用这些模型预训练好的权重及结构作为特征提取器。
这里结合 nlp 中的 embedding 的封装和使用介绍一下 tensorflow hub 的细节，安装方式看 github。主要参考了官网和 github：
https://tensorflow.google.cn/hub/
https://github.com/tensorflow/hub
我们先看看怎么使用一个别人为我们封装好的模型：
hub_module = hub.Module(self.get_temp_dir())
tokens = tf.constant(["cat", "lizard", "dog"])
embeddings = hub_module(tokens)
with tf.Session() as session:
session.run(tf.tables_initializer())
session.run(tf.global_variables_initializer())
self.assertAllClose(
session.run(embeddings),
[[1.11, 2.56, 3.45], [0.0, 0.0, 0.0], [1.0, 2.0, 3.0]])
代码的第一行表示加载一个封装好的 tensorflow hub 模型，参数可以是模型的路径也可以是一个保存有模型的 http 地址。
第二行创建了一个包含三个字符串的张量。
第三行调用了刚刚创建的 hub 模型，字符串张量作为模型的输入，embeddings就是模型的输出了。
调用 session.run 就能得到具体的 embeddings 输出值。
我第一次见到这个 demo 的时候惊喜之处在于，我们往常使用 embedding 时都需要将舒服的单词转换成相应的 id，将 id 作为输入查询相应的 embedding。难道这里把转换操作也封装进 hub module 里的？为了满足好奇心我们可以看一下如何封装这样一个 module。
hub module 封装：
def module_fn():
"""Spec function for a token embedding module."""
tokens = tf.placeholder(shape=[None], dtype=tf.string, name="tokens")
embeddings_var = tf.get_variable(
initializer=tf.zeros([vocab_size + num_oov_buckets, embeddings_dim]),
name=EMBEDDINGS_VAR_NAME,
dtype=tf.float32)
lookup_table = tf.contrib.lookup.index_table_from_file(
vocabulary_file=vocabulary_file,
num_oov_buckets=num_oov_buckets,
)
ids = lookup_table.lookup(tokens)
combined_embedding = tf.nn.embedding_lookup(params=embeddings_var, ids=ids)
hub.add_signature("default", {"tokens": tokens},
{"default": combined_embedding})
spec = hub.create_module_spec(module_fn)
with tf.Graph().as_default():
m = hub.Module(spec)
p_embeddings = tf.placeholder(tf.float32)
load_embeddings = tf.assign(m.variable_map[EMBEDDINGS_VAR_NAME],
p_embeddings)
with tf.Session() as sess:
sess.run([load_embeddings], feed_dict={p_embeddings: embeddings})
m.export(export_path, sess)
根据上面代码，创建 hub module 的流程如下：
1、调用 hub.create_module_spec 创建一个 spec，函数的参数是 module 的计算图创建函数
2、调用 hub.Module 创建一个 module 对象，参数是上一步创建的 spec
3、在 session 中训练模型，这个 demo 里面没有训练，而是直接利用 tf 的赋值操作将一个 numpy 矩阵赋值给了模型的参数
4、调用 export 函数将当前的 session 保存到某个路径中
在 module_fn 函数的最后调用了 hub.add_signature，第一个参数是创建的这个 hub module 的名称，第二个参数是 module 的输入，它是一个字典，支持多个输入，第三个参数是 module 能提供的输出，同样也是字典，支持输出多个数据。
可以看到，module 的输入定义为一个 字符创类型的 placeholder，然后利用 index_table_from_file 创建了一个 lookup_table，这个 table 就可以将字符串转化为相应的 id，这里就解答了之前的好奇。
总结
第一次看到 tensorflow hub 就觉得很优雅，以前做 nlp 的工作会花大量的时间在准备数据上，而利用 tf hub 以后所有的任务都可以使用统一的 embedding module，并且可以直接将字符串作为输入，不用再手动转换。
hub module 在使用时还能设定为参数可训练或者参数不可训练，这样对于不同的任务就能有更灵活的选择。对于一些训练样本较少的情况，可以冻结底层 module 的参数，做完全的迁移学习。
首次尝试 tensorflow hub 还是相当欣喜的，以后也会尽量使用。
