8a2748d9aa3f4de31aad46e765f48874
Swift for TensorFlow, First Class Machine Learning in Swift

本文是我在 atSwift Conf 2018 上的分享的文字版。内容和 keynote 高度重合,不过因为是文章,一些地方解释得会更加到位一些。如果看过 keynote 的同学可以不用看了。

背景

首先一个问题是:为什么要关注 Swift 在机器学习领域的应用?
从语言层面来看,Swift 兼具了脚本语言快速开发迭代的特性以及 low level 的编译型语言的性能。而效率和性能是后端开发考虑的两个主要方面,所以这也不能解释为什么 Server-Side Swift 受到越来越多的关注。概括来说,Swift 具备以下优势:

  1. 高性能;
  2. 开源并且社区非常活跃,语言本身也在高速发展与迭代;
  3. 开发效率很高,并且语言级别的静态类型和异常处理使得代码非常安全;

Next Big Thing on Server

既然说到 server 开发,我们不妨认真地思考一下最近几年的 server 开发都在讨论什么,云计算和容器化极大程度的提高运维效率之外,最重要的一大改变就是基于机器学习的应用呈现出了井喷式的发展,其中主要包括以下几个方面:

  1. 推荐系统,几乎运用到了互联网的所有服务中;
  2. 文本分析,推荐系统的基础能力之一,苹果在 iOS 中也非常重视 NLP 能力的完善与建设;
  3. 音视频理解,音视频由于其数据的特殊性,长久以来我们只能从信号学层面做一些粗浅的分析,近年来得益于卷积神经网络在计算机视觉领域的成功应用,越来越多的公司在使用这些技术尝试从视频中提取更多的结构化信息;
  4. 超分辨率,简单的说就是小图变大图,曾经我们普遍使用双线性插值和双三次插值,近年来通过试用卷积神经网络技术能够更好地猜出变大后图片的细节,也有广泛的应用。
  5. 聊天机器人,这个就不用说了,智能化客服基本是已经产业化的技术,虽然有时候会被调侃为人工智障,但不可忽视的是这项技术已经切实为很多公司节省了客服的人力成本。

目前机器学习应用主要的开发语言是 C++ 和 Python,Python 用于快速验证想法,验证之后用 C++实现后部署到线上系统,这两门语言的优缺点都比较明显,C++性能很强,但门槛高,开发效率低,Python 容易写,但也非常容易写出很多垃圾代码。所谓动态类型一时爽,重构起来火葬场。

我们不禁会想,既然 Swift 是一门博众家之长的语言,那假如它能提供很好的机器学习能力,岂不是非常完美?

本文的主角解决的就是这个问题 —— Swift for TensorFlow

Swift for TensorFlow 概览

Swift for Tensor, 简称 TFiwS( 为什么这么叫,官方并没有解释,我大胆的猜测一下,应该是正读看起来像TensorFlow for Swift, 倒过来读就是 Swift), 文章接下来会简称 TFS.

TFS 是今年 Chris 在2018 TensorFlow Dev Submit 上发布的,到现在也没满一年,是一个非常年轻的技术,现在也不算正式发布,但已经开源了。目前正在 active developing 的阶段,基本上每两周会有一个新的版本放出来。

TFS 有如下几个优势:

  1. 可以融入完整的 TF 生态系统,众所周知目前 TF 的生态是最完善的;
  2. 不仅仅只是 TensorFlow 的 API binding,而是提供了 First class machine learning 的能力;
  3. 基于2中有趣的特性(后面会讲),使得 TFS 同时支持跑 CPU、GPU 以及远程的 TPU;
  4. 兼具 可用性高性能
  5. 最后一点,支持 Xcode Playground,使得天生具备“Notebook”一样的感觉;

简单的例子:

那个诡异的实心圆代表矩阵乘法,等价于 matmul

那到底什么才是所谓的“First class machine learning”, 以及为什么说 TFS 是兼具可用性和性能呢? 我们还需要看 TF 的两种模式。

TensorFlow 的两种经典模式

  1. 图模式。代码如下:
#!/usr/bin/python
#coding=utf8
"""
# Author: aaron
# Created Time : 2018-09-01 16:50:36

# File Name: graph.py
# Description:

"""

import tensorflow as tf

x = tf.constant([[1,2,3], [4,5,6]])
xt = tf.transpose(x)
y = tf.matmul(x, xt)

with tf.Session() as sess:
    print sess.run(y)

图模式的一个重要特点就是:Lazy Evaluation。比如在执行 xt = tf.transpose(x) 的时候,实际上并没有触发矩阵转正的运算,而是只是生成了一个名为”转置”的运算节点,添加到了计算图中。最后,当执行 sess.run(y)的时候,所有计算才开始运行。

由于这个独特的模式,图模式的优点很明显,那就是性能很强,因为计算时已经知道了所有的计算节点,可以做很多优化,同理,缺点也明显,可用性很差,先建图后计算的模式并不符合直觉,而且只能使用节点支持的运算来构建计算图。Debug 起来很痛苦,因为你并不能通过 print x 这样的形式来 debug 一些中间变量(执行到 print 的时候 x 还没有 value) 。

  1. Eager Execution 模式。示例代码如下:
#!/usr/bin/python
#coding=utf8
"""
# Author: aaron
# Created Time : 2018-09-01 16:59:01

# File Name: eager.py
# Description:

"""

import tensorflow as tf
tf.enable_eager_execution()

x = tf.constant([[1,2,3],[4,5,6]])
y = tf.matmul(x, tf.transpose(x))
print(y)

顾名思义,Eager 模式与图相反,计算是立即发生的,整个过程并不会构建图。所以可以使用自然的流程控制,比如 if 语句来书写模型,也可以在执行的过程中插入 print x 来获取中间值,帮助我们 debug。

Eager 模式更符合直觉,易于理解,易于 debug。但因为没有 lazy,所以很难做优化,性能并不好。

第三种模式 — TFS 模式

以下是一段 Swift 代码:

    let x : Tensor<Float> = [[1,2,3], 
                             [4,5,6]]
    var y = Tensor<Float>(zeros: 
                            x.shape)
    print(y)
    if x.sum() > 100{
        y = x
    }else{
        y = x • x
    }
    print(y)

上面的代码看起来是 eager 的,但其实跑起来是以 graph 的形式跑的,所以具备图模式的所有优化。TFS 实现了一个改进版的 Swift 编译器,在编译阶段会自动分析代码中的 tensor 运算,并翻译成计算图,最终执行的时候就以图的模式运营。而且这个过程对程序员透明,也就是说:

程序员只需要当它是 eager 的模式来写代码,最终跑起来就和图模式一样快了

听起来是不是不太现实? 这是什么做到的呢? 我们进入下一个章节。

TFS 模式的原理:Program Slicing

上文提到,TFS 会在编译期间分析出所有的 tensor 运算,并生成计算图,我们先来看看具体是怎么做的,计算图由 TensorFlow Runtime 运行的。

简单的例子

先从一段简单的代码开始:

typealias FloatTensor = Tensor<Float>

func linear(x : FloatTensor, w : FloatTensor, b : FloatTensor) -> FloatTensor
{
    let tmp = w • x
    let tmp2 = tmp + b
    return tmp2
}

代码很直观,计算 wx + b ,只是以矩阵的形式。为了简单期间,我们没有使用变量,而是都用的常量。

如果是针对上面的代码,要怎么生成计算图的? 既然 Swift 没有 runtime,我们不妨来看看其对应的 AST。

可以看到,所有 Tensor 类型的值和变量都可以从 AST 中找出来,也就是说,我们只需要遍历这个 AST,就可以得到所有 Tensor 相关的常量和对应的操作,自然就可以构建计算图。如下图所示




以看到,这种在编译期间生成计算图的计算几乎强依赖于 Swift 的静态类型系统,所以这活儿 Python 还真不一定适合,Chris 在 TFiwS 的白皮书也介绍了为什么使用 Swift,感兴趣的可以点这里:https://github.com/tensorflow/swift/blob/master/docs/WhySwiftForTensorFlow.md

生成计算图后,代码该怎么执行呢?我们只需要将图代码移除后,替换为启动图计算的代码即可,上述代码变换后如下所示:

typealias FloatTensor = Tensor<Float>
func linear(x : FloatTensor, w : FloatTensor, b : FloatTensor) -> FloatTensor
{
    let result = execTensorFlowGraph("graph_generate_before")
    return result
}

上述代码运行时,直接调用 TensorFlow Runtime 执行刚才生成的计算图,取得执行结果后返回

复杂的例子

现在我们来看一个复杂一点的例子:

func linear(x : FloatTensor, w : FloatTensor, b : FloatTensor) -> FloatTensor
{
    let tmp = matmul(x, w)
    let tmp2 = tmp + b
    print(tmp2)
    let tmp3 = tmp2 * magicNumberGenerateFromTensor(x: tmp2)
    return tmp3
}

func magicNumberGenerateFromTensor(x : FloatTensor) -> Float
{
    return 3.0
}

和上一个例子不同的是,这里我们添加了一个真 eager 模式的代码,print(tmp2) , 不仅如此,在 tmp3的计算中,我们还耦合了 tensor 的逻辑和普通的本地运算: magicNumberGenerateFromTensor,在这样的前提下,我们生成的图是什么样的呢?

tensor逻辑,代表最终要生成 graph,由 TensorFlow Runtime 运行的逻辑,或者简单的理解为 Tensor 类的相关操作都是 tensor 逻辑。
而本地运算,则是由 Swift Runtime 执行的非 tensor 逻辑的代码。

之前的思路,我们可以得到上述的图,但存在两个问题:

  1. 我们需要打印出 tmp2 ,也就是 wx+b的值,但是这个图是以一个整体由 TensorFlow 执行的,如何自动拿到某个中间的值?
  2. magic 是由本地逻辑计算得到,TensorFlow 在执行这张图的时候,怎么知道 magic 的值到底是多少呢?

Program Slicing

接下来就是我们的 Program Slicing 技术登场了,这并不是一项新技术,但这应该是它第一次被用来干这事儿,详细介绍大家可以自行 wikipedia。

大家看完上面的问题,会发现一个终极问题是,本地代码和 graph 似乎有双向通信的需求。TFS 实现这样的双向通信,主要是通过 Program Slicing 技术,将原始代码转换为两份不同的代码,一份用来生成图,一份用来在本地跑。两份代码互相通信,完成计算。

我们再 review 一下原始代码,然后开始 slicing!

原始代码:
```swift
func linear(x : FloatTensor, w : FloatTensor, b : FloatTensor) -> FloatTensor
{
let tmp = matmul(x, w)
let tmp2 = tmp + b
print(tmp2)
let tmp3 = tmp2 * magicNumberGenerateFromTensor(x: tmp2)
return tmp3
}

top Created with Sketch.