# XLA 设备上的 PyTorch

> 原文： [http://pytorch.org/xla/](http://pytorch.org/xla/)

PyTorch 使用 [torch_xla 软件包](https://github.com/pytorch/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-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 障碍。

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

请参阅[完整的并行处理示例](https://github.com/pytorch/xla/blob/master/test/test_train_mp_mnist.py)，以获取更多关于在具有并行处理功能的多个 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 设备上使用多线程训练网络的更多信息，请参见[完整的多线程示例](https://github.com/pytorch/xla/blob/master/test/test_train_mnist.py)。

## 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](https://en.wikipedia.org/wiki/Bfloat16_floating-point_format) 数据类型。 实际上，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 存储库](https://github.com/pytorch/xla/)中找到。 在[此处](https://github.com/pytorch-tpu/examples)可以找到在 TPU 上运行网络的更多示例。

# PyTorch / XLA API

## xla_model

* * *

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

返回 XLA 设备的给定实例。

参数

*   **n**  (_python：int_ _，_ _可选_）–要返回的特定实例(普通）。 如果指定，将返回特定的 XLA 设备实例。 否则，将返回 &lt;cite&gt;devkind&lt;/cite&gt; 的第一个设备。

*   **devkind** (_字符串..._ _，_ _可选_）–如果指定，则为 &lt;cite&gt;TPU&lt;/cite&gt; ，&lt;cite&gt;中的一个 GPU&lt;/cite&gt; 或 &lt;cite&gt;CPU&lt;/cite&gt; (当前未实现“ GPU” XLA 设备）。

退货

具有所请求实例的&lt;cite&gt;torch设备&lt;/cite&gt;。

* * *

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

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

Parameters

*   **devkind** (_string...__,_ _optional_) – If specified, one of &lt;cite&gt;TPU&lt;/cite&gt;, &lt;cite&gt;GPU&lt;/cite&gt; or &lt;cite&gt;CPU&lt;/cite&gt; (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 到 &lt;cite&gt;xrt_world_size(）&lt;/cite&gt;减 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`）–需要调用其 &lt;cite&gt;step(）&lt;/cite&gt;函数的&lt;cite&gt;torch.optim器&lt;/cite&gt;实例。 &lt;cite&gt;step(）&lt;/cite&gt;函数将使用名为 &lt;cite&gt;optimizer_args&lt;/cite&gt; 的参数调用。

*   **屏障** (_bool_ _，_ _可选_）–是否应在此 API 中发布 XLA 张量屏障。 如果使用 PyTorch XLA &lt;cite&gt;ParallelLoader&lt;/cite&gt; 或 &lt;cite&gt;DataParallel&lt;/cite&gt; 支持，则不需要这样做，因为 XLA 数据加载器迭代器 &lt;cite&gt;next(）&lt;/cite&gt;调用会发出屏障。 默认值：False

*   **optimizer_args**  (_dict_ _，_ _可选_）–为 &lt;cite&gt;optimizer.step(）&lt;/cite&gt;调用的命名参数字典。

Returns

&lt;cite&gt;Optimizer.step(）&lt;/cite&gt;调用返回的值相同。

## 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。

*   **设备**(&lt;cite&gt;torch设备&lt;/cite&gt;…）–必须将数据发送到的设备列表。 &lt;cite&gt;加载器&lt;/cite&gt;返回的第 i 个样本将发送到 &lt;cite&gt;devices [i％len(devices）]&lt;/cite&gt; 。

*   **batchdim**  (_python：int_ _，_ _可选_）–保留批大小的尺寸。 默认值：0

*   **fixed_batch_size**  (_bool_ _，_ _可选_）–确保发送给设备的所有批次大小均相同。 一旦发现不匹配的批处理大小，原始的&lt;cite&gt;加载程序&lt;/cite&gt;迭代就会停止。 默认值：False

*   **loader_prefetch_size**  (_python：int_ _，_ _可选_）–线程从[[ &lt;cite&gt;loader&lt;/cite&gt; ，由工作线程处理，这些工作线程会将数据上传到设备。 默认值：8

*   **device_prefetch_size**  (_python：int_ _，_ _可选_）–每个设备队列的最大大小，工作线程在其中存放张量， 已经发送到设备。 默认值：4

* * *

```
per_device_loader(device)¶
```

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

Parameters

**设备**(&lt;cite&gt;torch设备&lt;/cite&gt;）–正在请求设备整个装载程序。

Returns

&lt;cite&gt;设备&lt;/cite&gt;的数据加载器。

* * *

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

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

Parameters

*   **网络**(`torch.nn.Module`或可调用）–模型的网络。 &lt;cite&gt;torch.nn.Module&lt;/cite&gt; 的子类，或者是返回 &lt;cite&gt;torch.nn.Module&lt;/cite&gt; 子类的可调用对象。

*   **device_ids** (字符串…或`torch.device`…）–应在其上进行复制的设备的列表。 如果列表为空，则网络将在 PyTorch CPU 设备上运行。

* * *

```
__call__(loop_fn, loader, fixed_batch_size=False, batchdim=0)¶
```

进行一次 EPOCH 训练/测试。

Parameters

*   **loop_fn** (可调用_的_）–在分配给参与复制的每个设备的每个线程上调用的函数。 该函数将使用 &lt;cite&gt;def loop_fn(model，device_loader，device，context）&lt;/cite&gt;签名来调用。 其中&lt;cite&gt;模型&lt;/cite&gt;是传递到 &lt;cite&gt;DataParallel&lt;/cite&gt; 构造器的每个设备网络。 &lt;cite&gt;device_loader&lt;/cite&gt; 是 &lt;cite&gt;ParallelLoader&lt;/cite&gt; ，它将为当前&lt;cite&gt;设备&lt;/cite&gt;返回样本。 &lt;cite&gt;上下文&lt;/cite&gt;是每个线程/设备上下文，具有 &lt;cite&gt;DataParallel&lt;/cite&gt; 对象的生存期，并且 &lt;cite&gt;loop_fn&lt;/cite&gt; 可以使用它来存储需要 在不同的 EPOCH 中保持一致。

*   **fixed_batch_size**  (_bool_ _，_ _可选_）–参数传递给 &lt;cite&gt;ParallelLoader&lt;/cite&gt; 构造函数。 默认值：False

*   **batchdim**  (_python：int_ _，_ _可选_）–由&lt;cite&gt;加载器&lt;/cite&gt;返回的样本尺寸 批量大小。 默认值：0

Returns

每个设备上 &lt;cite&gt;loop_fn&lt;/cite&gt; 返回的值的列表。

* * *

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

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

Parameters

*   **fn** –参与复制的每个设备要调用的功能。 将调用该函数，第一个参数是复制中进程的全局索引，然后是 &lt;cite&gt;args&lt;/cite&gt; 中传递的参数。

*   **args** - &lt;cite&gt;fn&lt;/cite&gt; 的参数。

*   **nprocs** –复制的进程/设备数。 目前，如果指定，则可以为 1 或最大设备数。

*   **join** –呼叫是否应等待生成的进程完成而阻塞。

*   **守护程序** –产生的进程是否应设置&lt;cite&gt;守护程序&lt;/cite&gt;标志(请参阅 Python 并行处理 API）。

Returns

&lt;cite&gt;torch.multiprocessing.spawn&lt;/cite&gt; API 返回的同一对象。

## 实用程序

* * *

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

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

可以代替 PyTorch &lt;cite&gt;DataLoader&lt;/cite&gt; 生成合成数据。

Parameters

*   **数据** –在每个迭代器步骤应返回的数据。

*   **sample_count** –要返回的&lt;cite&gt;数据&lt;/cite&gt;个样本的最大数量。