近日,谷歌宣布云端加速器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训练这些模型,大部分的情况仅需修改少量代码即可实现。
,