深入浅出XDL(三):framework

前一篇文章我们分析了xdl的稀疏API embedding的实现。同时也发现了一些的新的问题,首先来回顾一下。

embedding接口在模型中创建几类Op:Variable Op, 参数拉取op, Combiner Op; embedding接口相关的几个Op都继承于XDL自己定义的OpKernel,包括PsRegisterVariableOp, PsSparsePullOp, PsPullOp。我们知道,tensorflow的模型graph,只能包含继承于tensorflow Op的Op。那么embedding接口引入的这几个Op是如何与tensorflow的Op一起组成模型的呢?

并且,同时通过阅读源码(代码位置:xdl/xdl/core/)发现,xdl自己实现了一套新的runtime framework,包括graph, device, op, executor,tensor等等核心类的定义。对tensorflow的核心概念都有自己的一套实现。那么这套新的runtime framework与tensorflow如何一起工作的呢?

首先先看一下完整的XDL例子

样例

例子来自于xdl自带的example(代码位置:xdl/examples/deepctr/deepctr.py)

def train():
    ...
    ## 省略掉了数据加载部分. 注意emb1和emb2为xdl op,具体为PsSparsePullOp
    emb1 = xdl.embedding('emb1', batch['sparse0'], xdl.TruncatedNormal(stddev=0.001), 8, 1024, vtype='hash')
    emb2 = xdl.embedding('emb2', batch['sparse1'], xdl.TruncatedNormal(stddev=0.001), 8, 1024, vtype='hash')
    
    ## loss, train_op 均为xdl op
    loss = model(batch['deep0'], [emb1, emb2], batch['label'])
    train_op = xdl.SGD(0.5).optimize()
    log_hook = xdl.LoggerHook(loss, "loss:{0}", 10)
    sess = xdl.TrainSession(hooks=[log_hook])
    while not sess.should_stop():
        sess.run(train_op)
        
@xdl.tf_wrapper()
def model(deep, embeddings, labels):
    
    ## input 为tensorflow concat op, 这里有个问题,model函数的输入
    ## deep, embedding, labels均为xdl op, 为什么可以直接作为tensorflow
    ## concat op的输入呢 ? 
    input = tf.concat([deep] + embeddings, 1)
    
    
    ## fc1, fc2, fc3, y 均为tensorflow op, 并且这些全连接网络层
    ## 会引入模型参数,这些参数是如果交由xdl的参数服务器管理的呢?并
    ## 没有看到显示的调用xdl api分配这些dense的参数。
    fc1 = tf.layers.dense(
        input, 128, kernel_initializer=tf.truncated_normal_initializer(
            stddev=0.001, dtype=tf.float32))
    fc2 = tf.layers.dense(
        fc1, 64, kernel_initializer=tf.truncated_normal_initializer(
            stddev=0.001, dtype=tf.float32))
    fc3 = tf.layers.dense(
        fc2, 32, kernel_initializer=tf.truncated_normal_initializer(
            stddev=0.001, dtype=tf.float32))
    y = tf.layers.dense(
        fc3, 1, kernel_initializer=tf.truncated_normal_initializer(
            stddev=0.001, dtype=tf.float32))
            
    ## loss也是tensorflow的op, 为什么上面说model返回的是xdl op呢?
    ## 上面函数train_op是xdl op,它是如果做参数训练的呢 ?
    loss = tf.losses.sigmoid_cross_entropy(labels, y)
    return loss

将上面的问题罗列一下:

  • xdl op为什么可以直接作为tensorflow concat op的输入
  • tf.layers.dense全连接网络的参数,如果分配到xdl的参数服务器的
  • loss是tensorflow op, 为什么说model返回的是xdl op

回答这些问题的关键点在于:

  • 装饰器xdl.tf_wrapper(代码位置:xdl/xdl/python/backend/tf/tf_backend.py)
  • hook函数get_variable(代码位置:xdl/xdl/python/backend/tf/tf_hook.py)
  • xdl backend op: TFBackendOp(代码位置:xdl/xdl/core/backend/tf/backend_op.cc)

装饰器xdl.tf_wrapper

python装饰器语法不是本文的重点,不熟悉的读者可以找相关文档了解一下。直接介绍xdl.tf_wrapper的作用

  • 为被装饰的函数(比如example中的model函数)的输入(一般为xdl op),生成对应的tensorflow placeholder op
  • 调用被装饰的函数,并传入原始输入对应的placeholder op,被装饰函数一般用来构建tensorflow graph,但此时被装饰函数的输入已经有从xdl op 换成了相对应的tensorflow placeholder了
  • 被装饰函数创建的tensorflow graph,以及原始输入,都会成为xdl TFBackendOp的输入, xdl TFBackendOp会加入到xdl graph
  • 目前为止,xdl graph中有了embedding相关的op, 以及TFBackendOp,接下来xdl session会驱动xdl graph的运行

hook函数get_variable

被装饰函数(比如example中的model函数)在构建tensorflow网络的时候,一般是模型的dense部分,同样也会申请tensorflow variable,这些variable同样需要放到xdl的参数服务器上管理,但是被装饰函数里并没有显示的调用xdl的参数分配函数,这是怎么做到的呢?

答案是xdl hook了tensorflow的api tf.get_variable,替换为xdl自己实现的get_variable函数。

xdl在的自己实现的get_variable函数内

  • 首先调用原始的tensorflow get_variable函数,添加tensorflow variable op到tensorflow graph
  • 然后为刚添加的tensorflow variable op, 在xdl graph里添加一个对应的xdl variable op并作为TFBackendOp的输入

TFBackendOp

TFBackendOp负责驱动tensorflow graph的执行,也就是说整个tensorflow graph都将以TFBaclendOp的形式,参与xdl graph的执行。

TFBackendOp是连接xdl graph与tensorflow graph的关键,从tensor流动角度分析:

  • 从xdl graph到tensorflow graph的tensor,都将作为TFBackendOp的input;比如前面添加的tensorflow placeholder, tensorflow variable的填充值,都将通过TFBackendOp的input传入
  • 从tensorflow graph到xdl graph的tensor,都将作为TFBackendOp的output

总结

至此,我们大致理解了XDL的运行机制

  • XDL有自己的独立的framework,XDL API创建的是XDL Graph,添加的是XDL Op,并由XDL Session驱动运行
  • Tensorflow API创建的是tensorflow graph, 添加的的是tensorflow op, 由XDL TFBackendOp调用tensorflow session驱动执行
  • TFBackendOp是连接两个Graph的桥梁,Tensor正式通过这个桥梁在Graph间流动的
发布了52 篇原创文章 · 获赞 105 · 访问量 7万+

猜你喜欢

转载自blog.csdn.net/gaofeipaopaotang/article/details/100933225