(原型) PyTorch 中的命名张量简介

作者: Richard Zou

命名张量旨在通过允许用户将显式名称与张量维度相关联来使张量更易于使用。在大多数情况下，接受维度参数的操作将接受维度名称，从而避免了按位置跟踪维度的需要。此外，命名张量使用名称在运行时自动检查 API 是否被正确使用，提供额外的安全性。名称也可以用于重新排列维度，例如，支持“按名称广播”而不是“按位置广播”。

本教程旨在作为对 1.3 版本发布中将包含的功能的指南。在本教程结束时，您将能够

• 创建具有命名维度的张量，以及删除或重命名这些维度
• 了解操作如何传播维度名称的基本原理

• 了解命名维度如何在两个关键领域使代码更清晰

  • 广播操作
  • 扁平化和非扁平化维度

最后，我们将通过使用命名张量编写一个多头注意力模块来将这些知识付诸实践。

PyTorch 中的命名张量受到 Sasha Rush 的启发，并与之协作完成。Sasha 在他 2019 年 1 月的博客文章 中提出了最初的想法和概念证明。

基础知识：命名维度

PyTorch 现在允许张量具有命名维度；工厂函数接受一个新的 names 参数，该参数将名称与每个维度相关联。这适用于大多数工厂函数，例如

• tensor
• empty
• ones
• zeros
• randn
• rand

在这里，我们构建一个具有名称的张量

import torch
imgs = torch.randn(1, 2, 2, 3, names=('N', 'C', 'H', 'W'))
print(imgs.names)

输出

('N', 'C', 'H', 'W')

与 最初的命名张量博客文章 不同，命名维度是有序的：tensor.names[i] 是 tensor 的第 i 个维度的名称。

有两种方法可以重命名 Tensor 的维度

# Method #1: set the .names attribute (this changes name in-place)
imgs.names = ['batch', 'channel', 'width', 'height']
print(imgs.names)

# Method #2: specify new names (this changes names out-of-place)
imgs = imgs.rename(channel='C', width='W', height='H')
print(imgs.names)

输出

('batch', 'channel', 'width', 'height')
('batch', 'C', 'W', 'H')

删除名称的首选方法是调用 tensor.rename(None)

imgs = imgs.rename(None)
print(imgs.names)

输出

(None, None, None, None)

无名张量（没有命名维度的张量）仍然像往常一样工作，并且在它们的 repr 中没有名称。

unnamed = torch.randn(2, 1, 3)
print(unnamed)
print(unnamed.names)

输出

tensor([[[-0.5647,  0.8112,  1.4354]],

        [[-1.1201, -2.5431,  0.1843]]])
(None, None, None)

命名张量不要求所有维度都被命名。

imgs = torch.randn(3, 1, 1, 2, names=('N', None, None, None))
print(imgs.names)

输出

('N', None, None, None)

由于命名张量可以与无名张量共存，因此我们需要一种很好的方法来编写支持命名和无名张量的命名张量感知代码。使用 tensor.refine_names(*names) 来细化维度并将无名维度提升为命名维度。细化维度被定义为满足以下约束的“重命名”

• 一个 None 维度可以被细化为具有任何名称
• 一个命名维度只能被细化为具有相同的名称。

imgs = torch.randn(3, 1, 1, 2)
named_imgs = imgs.refine_names('N', 'C', 'H', 'W')
print(named_imgs.names)

# Refine the last two dims to 'H' and 'W'. In Python 2, use the string '...'
# instead of ...
named_imgs = imgs.refine_names(..., 'H', 'W')
print(named_imgs.names)


def catch_error(fn):
    try:
        fn()
        assert False
    except RuntimeError as err:
        err = str(err)
        if len(err) > 180:
            err = err[:180] + "..."
        print(err)


named_imgs = imgs.refine_names('N', 'C', 'H', 'W')

# Tried to refine an existing name to a different name
catch_error(lambda: named_imgs.refine_names('N', 'C', 'H', 'width'))

输出

('N', 'C', 'H', 'W')
(None, None, 'H', 'W')
refine_names: cannot coerce Tensor['N', 'C', 'H', 'W'] to Tensor['N', 'C', 'H', 'width'] because 'W' is different from 'width' at index 3

大多数简单操作都会传播名称。命名张量的最终目标是所有操作都能以合理、直观的方式传播名称。在 1.3 版本发布时，已经添加了许多常见操作的支持；例如，以下是 .abs()

print(named_imgs.abs().names)

输出

('N', 'C', 'H', 'W')

访问器和约简

可以使用维度名称来引用维度，而不是位置维度。这些操作也会传播名称。索引（基本索引和高级索引）尚未实现，但已列入路线图。使用上面的 named_imgs 张量，我们可以执行以下操作

output = named_imgs.sum('C')  # Perform a sum over the channel dimension
print(output.names)

img0 = named_imgs.select('N', 0)  # get one image
print(img0.names)

输出

('N', 'H', 'W')
('C', 'H', 'W')

名称推断

名称通过一个称为**名称推断**的两个步骤过程在操作中传播

1. **检查名称**：运算符可以在运行时执行自动检查，以检查某些维度名称是否必须匹配。
2. **传播名称**：名称推断将输出名称传播到输出张量。

让我们来看一个非常小的示例，即添加两个没有广播的单维张量。

x = torch.randn(3, names=('X',))
y = torch.randn(3)
z = torch.randn(3, names=('Z',))

检查名称：首先，我们将检查这两个张量的名称是否匹配。两个名称匹配当且仅当它们相等（字符串相等）或至少其中一个为 None（None 本质上是一个特殊的通配符名称）。因此，这三个中唯一会报错的是 x + z

catch_error(lambda: x + z)

输出

Error when attempting to broadcast dims ['X'] and dims ['Z']: dim 'X' and dim 'Z' are at the same position from the right but do not match.

传播名称：通过返回两个名称中最细化的名称来统一这两个名称。对于 x + y，X 比 None 更细化。

print((x + y).names)

输出

('X',)

大多数名称推断规则都很简单，但其中一些规则可能会有意想不到的语义。让我们来看几个你可能遇到的情况：广播和矩阵乘法。

广播

命名张量不会改变广播行为；它们仍然按位置广播。但是，在检查两个维度是否可以广播时，PyTorch 也会检查这些维度的名称是否匹配。

这会导致命名张量阻止在广播操作期间进行意外对齐。在下面的示例中，我们将 per_batch_scale 应用于 imgs。

imgs = torch.randn(2, 2, 2, 2, names=('N', 'C', 'H', 'W'))
per_batch_scale = torch.rand(2, names=('N',))
catch_error(lambda: imgs * per_batch_scale)

输出

Error when attempting to broadcast dims ['N', 'C', 'H', 'W'] and dims ['N']: dim 'W' and dim 'N' are at the same position from the right but do not match.

如果没有 names，则 per_batch_scale 张量将与 imgs 的最后一位对齐，这不是我们想要的。我们实际上想通过将 per_batch_scale 与 imgs 的批次维度对齐来执行此操作。有关如何按名称对齐张量的“按名称显式广播”功能，请参见下面的内容。

矩阵乘法

torch.mm(A, B) 对 A 的第二位和 B 的第一位执行点积，返回一个具有 A 的第一位和 B 的第二位的张量。（其他矩阵乘法函数，如 torch.matmul、torch.mv 和 torch.dot，行为类似）。

markov_states = torch.randn(128, 5, names=('batch', 'D'))
transition_matrix = torch.randn(5, 5, names=('in', 'out'))

# Apply one transition
new_state = markov_states @ transition_matrix
print(new_state.names)

输出

('batch', 'out')

如你所见，矩阵乘法不会检查收缩维度是否具有相同的名称。

接下来，我们将介绍命名张量支持的两种新行为：按名称显式广播以及按名称扁平化和取消扁平化维度。

新行为：按名称显式广播

使用多维数据时，人们经常会抱怨需要 unsqueeze “虚拟”维度才能执行操作。例如，在我们之前的每批次缩放示例中，对于无名张量，我们会执行以下操作

imgs = torch.randn(2, 2, 2, 2)  # N, C, H, W
per_batch_scale = torch.rand(2)  # N

correct_result = imgs * per_batch_scale.view(2, 1, 1, 1)  # N, C, H, W
incorrect_result = imgs * per_batch_scale.expand_as(imgs)
assert not torch.allclose(correct_result, incorrect_result)

通过使用名称，我们可以使这些操作更安全（并且易于与维度数量无关）。我们提供了一个新的 tensor.align_as(other) 操作，该操作会置换张量的维度以匹配 other.names 中指定的顺序，并在适当的情况下添加大小为一的维度（tensor.align_to(*names) 也适用）。

imgs = imgs.refine_names('N', 'C', 'H', 'W')
per_batch_scale = per_batch_scale.refine_names('N')

named_result = imgs * per_batch_scale.align_as(imgs)
# note: named tensors do not yet work with allclose
assert torch.allclose(named_result.rename(None), correct_result)

新行为：按名称扁平化和取消扁平化维度

一个常见的操作是扁平化和取消扁平化维度。目前，用户使用 view、reshape 或 flatten 执行此操作；用例包括扁平化批次维度，并将张量发送到必须接收具有特定维度数量的输入的操作符（例如，conv2d 接收 4D 输入）。

为了使这些操作比 view 或 reshape 更有语义意义，我们引入了一个新的 tensor.unflatten(dim, namedshape) 方法，并更新了 flatten 以支持名称：tensor.flatten(dims, new_dim)。

flatten 只能扁平化相邻维度，但也能作用于非连续维度。必须向 unflatten 传递一个命名形状，它是一个 (dim, size) 元组列表，以指定如何取消扁平化该维度。可以在 flatten 中保存大小以供 unflatten 使用，但我们目前还没有这样做。

imgs = imgs.flatten(['C', 'H', 'W'], 'features')
print(imgs.names)

imgs = imgs.unflatten('features', (('C', 2), ('H', 2), ('W', 2)))
print(imgs.names)

输出

('N', 'features')
('N', 'C', 'H', 'W')

自动梯度支持

Autograd 目前忽略了所有张量上的名称，只将它们视为常规张量。梯度计算是正确的，但我们失去了名称提供的安全性。在路线图中，将介绍对 Autograd 的名称处理。

x = torch.randn(3, names=('D',))
weight = torch.randn(3, names=('D',), requires_grad=True)
loss = (x - weight).abs()
grad_loss = torch.randn(3)
loss.backward(grad_loss)

correct_grad = weight.grad.clone()
print(correct_grad)  # Unnamed for now. Will be named in the future

weight.grad.zero_()
grad_loss = grad_loss.refine_names('C')
loss = (x - weight).abs()
# Ideally we'd check that the names of loss and grad_loss match, but we don't
# yet
loss.backward(grad_loss)

print(weight.grad)  # still unnamed
assert torch.allclose(weight.grad, correct_grad)

输出

tensor([ 0.3588,  0.4460, -0.4983])
tensor([ 0.3588,  0.4460, -0.4983])

其他支持的功能（以及不支持的功能）

请参阅此处，以详细了解 1.3 版本支持的内容。

特别地，我们要重点介绍三个目前不支持的重要功能

• 通过 torch.save 或 torch.load 保存或加载命名张量
• 通过 torch.multiprocessing 进行多进程处理
• JIT 支持；例如，以下操作会导致错误

imgs_named = torch.randn(1, 2, 2, 3, names=('N', 'C', 'H', 'W'))


@torch.jit.script
def fn(x):
    return x


catch_error(lambda: fn(imgs_named))

输出

NYI: Named tensors are currently unsupported in TorchScript. As a  workaround please drop names via `tensor = tensor.rename(None)`.

作为解决方法，请在使用任何尚不支持命名张量的操作之前，通过 tensor = tensor.rename(None) 移除名称。

更长的示例：多头注意力

现在，我们将介绍一个完整的示例，用于实现一个常见的 PyTorch nn.Module：多头注意力。我们假设读者已经熟悉多头注意力；如需复习，请参阅 此解释 或 此解释。

我们改编了 ParlAI 中的多头注意力实现；具体而言，请参阅此处。阅读该示例中的代码；然后，与下面的代码进行比较，注意标记为 (I)、(II)、(III) 和 (IV) 的四个位置，在这些位置使用命名张量可以使代码更具可读性；我们将在代码块之后深入探讨每个位置。

import torch.nn as nn
import torch.nn.functional as F
import math


class MultiHeadAttention(nn.Module):
    def __init__(self, n_heads, dim, dropout=0):
        super(MultiHeadAttention, self).__init__()
        self.n_heads = n_heads
        self.dim = dim

        self.attn_dropout = nn.Dropout(p=dropout)
        self.q_lin = nn.Linear(dim, dim)
        self.k_lin = nn.Linear(dim, dim)
        self.v_lin = nn.Linear(dim, dim)
        nn.init.xavier_normal_(self.q_lin.weight)
        nn.init.xavier_normal_(self.k_lin.weight)
        nn.init.xavier_normal_(self.v_lin.weight)
        self.out_lin = nn.Linear(dim, dim)
        nn.init.xavier_normal_(self.out_lin.weight)

    def forward(self, query, key=None, value=None, mask=None):
        # (I)
        query = query.refine_names(..., 'T', 'D')
        self_attn = key is None and value is None
        if self_attn:
            mask = mask.refine_names(..., 'T')
        else:
            mask = mask.refine_names(..., 'T', 'T_key')  # enc attn

        dim = query.size('D')
        assert dim == self.dim, \
            f'Dimensions do not match: {dim} query vs {self.dim} configured'
        assert mask is not None, 'Mask is None, please specify a mask'
        n_heads = self.n_heads
        dim_per_head = dim // n_heads
        scale = math.sqrt(dim_per_head)

        # (II)
        def prepare_head(tensor):
            tensor = tensor.refine_names(..., 'T', 'D')
            return (tensor.unflatten('D', [('H', n_heads), ('D_head', dim_per_head)])
                          .align_to(..., 'H', 'T', 'D_head'))

        assert value is None
        if self_attn:
            key = value = query
        elif value is None:
            # key and value are the same, but query differs
            key = key.refine_names(..., 'T', 'D')
            value = key
        dim = key.size('D')

        # Distinguish between query_len (T) and key_len (T_key) dims.
        k = prepare_head(self.k_lin(key)).rename(T='T_key')
        v = prepare_head(self.v_lin(value)).rename(T='T_key')
        q = prepare_head(self.q_lin(query))

        dot_prod = q.div_(scale).matmul(k.align_to(..., 'D_head', 'T_key'))
        dot_prod.refine_names(..., 'H', 'T', 'T_key')  # just a check

        # (III)
        attn_mask = (mask == 0).align_as(dot_prod)
        dot_prod.masked_fill_(attn_mask, -float(1e20))

        attn_weights = self.attn_dropout(F.softmax(dot_prod / scale,
                                                   dim='T_key'))

        # (IV)
        attentioned = (
            attn_weights.matmul(v).refine_names(..., 'H', 'T', 'D_head')
            .align_to(..., 'T', 'H', 'D_head')
            .flatten(['H', 'D_head'], 'D')
        )

        return self.out_lin(attentioned).refine_names(..., 'T', 'D')

(I) 细化输入张量维度

def forward(self, query, key=None, value=None, mask=None):
    # (I)
    query = query.refine_names(..., 'T', 'D')

query = query.refine_names(..., 'T', 'D') 充当可强制执行的文档，并将输入维度提升为命名维度。它检查最后两个维度是否可以细化为 ['T', 'D']，从而防止可能在后续代码中出现静默或令人困惑的大小不匹配错误。

(II) 在 prepare_head 中操作维度

# (II)
def prepare_head(tensor):
    tensor = tensor.refine_names(..., 'T', 'D')
    return (tensor.unflatten('D', [('H', n_heads), ('D_head', dim_per_head)])
                  .align_to(..., 'H', 'T', 'D_head'))

首先要注意的是代码如何清楚地说明了输入和输出维度：输入张量必须以 T 和 D 维度结尾，输出张量以 H、T 和 D_head 维度结尾。

第二要注意的是代码如何清楚地描述了正在执行的操作。prepare_head 接收键、查询和值，并将嵌入维度拆分为多个头，最后重新排列维度顺序，使其成为 [..., 'H', 'T', 'D_head']。ParlAI 使用以下代码实现 prepare_head，使用 view 和 transpose 操作

def prepare_head(tensor):
    # input is [batch_size, seq_len, n_heads * dim_per_head]
    # output is [batch_size * n_heads, seq_len, dim_per_head]
    batch_size, seq_len, _ = tensor.size()
    tensor = tensor.view(batch_size, tensor.size(1), n_heads, dim_per_head)
    tensor = (
        tensor.transpose(1, 2)
        .contiguous()
        .view(batch_size * n_heads, seq_len, dim_per_head)
    )
    return tensor

我们的命名张量变体使用了一些操作，虽然更冗长，但比 view 和 transpose 更有语义意义，并且包含以名称形式的可强制执行的文档。

(III) 按名称显式广播

def ignore():
    # (III)
    attn_mask = (mask == 0).align_as(dot_prod)
    dot_prod.masked_fill_(attn_mask, -float(1e20))

mask 通常具有 [N, T] 维度（在自注意力的情况下）或 [N, T, T_key] 维度（在编码器注意力的情况下），而 dot_prod 具有 [N, H, T, T_key] 维度。为了使 mask 与 dot_prod 正确广播，我们通常会在自注意力的情况下对维度 1 和 -1 进行 unsqueeze，或在编码器注意力的情况下对维度 1 进行 unsqueeze。使用命名张量，我们只需使用 align_as 将 attn_mask 对齐到 dot_prod，不再需要担心在哪里 unsqueeze 维度。

(IV) 使用 align_to 和 flatten 进行更多维度操作

def ignore():
    # (IV)
    attentioned = (
        attn_weights.matmul(v).refine_names(..., 'H', 'T', 'D_head')
        .align_to(..., 'T', 'H', 'D_head')
        .flatten(['H', 'D_head'], 'D')
    )

这里，与 (II) 中一样，align_to 和 flatten 比 view 和 transpose 更具语义意义（虽然更冗长）。

运行示例

n, t, d, h = 7, 5, 2 * 3, 3
query = torch.randn(n, t, d, names=('N', 'T', 'D'))
mask = torch.ones(n, t, names=('N', 'T'))
attn = MultiHeadAttention(h, d)
output = attn(query, mask=mask)
# works as expected!
print(output.names)

输出

('N', 'T', 'D')

上面的代码按预期工作。此外，请注意，我们在代码中根本没有提到批次维度的名称。实际上，我们的 MultiHeadAttention 模块与批次维度的存在无关。

query = torch.randn(t, d, names=('T', 'D'))
mask = torch.ones(t, names=('T',))
output = attn(query, mask=mask)
print(output.names)

输出

('T', 'D')

结论

感谢阅读！命名张量仍在不断开发中；如果你有任何反馈或改进建议，请通过创建 一个问题 告知我们。