了“全球人工智能”拥有十多万AI产业用户,多名AI技术专家+多名AI创业企业高管+多名AI产业投资者核心用户来自:北大,清华,中科院,麻省理工,卡内基梅隆,斯坦福,哈佛,牛津,剑桥......以及谷歌,腾讯,百度,脸谱,微软,阿里,海康威视,英伟达......等全球名校和名企。
-数百位AI投资者喊你来融资
-加入AI企业高管群
-加入AI技术专家群
深度学习如今无处不在且必不可少。这次创新部分得益于可扩展的深度学习系统,比如TensorFlow、MXNet、Caffe和PyTorch。大多数现有系统针对窄范围的服务器级GPU进行了优化,并且需要在其他平台(如手机、IoT设备和专用加速器(FPGA、ASIC))上部署大量工作。随着深度学习框架和硬件后端数量不断增加,我们提出了一个统一的中间表征(IR)堆栈,可以弥补以生产力为中心的深度学习框架与性能或效率导向的硬件后端之间的差距。
我们很高兴宣布TVM的发行将解决这一问题。TVM是一个全新的框架,它可以:
优化CPU、GPU和其他专业化硬件的常规深度学习的计算量
自动转换计算图以最小化内存利用率,优化数据布局,融合计算模式
提供从现有的前端框架到裸机硬件的端到端编译,一直到浏览器可执行的Javascript
有了TVM的帮助,只需少量额外工作我们即可轻易地在手机端、嵌入式设备甚至浏览器上运行深度学习任务。TVM还为多硬件平台上的深度学习工作负载提供了统一的优化框架,其中包括依赖全新计算原语的专用加速器。
通过采用编译器社区的共同理念,并提供两个中间表征层,我们有效地将高级深度学习算法降低到多硬件后端。在今天的版本中,我们开源了包含x86、ARM、OpenCL、Metal、CUDA和Javascript优化原语的TVM工具包。目前,我们正积极地为专用硬件加速和英伟达的GEMM优化的Volta架构提供支持。
技术细节
TVM堆栈的目标在于提供一个可重复使用的工具链来将高级神经网络描述从深度学习框架前端向下编译为多硬件后端的低级机器代码。拿ApacheMXNet作为前端的例子,以下代码段演示了如何使用TVM将深层学习模型的高级描述编译为针对目标硬件定制的优化可执行模块。
挑战在于支持多个硬件后端,同时将计算、内存和能源足迹(energyfootprint)保持在最低水平。我们借鉴了编译器社区的共同理念,以弥合大量深度学习框架和硬件后端之间的差距:我们构建了一个由NNVM组成的两级中间层,一个可用于任务调度、内存管理以及TVM(用于优化计算内核的富有表现力的低级IR)的高级中间表征(IR)。
堆栈的第一级是基于计算图的表征。计算图是一个有向无环图,其表征作为节点的计算和作为边缘的数据流依赖性。这一表征非常强大:它允许我们将操作属性融入到计算图中,并指定转换规则以迭代优化计算图。这是一个被大多数现有深度学习框架的采用的惯常方法,包括TVM堆栈中的NNVM图表征、TensorFlowXLA和英特尔的ngraph。
很多强大的优化可由图优化框架支持。例如,我们提供了一个子线性内存优化功能,允许用户在单块GPU上训练层的ImageNetResNet。
然而,我们发现只基于IR的计算图并不足以应对支持不同硬件后端的挑战。原因是像卷积或矩阵乘法一样的图运算符可以非常不同的方式映射和优化每个硬件后端。这些特定硬件的优化在内存布局、并行线程模式、缓存访问模式和硬件原语选择方面的变数可能很大。我们希望能够以通用表征的方式明确地表达这些优化旋钮(optimizationknob),以有效地导航优化空间。
我们构建了一个低级表征以解决这一问题。该表征基于索引公式,并额外支持递推计算。
低级IR采用现有的图像处理语言(比如Halide或darkroom)的原理以制定一个富有表现力的深度学习DSL。TVM构建了由循环转换工具比如loopy、多面体分析启发的图优化。我们也从深度学习框架比如MXNet、TensorFlow、Theano使用的数据流描述语言汲取灵感,然后在调度阶段对TVM中描述的算法进行处理,以应用针对目标硬件后端定制的转换。
TVM包括通常出现在CPU优化框架中的标准转换原语。更重要的是,TVM通过利用线程协作模式、数据布局转换和强大的新计算原语,集成了针对GPU的全新优化原语。TVM与NNVM的结合为优化软件堆栈中深度学习工作负载提供了一个很好的机会,从而实现了图级别与运算符级别优化的联合计算。
多语言和平台支持
TVM的众多优势之一在于它可以支持多种语言和平台。TVM框架中有两个组件:编译器栈,其中包含完整的优化库以产生优化的机器代码;以及轻量级运行环境,提供在不同平台上部署编译模块所需的可移植性。
TVM目前支持嵌入式编译器堆栈的Python和C++接口。我们在设计中最大限度地重用框架,以便编译器堆栈的改进可以在Python和C++组件之间互换使用。我们还提供了一个轻量级运行环境,可以在包括Android、iOS、树莓派和Web浏览器在内的平台上直接运行诸如Javascript、Java、Python和C++等语言的TVM编译代码。
远程部署和执行
TVM支持使用TVMPRC进行交叉编译,测试嵌入式设备,这是一种轻量级界面,用于在远程嵌入式设备上部署和执行TVM交叉编译模块。我们TVM用户提供了易用的高级Python界面,用于在各种低级嵌入式设备上进行远程编译、优化,并测试深度学习算法。
性能
TVM目前仍然处在早期发展阶段,我们会在未来看到更多改进,但目前它已经展现出了强大的实力。
TVM为我们提供了灵活性,让我们可以探索不同硬件平台上,不同深度学习内核中丰富的优化空间。例如,TVM允许我们裁剪数据布局,为内核与平台融合模式需求。注意,构建基线库可用于更通用的问题,而TVM的优化内核通过自动调整过程对人们评估的工作负载进行了大量调整。TVM可以作为快速生成专属内核的绝佳工具。
本节列出的结果仍有提升空间。
树莓派
在结果的第一部分,我们比较了TVMCPU计划与nnpack在树莓派3B上执行ResNet任务的表现。由于时间有限,我们使用TVM实现直接卷积,而nnpack对3×3内核执行winograd转换。
我们可以看到,使用TVM自动调整内核,我们可以获得类似于nnpack中的手动优化内核的性能(实验在树莓派上进行)。
GPU结果
PS.该测试结果由LeyuanWang(AWS/UCDavis)、YuweiHu(图森未来)与WeitangLiu(AWS/UCDavis)提供。
为了进行概念证明,我们构建了一种端到端编译管道,可以将MXNet模型编译成TVM执行图。我们通过自动融合运算符并使TVM生成融合的内核,在图形节点之间和之中应用优化。我们对MobileNet在ImageNet上的运行表现进行基准测试,结果如下:
我们可以看到,TVM在速度上快过基准方法。更有意思的是,内核融合可以带来更快的速度提升。值得一提的是,TVM可以自行生成所有优化GPU内核,而不依赖于外部库,如CuDNN。
我们正在开展新的实验,期待获得更好的结果。