动态形状

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

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