基础知识背景

2019/12/20 posted in  TensorFlow

与其说TensorFlow是一套框架,更愿意说TensorFlow是一套语言,TensorFlow灵活性特别大,其本质在于构建一个计算图,计算图的各个节点按照计算逻辑,与资源配置情况分配到不同的计算设备上来进行计算,计算图整体的计算资源,比如计算设备CPU、GPU、线程池是框架本身来灵活配置的,而灵活性过大导致了TensorFlow无论在计算图、计算资源的适配上都需要花费额外的功夫才能保证模型训练效率,加上TensorFlow API的多种不同的魔性版本(比如CV场景,TensorFlow历史上支持多个完全不同的API封装,如Slim、Estimator、Keras等等)、以及文档上很多模棱两可的工作,在TensorFlow写好高效的模型即使是对一名资深的算法从业人员也是一件相对复杂的工作。

本文会在先通过支持多个业务的TensorFlow任务学习到知识来分享在很多业务场景下TensorFlow效果不高的原因,基于这样的背景,我们希望和算法同学分享下在云音乐推荐场景下,TensorFlow标准化的一些工作。

TensorFlow低效的几个原因分析

低效数据读取

数据读取部分一直是业务同学特别忽略的一个过程,通常业务数据,尤其是推荐、搜索场景下的业务数据,整体量差异很大,从数G到数T不止,如果预先拉取数据到本地进行存储,会额外增加耗时,TF本身提供了HDFS的数据访问方法,也提供了一套相应地高效地读取HDFS数据的方法,但是存在各种各样的问题:

  1. 读取数据接口API未能明显标识那些是python级别的接口函数、哪些是C++的接口函数,误用python数据读取接口时,由于python本身全局锁问题,导致性能极低;
  2. 对数据进行操作时,由于本身python在这方面的便利以及算法同学对numpy、scipy、pandas比较熟悉,容易引入python的操作,也会造成1中遇到的问题;
  3. 使用Dataset的API时,由于一般在推荐场景下数据存储在hdfs这类存储系统上,一般的读取接口没有做专项优化,未使用优化过的数据读取API;
  4. 未合理进行向量化操作,具体比如对dataset做map操作时,应在batch前还是batch后,batch后是向量化操作,而batch前是单个单个处理,函数调用次数前者远低于后者;
  5. 在配置环境来读取hdfs文件时,我们发现,hadoop环境会默认配置MALLOC_ARENA_MAX环境变量,这个变量控制malloc的内存池,embedding_lookup_sparse的uniqueOp在被hadoop限制MALLOC_ARENA_MAX=4后会受很大影响;

频繁移动你的数据

在机器学习系统中,要想程序跑的快,有两个原则:

  1. 尽量减少数据的移动;
  2. 数据离计算尽可能近;

在写TF时,由于对API的不熟悉,很容易会造成数据的移动,通常这样不会产生太大的问题,但是在某些场景下,会导致大量的数据拷贝,不同设备之间的拷贝不仅会占用设备间的数据通道,也会耗费大量的原本可以用来计算的资源。
比较常见的问题是Graph内外数据的拷贝,如下code:

import tensorflow as tf
ds = read_from_tfrecord_fun(path)
with tf.Session() as sess:
  data,label = sess.run(ds)
    
... # build your model
train_op = build_model()
with tf.Session() as sess:
  while ...:
    sess.run(train_op, feed_dict={"data": data, "label": label})

这部分代码相当于把整个模型的训练分为两个部分:第一个部分将数据从tfrecord文件中拿出,第二个部分将从tfrecord拿出的数据再feed进训练网络。
从功能上这个没有什么问题,但是仔细想想,这里涉及到graph到python内存的拷贝,然后从python内存拷贝到训练的graph中,并且由于python本身性能和GIL的影响,这里的耗时极大,并且这个拷贝是随着训练一致存在,假设你数据大小为100G,训练10个epoch,整个训练过程相当于不停地从graph内拷贝到python进程空间,然后从python进程空间拷贝到graph内,整个数据量为2*10*100G=2T,这部分的耗时相对于训练时间占很大比例。

资源非必需原因消耗

这一块有大概几个方面:
Graph自增节点
如下代码

def build_model(model_path):
    model_input = tf.placeholder('float32', [1, IMAGE_HEIGHT, IMAGE_WIDTH, COLOR_CHANNELS])
    ...... 
    return model_input,vec1
    
def get_ops(vec):
    ...
    return new_vec

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    img_input, vec1_out = build_model("path/to/model")
    for cur_img_path in imgs_path_list: 
        cur_img = load_image(cur_img_path)
        vec1_out = sess.run(vec1, feed_dict = {img_input:cur_img})      # 正向传播输出模型中的vec1
        new_vec = get_ops(vec1_out)       

在对图像列表循环操作时,一直在不停地调用get_ops来产生新的节点,graph一直在变大,通常这个过程不会影响代码本身的正确性,但是由于一直在不停地扩展graph,会造成计算机资源一直在扩展graph,本身模型的计算会被大大地阻碍。

cache()
在使用dataset读取数据时,提供一个cache()的api,用于缓存读到之后处理完成的数据,在第一个epoch处理完成之后就会缓存下来,若cache()中未指明filename,则会直接缓存到内存,相当于把数据集所有数据缓存到了内存,直至内存爆掉;

不合理的gpu_config
TensorFlow在调用gpu时,由于仅在memory fraction层面来做了相关的切分来区分某个gpu被各自任务占用,这里会有很多问题,比如某个人物memory比较小,但是会一直占用gpu的kernel来进行计算,另外还有目前一些对于gpu的virtualDevices也会造成任务的速率有很大的问题,如下为config.proto中关于gpu的配置相关内容。

message GPUOptions {
  double per_process_gpu_memory_fraction = 1;
  bool allow_growth = 4;
  string allocator_type = 2;
  int64 deferred_deletion_bytes = 3;
  string visible_device_list = 5;
  int32 polling_active_delay_usecs = 6;
  int32 polling_inactive_delay_msecs = 7;
  bool force_gpu_compatible = 8;
  message Experimental {
    message VirtualDevices {
            repeated float memory_limit_mb = 1;
    }
  repeated VirtualDevices virtual_devices = 1;
  bool use_unified_memory = 2;
  int32 num_dev_to_dev_copy_streams = 3;
  string collective_ring_order = 4;
  bool timestamped_allocator = 5;
  int32 kernel_tracker_max_interval = 7;
  int32 kernel_tracker_max_bytes = 8;
  int32 kernel_tracker_max_pending = 9;
  }
  Experimental experimental = 9;
}

使用TensorFlow做模型计算之外的事情

TensorFlow在宣传时,期望all in TensorFlow, ALL In XX是所有开源框架的目标,但是是不可能的,所有系统都是一步一步的tradeoff的过程,在一方面性能的优势,必然在设计上会有另一方面的劣势。TensorFlow也不例外,尽管TensorFlow设计了很多data preprocess的api,这些api尽管在功能实现上完成了某些machine learning场景下基本的data preprocess的功能,但是从整个系统架构上来看,TensorFlow来进行data preprocess的工作本身并没有太多优势,相反会由于种种原因会造成很多额外的资源损耗与问题:

数据无法可视化

TensorFlow Data Preprocess Pipeline因为直接在graph内,其处理无法可视化,就是一些字节流,无法可视化,也无法验证准确性,机器学习是一个try and minimize error,如果仅能训练处数据而看不到数据,无异于让瞎子指挥交通,这一点上目前TensorFlow Data Process Pipeline并没有特别好的解决方案,传统的大数据解决方案相对来说会更具可行性。

更多适合data preprocess工具

其实,在数据处理上,有更多适合在Spark、Flink上去做,有几个方面的原因:

  1. 底层机器选择上,Spark、Flink的集群架构被用来设计做data preprocess,无论在性能、数据质量、各种指标监控上都有很成熟的经验,而深度学习主机并不适合来做data preprocess的工作;
  2. 一部分data preprocess的工作,尤其是计算复杂的data transform的工作可以计算完成之后cache到数据库当中,线上计算时不需要额外计算,而是直接缓存,如分桶、ID化这些工作;
  3. 被缓存的数据是one-stage finish的,而集成在TF中,后续模型上线这部分耗时会一直在存在,优化空间极易碰到天花板;
  4. 在实时场景下,我们通常会利用各种大数据工具来完成比如实时数据的校验,通过snapshot收集实时推荐反馈后的行为来形成正负样本,这块目前有很多有效的工具来完成,而使用TensorFlow Data Pipeline目前看来没有特别好的解决方案;

超参设计不合理

这类问题不能称得上是问题,只是算法在设计的时候应该考虑系统本身能力来做相关的优化,比如在考虑模型计算时间太长时,应该尽力保证large batch,这样会减少模型update的次数,从而从数据移动、数据拷贝的角度上来减少相关操作,关于large batch的training,facebook有相关的文章Training ImageNet in 1 Hour - Facebook Research来证明其实通过修改参数以及增加warm up可以达到STOA的指标。
另外包括模型结构、embedding、优化器,这类操作在某些场景下,也会严重影响模型训练的数据。这里分享一个case,关于adam优化器与embedding的一个TensorFlow的issue,当模型中包括一个embedding层时,比如亿级别的feature embedding,理论上模型只需更新其中出现过的feature对应的embedding,但是TensorFlow原始的adam会更新亿级别feature embedding中所有的feature embedding,而不是更新出现过的id embedding,当然TF针对此场景有LazyAdamOptimizer,但是坦白来讲效率也不高。因此,超参的不合理也是影响TensorFlow低效训练的因素之一。

Ironbaby: TensorFlow Recommendation Framework

由于TensorFlow做推荐相关算法开发有以上很多零零碎碎的问题,我们组开发了一套基于TensorFlow的推荐框架:Ironbaby,Ironbaby通过封装基本的操作接口,来固化下来一些比较对计算系统比较友好的操作,一方面能够有效地提高算法的运行效率;另一方面,由于标准化基本操作之后,能够更简便的对系统进行二次开发,比如机器学习平台的容器化改造,比如更方便地对接精排系统提供上线服务等等。
Ironbaby主要从以下几个方面来简化我们模型的构造以及训练过程:

  1. Ironbaby支持Sparse、Dense、Sequence三类数据接口,其中Sparse为稀疏类特征,Dense为稠密特征,Sequence类为不定长但有先后顺序关系的数据,每一个接入Ironbaby的任务,需保证数据处理为这三类之一,Ironbaby支持这三类数据的有效读写与处理,支持包括本地磁盘、cephfs、hdfs、kafka的数据读取方案,也提供local script以及spark script脚本来完成转换,其中spark仅需要在保存tfrecord时,保证column的type即可;
  2. Ironbaby推荐使用配置文件来完成任务模型信息的配置,配置文件的方式能够有效地统一平台对于任务的管理方式,比如任务的定时调度、任务信息监控等等;
  3. 更抽象层的封装,TensorFlow本身有多套相关的api来完成这些工作,但是其文档组织太过于繁琐且复杂,Ironbaby目前专注于推荐场景,对一些经常使用的部分进行了比较好的封装,比如cross module,可以通过简单配置来完成不同cross模块的支持;
  4. 基于estimator + tfrecord的高效数据读取方案,既能保证高效地数据读取、模型训练(杜绝了graph 内外拷贝的风险),也能够通过saved model的方式一键部署在业务系统完成在线推理;
  5. Ironbaby目前未使用任何data preprocess的api,我们任务进入Ironbaby的数据为处理好的数据,Ironbaby只负责模型计算的部分,其他如数据预处理应该有其他的更有效的数据处理工具来完成;
  6. Ironbaby严格控制TensorFlow中不同版本各类api的使用,比如不会频繁引入contrib中的api,所有Ironbaby中使用的api均是测试成功,能够保证基础性能,Ironbaby能够有效缓冲TensorFlow频繁更新与业务代码历史包袱的矛盾;
  7. Ironbaby通过标准化我们的模型训练,也能从另一方面去方便标准化我们模型训练周边相关的服务,比如模型监控、业务数据的监控、在线推理数据监控等等;

最后一点可能是所有规范化SDK的作用,业务同学如果能够使用统一APi完成模型的构建,在需要工程、框架同学支持时,能够有效地减少团队间沟通成本,没有规范的模型代码,需要更多地时间去消费其与问题不相关的逻辑,日积月累,这样的沟通成本对于整个团队都是不可忽视的。

目前Ironbaby已支持的功能

  • 基础模型的封装:如DeepFM, Wide and Deep, XDeepFM, fibinet等模型已经实现,仅需要修改配置文件即可完成模型的训练、评估、离线预测、导出等功能;
  • 支持扩展自定义网络结构:用户仅需要继承BabyModel,然后重写build_model_fn即可;
  • 默认支持TensorBoard,能够有效地观测模型训练参数;
  • 支持模型运行的训练、评估、离线预测、导出成saved_model提供给精排在线推理;
  • 支持TensorFlow的parameter server分布式协议,仅需要配置文件配置好相关目录,后续会基于goblin打通k8s,完成资源的动态分配;

后言

任何对于Ironbaby感兴趣的想要尝试的同学欢迎联系我们团队,我们会提供持续的解决方案与相关能力的培训,目前已完成的文档,后续我们也会在我们gitlab项目主页上开放我们的roadmap,欢迎大家多多讨论