PyTorch XLA 设备上的 PyTorch

原文： 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.float和torch.double）的方式有所不同。 此行为由XLA_USE_BF16环境变量控制：

• 默认情况下，TPU 上的torch.float和torch.double均为torch.float。
• 如果设置了XLA_USE_BF16，则 TPU 上的torch.float和torch.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>个样本的最大数量。