谷歌云TPU上可以用Julia啦!0.23秒跑100张图片,Jeff Dean点赞推荐

伊瓢 发自 凹非寺 量子位 报道 | 公众号 QbitAI

不久前,Julia Computing官方放出了一篇论文,展示将Julia代码和机器学习模型编译到谷歌云TPU的方法,可以实现在0.23秒内完成100张图片VGG19正向传递。

Jeff Dean看到这个方法后,开心的转了一作的Twitter,并且评价:真是又快又容易呀!

该方法应用的正是谷歌上个月给TPU开放权限的XLA(Accelerated Linear Algebra)编译器。

XLA嵌入

首先,需要定义动态语义和静态嵌入,定义一个运行时结构来嵌入XLA值

 1const AA{T, N} = AbstractArray{T, N}
 2struct XRTArray{T, Shp, N} <: AA{T, N}
 3storage::XRTAllocation
 4# XRTArrays are constructable by
 5# conversion from regular arrays
 6function XRTArray(
 7a::Array{T, N}) where {T, N}
 8new{T, size(A), N}(transfer(a))
 9end
10# XRTArrays are constructable from a
11# remote memory allocation if
12# (T, Dims, N) are specified
13function XRTArray{T, Dims, N}(
14a::XRTAllocation) where {T, Dims, N}
15new{T, Dims, N}(a)
16end
17end
XRTArray的定义

假设我们有一个示例XLA操作’Foo’采用一个静态操作数(例如一个整数)和两个动态操作数。 我们将声明此嵌入如下:

1struct HloFoo <: HloOp{:foo}
2static_operand::Int
3end
4
5function (hlo::HloFoo)(dop1::XRTArray,
6dop2::XRTArray)
7execute(hlo, dynamic_op1, dynamic_op2)
8end

手动构建XLA嵌入:

 1# An HLO operand that generates a random
 2# uniform random number of the specificed
 3# shape and element type:
 4struct HloRng <: HloOp{:rng}
 5Type
 6Shape
 7end
 8
 9"""A function that adds random numbers to
10each entry of a 1000x1000 matrix"""
11@eval function add_rand_1000x1000(
12A::XRTArray{Float32, (1000, 1000), 2}
13random = $(HloRng(Float32,
14(1000, 1000)))()
15result = $(HloAdd())(random, A)
16return result
17end
18

将Julia映射到XLA

现在,可以将Julia代码编译到XLA,不过Julia不是用HLO运行的,而是根据Julia库提供的功能编写的。

好在,Julia使用多个调度可以很容易地表达如何用HLO实现标准库抽象。下面是一个简单的例子:

 1# Matrix-Matrix and Matrix-Vector product
 2function Base.:*(A::XRTMatrix,
 3B::Union{XRTMatrix, XRTArray})
 4ddots = DimNums((1,), (0,), (), ())
 5HloDot(ddots)(A, B)
 6end
 7Base.transpose(A::XRTArray) =
 8HloTranspose((1,0))(A)
 9# Scalar addition
10Base.:+(A::XRTArray{T, (), 0},
11B::XRTArray{T, (), 0})
12where {T<:XLAScalar} =
13GenericHloOp{:add}(T, ())(A, B)

结果比对

论文中举了VGG19正向传递和反向传递的例子。

VGG19正向传递使用Metalhead软件包Mike Innes&Contributors(2018)中的VGG19实现,它利用了Flux Innes&Contributors(2017)框架将熟悉的机器学习层(卷积层,密集层)转换为线性代数运算。

但是,重要的是,Flux框架中的每一层都只是一个常规函数,而这些函数又调用常规线性代数运算。因此用Flux表达的机器学习模型只需要简单的Julia函数,适用于本文所述的方法。

VGG反向传递则使用基于Zygote.jl编译器的AD框架Innes(2018)。Zygote对Julia代码进行操作,其输出也是Julia函数,适合重新引入Zygote以获得更高阶导数,但也适合编译到TPU。

这张表格展示编译到XLA后,Metalhead.jl VGG19正向传递和反后传递的指令计数细分。

传送门

Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUs 作者:Keno Fischer, Elliot Saba https://arxiv.org/abs/1810.09868

原文发布于微信公众号 - 量子位(QbitAI)

原文发表时间:2018-10-25

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏生信技能树

二代测序数据拼接之原理篇

前前后后接触了一些基因组和转录组拼接的工作,而且后期还会持续进行。期间遇到了各种各样莫名其妙的坑,也尝试了一些不同的方法和软件,简单做一个阶段性小结,本篇是原理...

1.8K50
来自专栏云时之间

NLP系列学习:常用的语言平滑模型

语言模型常见的平滑算法就那几种,一般的教程都不提分几种的模式、分类。 不过在MIT的NLP课程ppt中总结说有三种模式:Discounting, Interp...

33560
来自专栏数说工作室

【SAS Says】基础篇:8. 相关、回归等初步统计

SAS是一个专业的统计软件,前面我们介绍了很多数据管理、输出美化的东西,本节终于要介绍一点SAS做统计的知识了,不过,在基础篇中我们只大概介绍一下,更多统计分析...

34660
来自专栏数据结构与算法

P3150 pb的游戏(1)

题目背景 (原创) 有一天 pb和zs玩游戏 你需要帮zs求出每局的胜败情况 题目描述 游戏规则是这样的: 每次一个人可以对给出的数进行分割,将其割成两个非零自...

35370
来自专栏xingoo, 一个梦想做发明家的程序员

推荐系统那点事 —— 基于Spark MLlib的特征选择

在机器学习中,一般都会按照下面几个步骤:特征提取、数据预处理、特征选择、模型训练、检验优化。那么特征的选择就很关键了,一般模型最后效果的好坏往往都是跟特征的选...

30490
来自专栏null的专栏

优化算法——遗传算法

与遗传算法的第一次接触 遗传算法是我进入研究生阶段接触的第一个智能算法,从刚开始接触,到后来具体去研究,再到后来利用遗传算法完成了水利水电的程序设计比赛,整个过...

2.2K60
来自专栏null的专栏

数据结构和算法——动态规划

一、动态规划的思想     动态规划(dynamic programming)是一种算法设计的思想,主要是将一个问题划分成几个更小的问题,并对这样更小的问题进行...

34640
来自专栏数说工作室

【SAS Says】基础篇:基本统计、相关分析与回归分析

特别说明:本节【SAS Says】基础篇:SAS宏初步,用的是数说君学习《The little SAS book》时的中文笔记,我们认为这是打基础的最好选择 S...

49550
来自专栏行者常至

007.python科学计算库matplotlib(下)

版权声明:本文为博主原创文章,允许转载,请标明出处。 https://blog.csdn.net/qwdafedv/article/deta...

26420
来自专栏前端说吧

JS-缓冲运动基础结构

课程来源路径:智能社得开发课程:https://ke.qq.com/webcourse/index.html#course_id=152997&term_id=...

9710

扫码关注云+社区

领取腾讯云代金券