阅读(554) (0)

PyTorch XLA 设备上的 PyTorch

2020-09-11 10:49:09 更新
原文: http://pytorch.org/xla/

PyTorch 使用 torch_xla 软件包在 XPU 设备(如 TPU)上运行。 本文档介绍了如何在这些设备上运行模型。

创建 XLA 张量

PyTorch / XLA 向 PyTorch 添加了新的xla设备类型。 此设备类型的工作方式与其他 PyTorch 设备类型一样。 例如,以下是创建和打印 XLA 张量的方法:

import torch
import torch_xla
import torch_xla.core.xla_model as xm


t = torch.randn(2, 2, device=xm.xla_device())
print(t.device)
print(t)

此代码应该看起来很熟悉。 PyTorch / XLA 使用与常规 PyTorch 相同的界面,但有一些附加功能。 导入torch_xla会初始化 PyTorch / XLA,xm.xla_device()会返回当前的 XLA 设备。 根据您的环境,这可能是 CPU 或 TPU。

XLA 张量是 PyTorch 张量

可以像 CPU 或 CUDA 张量一样在 XLA 张量上执行 PyTorch 操作。

例如,可以将 XLA 张量添加在一起:

t0 = torch.randn(2, 2, device=xm.xla_device())
t1 = torch.randn(2, 2, device=xm.xla_device())
print(t0 + t1)

或乘以矩阵:

print(t0.mm(t1))

或与神经网络模块一起使用:

l_in = torch.randn(10, device=xm.xla_device())
linear = torch.nn.Linear(10, 20).to(xm.xla_device())
l_out = linear(l_in)
print(l_out)

与其他设备类型一样,XLA 张量仅可与同一设备上的其他 XLA 张量一起使用。 所以代码像

l_in = torch.randn(10, device=xm.xla_device())
linear = torch.nn.Linear(10, 20)
l_out = linear(l_in)
print(l_out)
## Input tensor is not an XLA tensor: torch.FloatTensor

由于 torch.nn.Linear 模块在 CPU 上,因此将引发错误。

在 XLA 设备上运行模型

建立新的 PyTorch 网络或转换现有网络以在 XLA 设备上运行仅需要几行 XLA 专用代码。 以下代码片段突出显示了在单个设备,具有 XLA 并行处理功能的多个设备或具有 XLA 多线程的多个线程上运行时的这些行。

在单个 XLA 设备上运行

以下代码片段显示了单个 XLA 设备上的网络训练:

import torch_xla.core.xla_model as xm


device = xm.xla_device()
model = MNIST().train().to(device)
loss_fn = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)


for data, target in train_loader:
  optimizer.zero_grad()
  data = data.to(device)
  target = target.to(device)
  output = model(data)
  loss = loss_fn(output, target)
  loss.backward()


  xm.optimizer_step(optimizer, barrier=True)

此代码段突出显示了切换模型以在 XLA 上运行非常容易。 模型定义,数据加载器,优化器和训练循环可在任何设备上运行。 唯一的 XLA 特定代码是几行代码,这些代码获取 XLA 设备并以屏障进入优化程序。 在每次训练迭代结束时调用xm.optimizer_step(optimizer, barrier=True)都会使 XLA 执行其当前图形并更新模型的参数。 有关 XLA 如何创建图形和运行操作的更多信息,请参见 XLA Tensor Deep Dive 。

在具有并行处理功能的多个 XLA 设备上运行

通过在多个 XLA 设备上运行,PyTorch / XLA 可以轻松加速训练。 以下代码段显示了如何:

import torch_xla.core.xla_model as xm
import torch_xla.distributed.parallel_loader as pl
import torch_xla.distributed.xla_multiprocessing as xmp


def _mp_fn(index):
  device = xm.xla_device()
  para_loader = pl.ParallelLoader(train_loader, [device])


  model = MNIST().train().to(device)
  loss_fn = nn.NLLLoss()
  optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)


  for data, target in para_loader.per_device_loader(device):
    optimizer.zero_grad()
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward()
    xm.optimizer_step(optimizer)


if __name__ == '__main__':
  xmp.spawn(_mp_fn, args=())

此多设备代码段和先前的单设备代码段之间存在三个区别:

  • xmp.spawn()创建分别运行 XLA 设备的进程。
  • ParallelLoader将训练数据加载到每个设备上。
  • xm.optimizer_step(optimizer)不再需要障碍。 ParallelLoader 自动创建用于评估图形的 XLA 障碍。

模型定义,优化器定义和训练循环保持不变。

请参阅完整的并行处理示例,以获取更多关于在具有并行处理功能的多个 XLA 设备上训练网络的信息。

通过多线程在多个 XLA 设备上运行

使用进程(请参见上文)在多个 XLA 设备上运行比使用线程更可取。 但是,如果您想使用线程,则 PyTorch / XLA 具有DataParallel接口。 以下代码片段显示了具有多个线程的相同网络训练:

import torch_xla.core.xla_model as xm
import torch_xla.distributed.data_parallel as dp


devices = xm.get_xla_supported_devices()
model_parallel = dp.DataParallel(MNIST, device_ids=devices)


def train_loop_fn(model, loader, device, context):
  loss_fn = nn.NLLLoss()
  optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)


  model.train()
  for _, (data, target) in loader:
    optimizer.zero_grad()
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward()
    xm.optimizer_step(optimizer)


for epoch in range(1, num_epochs + 1):
  model_parallel(train_loop_fn, train_loader)

多线程和并行处理代码之间的唯一区别是:

  • 使用xm.get_xla_supported_devices()在同一过程中获取多个设备。
  • 该模型包装在dp.DataParallel中,并通过了训练循环和数据加载器。

有关在多 XLA 设备上使用多线程训练网络的更多信息,请参见完整的多线程示例

XLA Tensor 深潜

使用 XLA 张量和设备仅需要更改几行代码。 但是,即使 XLA 张量的行为很像 CPU 和 CUDA 张量,其内部结构也不同。 本节描述了 XLA 张量独特的原因。

XLA 张量是懒惰的

CPU 和 CUDA 张量立即启动操作或急切启动。 另一方面,XLA 张量是惰性。 他们将操作记录在图形中,直到需要结果为止。 这样推迟执行,XLA 可以对其进行优化。 例如,多个单独操作的图形可能会融合为一个优化操作。

懒惰执行通常对调用者不可见。 当在 XLA 设备和 CPU 之间复制数据时,PyTorch / XLA 自动构建图形,将它们发送到 XLA 设备,并进行同步。 采取优化程序步骤时插入屏障会显式同步 CPU 和 XLA 设备。

XLA 张量和 bFloat16

当在 TPU 上运行时,PyTorch / XLA 可以使用 bfloat16 数据类型。 实际上,PyTorch / XLA 在 TPU 上处理浮点类型(torch.floattorch.double)的方式有所不同。 此行为由XLA_USE_BF16环境变量控制:

  • 默认情况下,TPU 上的torch.floattorch.double均为torch.float
  • 如果设置了XLA_USE_BF16,则 TPU 上的torch.floattorch.double均为bfloat16

TPU 上的 XLA 张量将始终报告其 PyTorch 数据类型,而不管其使用的实际数据类型是什么。 这种转换是自动且不透明的。 如果将 TPU 上的 XLA 张量移回 CPU,它将从其实际数据类型转换为其 PyTorch 数据类型。

内存布局

XLA 张量的内部数据表示对于用户而言是不透明的。 它们不公开其存储,并且它们总是看起来是连续的,这与 CPU 和 CUDA 张量不同。 这使 XLA 可以调整张量的内存布局以获得更好的性能。

将 XLA 张量移入和移出 CPU

XLA 张量可以从 CPU 移到 XLA 设备,也可以从 XLA 设备移到 CPU。 如果移动了视图,则其视图的数据将被复制到另一台设备,并且不会保留视图关系。 换句话说,将数据复制到另一设备后,它与先前的设备或其上的任何张量都没有关系。

保存和加载 XLA 张量

在保存之前,应将 XLA 张量移至 CPU,如以下代码段所示:

import torch
import torch_xla
import torch_xla.core.xla_model as xm


device = xm.xla_device()


t0 = torch.randn(2, 2, device=device)
t1 = torch.randn(2, 2, device=device)


tensors = (t0.cpu(), t1.cpu())


torch.save(tensors, 'tensors.pt')


tensors = torch.load('tensors.pt')


t0 = tensors[0].to(device)
t1 = tensors[1].to(device)

这使您可以将加载的张量放置在任何可用设备上。

根据以上有关将 XLA 张量移至 CPU 的说明,使用视图时必须格外小心。 建议不要在保存张量并将其移至目标设备后重新创建视图,而不必保存视图。

可以直接保存 XLA 张量,但不建议这样做。 XLA 张量始终会加载回保存它们的设备,如果该设备不可用,加载将会失败。 与所有 PyTorch 一样,PyTorch / XLA 正在积极开发中,这种行为将来可能会改变。

进一步阅读

其他文档可在 PyTorch / XLA 存储库中找到。 在此处可以找到在 TPU 上运行网络的更多示例。

PyTorch / XLA API

xla_model

torch_xla.core.xla_model.xla_device(n=None, devkind=None)¶

返回 XLA 设备的给定实例。

参数

  • n (python:int , 可选)–要返回的特定实例(普通)。 如果指定,将返回特定的 XLA 设备实例。 否则,将返回 <cite>devkind</cite> 的第一个设备。
  • devkind (字符串... , 可选)–如果指定,则为 <cite>TPU</cite> ,<cite>中的一个 GPU</cite> 或 <cite>CPU</cite> (当前未实现“ GPU” XLA 设备)。

退货

具有所请求实例的<cite>torch设备</cite>。

torch_xla.core.xla_model.get_xla_supported_devices(devkind=None, max_devices=None)¶

返回给定类型的受支持设备的列表。

Parameters

  • devkind (string...__, optional) – If specified, one of <cite>TPU</cite>, <cite>GPU</cite> or <cite>CPU</cite> (the 'GPU' XLA device is currently not implemented).
  • max_devices (python:int , 可选)–此类设备的最大返回数量。

Returns

设备字符串列表。

torch_xla.core.xla_model.xrt_world_size(defval=1)¶

检索参与复制的设备数。

Parameters

defval (python:int , 可选)–如果没有可用的复制信息,则返回默认值。 默认值:1

Returns

参与复制的设备数。

torch_xla.core.xla_model.get_ordinal(defval=0)¶

检索当前进程的复制序号。

序数范围从 0 到 <cite>xrt_world_size()</cite>减 1。

Parameters

defval (python:int , 可选)–如果没有可用的复制信息,则返回默认值。 默认值:0

Returns

当前进程的复制序号。

torch_xla.core.xla_model.is_master_ordinal()¶

检查当前进程是否为主序(0)。

Returns

一个布尔值,指示当前进程是否是主序。

torch_xla.core.xla_model.optimizer_step(optimizer, barrier=False, optimizer_args={})¶

运行提供的优化器步骤并发出 XLA 设备步骤计算。

Parameters

  • 优化器(torch.Optimizer)–需要调用其 <cite>step()</cite>函数的<cite>torch.optim器</cite>实例。 <cite>step()</cite>函数将使用名为 <cite>optimizer_args</cite> 的参数调用。
  • 屏障 (bool , 可选)–是否应在此 API 中发布 XLA 张量屏障。 如果使用 PyTorch XLA <cite>ParallelLoader</cite> 或 <cite>DataParallel</cite> 支持,则不需要这样做,因为 XLA 数据加载器迭代器 <cite>next()</cite>调用会发出屏障。 默认值:False
  • optimizer_args (dict , 可选)–为 <cite>optimizer.step()</cite>调用的命名参数字典。

Returns

<cite>Optimizer.step()</cite>调用返回的值相同。

Distributed

class torch_xla.distributed.parallel_loader.ParallelLoader(loader, devices, batchdim=0, fixed_batch_size=False, loader_prefetch_size=8, device_prefetch_size=4)¶

使用背景数据上传包装现有的 PyTorch DataLoader。

Parameters

  • 加载器(torch.utils.data.DataLoader)–要包装的 PyTorch DataLoader。
  • 设备(<cite>torch设备</cite>…)–必须将数据发送到的设备列表。 <cite>加载器</cite>返回的第 i 个样本将发送到 <cite>devices [i%len(devices)]</cite> 。
  • batchdim (python:int , 可选)–保留批大小的尺寸。 默认值:0
  • fixed_batch_size (bool , 可选)–确保发送给设备的所有批次大小均相同。 一旦发现不匹配的批处理大小,原始的<cite>加载程序</cite>迭代就会停止。 默认值:False
  • loader_prefetch_size (python:int , 可选)–线程从[[ <cite>loader</cite> ,由工作线程处理,这些工作线程会将数据上传到设备。 默认值:8
  • device_prefetch_size (python:int , 可选)–每个设备队列的最大大小,工作线程在其中存放张量, 已经发送到设备。 默认值:4
per_device_loader(device)¶

检索给定设备的加载程序对象。

Parameters

设备(<cite>torch设备</cite>)–正在请求设备整个装载程序。

Returns

<cite>设备</cite>的数据加载器。

class torch_xla.distributed.data_parallel.DataParallel(network, device_ids=None)¶

使用线程以复制模式启用模型网络的执行。

Parameters

  • 网络(torch.nn.Module或可调用)–模型的网络。 <cite>torch.nn.Module</cite> 的子类,或者是返回 <cite>torch.nn.Module</cite> 子类的可调用对象。
  • device_ids (字符串…或torch.device…)–应在其上进行复制的设备的列表。 如果列表为空,则网络将在 PyTorch CPU 设备上运行。
__call__(loop_fn, loader, fixed_batch_size=False, batchdim=0)¶

进行一次 EPOCH 训练/测试。

Parameters

  • loop_fn (可调用的)–在分配给参与复制的每个设备的每个线程上调用的函数。 该函数将使用 <cite>def loop_fn(model,device_loader,device,context)</cite>签名来调用。 其中<cite>模型</cite>是传递到 <cite>DataParallel</cite> 构造器的每个设备网络。 <cite>device_loader</cite> 是 <cite>ParallelLoader</cite> ,它将为当前<cite>设备</cite>返回样本。 <cite>上下文</cite>是每个线程/设备上下文,具有 <cite>DataParallel</cite> 对象的生存期,并且 <cite>loop_fn</cite> 可以使用它来存储需要 在不同的 EPOCH 中保持一致。
  • fixed_batch_size (bool , 可选)–参数传递给 <cite>ParallelLoader</cite> 构造函数。 默认值:False
  • batchdim (python:int , 可选)–由<cite>加载器</cite>返回的样本尺寸 批量大小。 默认值:0

Returns

每个设备上 <cite>loop_fn</cite> 返回的值的列表。

torch_xla.distributed.xla_multiprocessing.spawn(fn, args=(), nprocs=None, join=True, daemon=False)¶

启用基于并行处理的复制。

Parameters

  • fn –参与复制的每个设备要调用的功能。 将调用该函数,第一个参数是复制中进程的全局索引,然后是 <cite>args</cite> 中传递的参数。
  • args - <cite>fn</cite> 的参数。
  • nprocs –复制的进程/设备数。 目前,如果指定,则可以为 1 或最大设备数。
  • join –呼叫是否应等待生成的进程完成而阻塞。
  • 守护程序 –产生的进程是否应设置<cite>守护程序</cite>标志(请参阅 Python 并行处理 API)。

Returns

<cite>torch.multiprocessing.spawn</cite> API 返回的同一对象。

实用程序

class torch_xla.utils.utils.SampleGenerator(data, sample_count)¶

迭代器,它返回给定输入数据的多个样本。

可以代替 PyTorch <cite>DataLoader</cite> 生成合成数据。

Parameters

  • 数据 –在每个迭代器步骤应返回的数据。
  • sample_count –要返回的<cite>数据</cite>个样本的最大数量。