Google在雲計算加速器TPU進一步支持深度學習函數庫PyTorch,大多數的PyTorch程序,僅需使用Python組件PyTorch/XLA,進行最小程度的修改,就能夠使用TPU來加速機器學習工作負載,且在TPU還不支持的操作,自動退回到以CPU運算,而PyTorch/XLA也會生成詳細的報告,助PyTorch用戶發現應用程序瓶頸,在Cloud TPU上,能夠高性能的執行機器學習工作負載。

在2019年PyTorch開發者大會上,臉書、Google以及Salesforce Research的工程師開始了PyTorch-TPU項目,目標是要讓PyTorch社羣,可以更容易運用Cloud TPU高效運算能力,因此該團隊創建了PyTorch/XLA組件,結合PyTorch API與XLA線性代數編譯器,讓PyTorch可以連接到Cloud TPU,將TPU核心作爲設備使用。

XLA是Google在2018年推出的優化編譯器,通過優划算法運算,可以提升機器學習模型的執行速度,XLA支持處理器以及繪圖卡,也能夠在Google TPU上執行,通過XLA可以讓Cloud TPU更容易支持PyTorch。

最新發布的PyTorch/XLA,使Cloud TPU正式支持PyTorch 1.6,在基礎設施層提供模型平行運算能力,允許開發者可以在多個TPU核心上,分散大規模嵌入表,因此有能力創建許多過去無法完成的應用。另外,Google在雲計算上發佈的深度學習虛擬機(DLVM)鏡像文件,其中包括PyTorch等各種深度學習框架,而PyTorch/XLA 1.6現在預裝在DLVM中,並且已經對Cloud TPU進行優化,供用戶快速激活開始工作。

爲了讓開發者更好上手PyTorch/XLA,Google雲計算還提供一組常用深度學習模型開源實例,以及相關的教程,這些模型包括使用ImageNet資料集的圖像分類任務ResNet -50,還有Transformer、RoBERTa以及BERT等知名模型,也能使用PyTorch 1.6才加入,由臉書開發的深度學習推薦模型(DLRM)。官方提到,在Cloud TPU訓練這些模型,大部分的情況僅需要很少的程序代碼更改。

艾倫人工智能研究所(Allen Institute for AI,A12)最近在多個項目中,開始使用PyTorch/XLA以及Cloud TPU進行研究,他們利用PyTorch/XLA在最新的語言模型中,加入可視化組件,以改善語言理解能力,A12科學家Matthew Peters認爲,雖然PyTorch/XLA目前還只是新技術,但提供已經投入使用PyTorch的組織,一個加速訓練機器學習模型的新運算平臺。

相關文章