近日,谷歌宣布雲端加速器TPU開始進一步支持深度學習函式庫PyTorch,僅需使用Python套件PyTorch/XLA,進行最小程度的修改,就能夠使用TPU來加速機器學習工作負載,且在TPU還不支援的操作,自動退回到以CPU運算,同時PyTorch/XLA也會生成詳細的報告,幫助使用者發現應用程序瓶頸,高效執行機器學習工作負載。
XLA是Google在2018年推出的最佳化編譯器,通過最佳化算法運算,可以提升機器學習模型的執行速度,XLA支援處理器以及繪圖卡,也能夠在Google TPU上執行,并可以讓Cloud TPU更容易支持PyTorch。
在最新的PyTorch/XLA上,使Cloud TPU正式支援PyTorch 1.6,在基礎設施層提供模型平行運算能力,允許開發者可以在多個TPU核心上,分散大規模嵌入表,因此有能力創建許多過去無法實現的應用。
另外,谷歌還在雲端上發布的深度學習虛拟機器(DLVM)映像檔,其中包括PyTorch等各種深度學習框架,而 PyTorch/XLA 1.6現在預裝在DLVM中,并且已經對Cloud TPU進行最佳化。
為了讓開發者更好上手PyTorch/XLA,Google雲端還提供一組常用深度學習模型開源實作,以及相關的教程,這些模型包括使用ImageNet資料集的圖像分類任務ResNet-50,還有Transformer、RoBERTa以及BERT等知名模型,也能使用PyTorch 1.6才加入,由臉書開發的深度學習推薦模型(DLRM)。官方提到,在Cloud TPU訓練這些模型,大部分的情況僅需修改少量代碼即可實現。
,