PyTorch 中的命名张量简介（原型）

原文：https://pytorch.org/tutorials/intermediate/named_tensor_tutorial.html

作者： Richard Zou

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

本教程旨在作为 1.3 启动中将包含的功能的指南。 到最后，您将能够：

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

• 了解命名尺寸如何在两个关键区域实现更清晰的代码：

  • 广播操作
  • 重塑和展开尺寸

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

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

基础知识：命名维度

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

• 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.7420, -0.3646,  0.1424]],
        [[-0.6065, -1.4888,  0.2935]]])
(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. 传播名称：名称推断将输出名称传播到输出张量。

让我们看一个非常小的例子，添加 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的第二个暗角的张量。 （其他matmul函数，例如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.5398, 0.7907, 0.7784])
tensor([0.5398, 0.7907, 0.7784])

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

有关 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 使用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正确广播，我们通常会在自注意的情况下将的调暗1和-1压下，在编码器的情况下，我们将unsqueeze调暗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')

总结

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

脚本的总运行时间：（0 分钟 0.094 秒）

下载 Python 源码：named_tensor_tutorial.py

下载 Jupyter 笔记本：named_tensor_tutorial.ipynb

由 Sphinx 画廊生成的画廊