实践教程｜PyTorch与torch-xla的桥接

文章预览

↑ 点击 蓝字  关注极市平台 作者丨 liuyaohui 来源丨GiantPandaCV 编辑丨极市平台 极市导读   文章从XLATensor开始的溯源、注册PyTorch库实现、从PyTorch调用到torch_xla三个方面来介绍PyTorch与torch-xla的桥接  >> 加入极市CV技术交流群，走在计算机视觉的最前沿 前言 XLA (Accelerated Linear Algebra)是一个开源的机器学习编译器，对PyTorch、Tensorflow、JAX等多个深度学习框架都有支持。最初XLA实际上是跟Tensorflow深度结合的，很好地服务了Tensorflow和TPU，而与XLA的结合主要依赖于社区的支持，即torch-xla。 torch-xla在支持XLA编译的基础上，较大限度地保持了PyTorch的易用性，贴一个官方的DDP训练的例子：  import torch.distributed as dist   -import torch.multiprocessing as mp   +import torch_xla.core.xla_model as xm   +import torch_xla.distributed.parallel_loader as pl   +import torch_xla.distributed.xla_multiproce ………………………………