是否能超越Python不好說,不過Julia語言發展確實非常迅速,它可以視為同時具備了Python的靈活性與C的速度,但目前TensorFlow和PyTorch等框架官方都不支持Julia語言。因此近日有研究者借助XLA底層編譯器為Julia構建TPU支持,他們表示該方法能夠將Julia程序編寫的VGG19模型融合到TPU可執行文件中,并調用TPU實現高效計算。而Google.ai的負責人JeffDean在推特上也表示「Julia+TPUs=fastandeasilyexpressibleMLcomputations!」
1.引言
過去的幾年里推動機器學習技術穩步發展的根本性改變之一是訓練和優化機器學習模型的巨大計算力。許多技術都是很年前就已經提出,唯有近幾年提升的計算力可以為現實世界的問題提供足夠優質的解決方案。這些計算能力的很大一部分是通過GPU獲取的,其針對向量的計算能力最初是為圖形而設計的,但機器學習模型通常需要執行復雜的矩陣運算,因此GPU同樣表現出了非常好的性能。
這些方法及GPU在現實世界,尤其是在機器學習領域的成功引發了硬件設計者的一系列創新,他們致力于為機器學習工作負載研發新的加速器。然而,盡管GPU很長一段時間都在CUDA等軟件系統發力,但這些庫通常不會擴展到新的非GPU加速器,為這些加速器開發軟件仍然是一大挑戰。
2017年,谷歌宣布他們將通過云服務向大眾提供他們專有的TPU機器學習加速器。最初,TPU的使用局限于根據谷歌TensorFlow機器學習框架編寫的應用。幸運的是,2018年9月,谷歌通過底層XLA(AcceleratedLinearAlgebra)編譯器的IR開放了TPU的訪問權限。這個IR是一個通用的優化編譯器,用于表達線性代數基元的任意計算,因此為使用TPU的非TensorFlow用戶以及非機器學習工作負載提供了良好的基礎。
在本文中,我們介紹了使用這個接口編譯通用Julia代碼的初步工作,它們可以進一步訪問谷歌云的TPU。這一方法與TensorFlow(Abadietal.,2016)采用的方法形成對比,后者沒有編譯python代碼,而是先用Python構建一個計算圖,然后再對這個計算圖進行編譯。它在美學上類似于JAX(Frostigetal.,2018),JAX的目標是通過跟蹤和Offload高級數組運算來OffloadPython本身編寫的計算。然而重要的是,我們不依賴于追蹤,而是利用Julia的靜態分析和編譯能力來編譯整個程序,包括傳遞到設備端的所有控制流。
值得一提的是,我們的方法允許用戶在編寫模型時充分利用Julia語言的表現力。這些表現力主要體現在一些高級特征上,如多重派發、高階函數和現有庫,如微分方程求解器(Rackauckas&Nie,2017)和通用線性代數例程等。由于只在純Julia代碼上運行,所以它也與Zygote.jl(Innes,2018)自動微分工具兼容,該工具能執行自動微分作為高級編譯過程。總的來說,我們能夠編譯使用Flux機器學習框架編寫的完整機器學習模型,將模型的前向、反向傳播及訓練回路融合成一個可執行文件,并Offload到TPU中。
論文:AutomaticFullCompilationofJuliaProgramsandMLModelstoCloudTPUs
論文鏈接:https://arxiv.org/abs/1810.09868
摘要:谷歌的云TPU是一種前景廣闊的新型機器學習工作負載硬件架構,近年來已經成就了谷歌很多里程碑式的機器學習突破。如今,谷歌已經在其云平臺上為大眾提供TPU,最近又進一步開放,允許非TensorFlow前端使用。我們描述了一種通過這一新API及谷歌XLA編譯器將Julia程序的適當部分Offload到TPU的方法和實現。我們的方法能夠將Julia程序編寫的VGG19模型及其正向傳播完全融合到單個TPU可執行文件中,以便Offload到設備上。我們的方法與Julia代碼上現有的基于編譯器的自動微分技術很好地結合在一起,因此也能夠自動獲得VGG19反向傳播并采用類似的方法將其Offload到TPU。使用我們的編譯器訪問TPU,我們能夠在0.23秒內完成批量為100張圖像的VGG19前向傳播,而CPU上的原始模型則需要52.4s。我們的實現僅需不到1000行的Julia代碼,無需根據TPU對核心Julia編譯器或任何其他Julia包進行特有的更改。
7結果
7.4在TPU上進行評估
圖2:不同批大小對應的VGG19前向傳播時長。
圖3:被編譯為XLA后,Metalhead.jlVGG19前向傳播和反向傳播的指令數分解。