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

作者: Richard Zou

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

本教程旨在作为1.3版本发布功能的指南。通过本教程，您将能够：

• Create Tensors with named dimensions, as well as remove or rename those dimensions
• Understand the basics of how operations propagate dimension names

• See how naming dimensions enables clearer code in two key areas:

  • Broadcasting operations
  • Flattening and unflattening dimensions

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

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)来细化维度并将未命名维度提升为命名维度。细化维度被定义为具有以下约束的“重命名”：

• A None dim can be refined to have any name
• A named dim can only be refined to have the same name.

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. Check names: an operator may perform automatic checks at runtime that check that certain dimension names must match.
2. Propagate names: name inference propagates output names to output tensors.

让我们通过一个非常小的例子来添加两个没有广播的一维张量。

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输入）。

为了使这些操作在语义上比视图或重塑更有意义，我们引入了一个新的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])

其他支持（和不支持）的功能

See here 查看1.3版本支持的详细内容。

特别是，我们想指出三个目前不支持的重要功能：

• Saving or loading named tensors via torch.save or torch.load
• Multi-processing via torch.multiprocessing
• JIT support; for example, the following will error

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 接受 key、query 和 value，并将嵌入维度分割成 多个头，最后重新排列维度顺序为 [..., 'H', 'T', 'D_head']。 ParlAI 使用 view 和 transpose 操作 实现了 prepare_head，如下所示：

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，我们通常会在自注意力机制的情况下 unsqueeze 维度 1 和 -1，或者在编码器 注意力机制的情况下 unsqueeze 维度 1。使用命名张量时，我们只需使用 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')

结论

感谢您的阅读！命名张量仍在开发中；如果您有反馈和/或改进建议，请通过创建问题告知我们。