动态形状

动态形状是指张量形状的可变性质，其形状取决于另一个上游张量的值。例如

>>> import torch, torch_xla
>>> in_tensor  = torch.randint(low=0, high=2, size=(5,5), device='xla:0')
>>> out_tensor = torch.nonzero(in_tensor)

`out_tensor` 的形状取决于 `in_tensor` 的值，并受 `in_tensor` 形状的限制。换句话说，如果您执行

>>> print(out_tensor.shape)
torch.Size([<=25, 2])

您会看到第一个维度取决于 `in_tensor` 的值，其最大值为 25。我们将第一个维度称为动态维度。第二个维度不依赖于任何上游张量，因此我们将其称为静态维度。

动态形状可以进一步分为有界动态形状和无界动态形状。

• 有界动态形状：指其动态维度受静态值限制的形状。它适用于需要静态内存分配的加速器（例如 TPU）。

• 无界动态形状：指其动态维度可以无限大的形状。它适用于不需要静态内存分配的加速器（例如 GPU）。

目前，仅支持有界动态形状，且该功能处于实验阶段。

有界动态形状

目前，我们支持在 TPU 上使用动态大小输入的 MLP（多层感知机）模型。

此功能由标志 `XLA_EXPERIMENTAL="nonzero:masked_select"` 控制。要启用此功能运行模型，您可以执行

XLA_EXPERIMENTAL="nonzero:masked_select:masked_scatter" python your_scripts.py

以下是我们在运行 MLP 模型 100 次迭代时获得的一些数据

	无动态形状	有动态形状
端到端训练时间	29.49	20.03
编译次数	102	49
编译缓存命中次数	198	1953

动态形状的动机之一是减少形状在迭代之间不断变化时产生的过度重新编译次数。从上图可以看出，编译次数减少了一半，从而缩短了训练时间。

要试用此功能，请运行

XLA_EXPERIMENTAL="nonzero:masked_select" PJRT_DEVICE=TPU python3 pytorch/xla/test/ds/test_dynamic_shape_models.py TestDynamicShapeModels.test_backward_pass_with_dynamic_input

有关我们未来计划如何在 PyTorch/XLA 中扩展动态形状支持的更多详细信息，请随时查阅我们的 RFC。