PyTorch-20 命名张量(Named Tensors)的介绍

Posted 2023-02-23 ScorpioDoctor

请到这里去查看图文教程：http://studyai.com/pytorch-1.4/intermediate/named_tensor_tutorial.html

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

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

创建具有命名维度的张量，并移除或重命名这些维度。
了解操作/算子(operations)如何传播维度名称的基础知识

请参见命名维度如何在两个关键领域实现更清晰的代码：
        广播操作(Broadcasting operations)
        展平和收缩维度(Flattening and unflattening dimensions)

最后，我们将通过使用named-tensors构建多头注意模块(multi-head attention module)来学习实践命名张量的这些知识点。

PyTorch中的命名张量是由Sasha Rush启发并与 Sasha Rush 合作完成的。 Sasha在他的博客 January 2019 blog post 中提出了最初的想法和概念证明。 基础: 命名维度(named dimensions)

PyTorch现在允许张量具有命名维度；工厂函数(factory functions)采用一个新的“names”参数，该参数将把每个维度与一个名称相关联。 这一方法在很多工厂函数中都可以使用：

tensor
empty
ones
zeros
randn
rand

现在我们构造一个伴有名称的张量:

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

命名维度是这样排序的: tensor.names[i] is the name of the i th dimension of tensor.

有两种方法去重新命名 Tensor 的维度(dimensions):

方法 #1: 设置 .names 属性(attribute) (这种方法可以原位修改指定维度的名称)

imgs.names = ['batch', 'channel', 'width', 'height'] print(imgs.names)

方法 #2: 指定新的names (this changes names out-of-place)

imgs = imgs.rename(channel='C', width='W', height='H') print(imgs.names)

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

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

未命名张量 (tensors with no named dimensions) 仍然可以正常工作， 并且在他们的 “repr” 中也没有名称

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

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

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

因为命名张量可以与未命名张量共存，所以我们需要一种很好的方法来编写命名的张量感知代码(named tensor-aware code)， 它可以同时感知命名的和未命名的张量。 使用 tensor.refine_names(*names) 来改善维度命名情况，把未命名的维度都变成命名维度(lift unnamed dims to named dims)。 改善一个维度(Refining a dimension)是指带有下列约束条件的一个重命名(“rename”)操作：

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')

尝试将一个 已经存在名称的维度 的 名称 修改成 别的名称

catch_error(lambda: named_imgs.refine_names('N', 'C', 'H', 'width'))

大多数简单的操作可以传播名称. 命名张量的终极目标是所有操作都可以以合理的、直观的方式传播名称(propagate names). 在1.3版本中增加了对许多常见操作的支持, 比如： .abs() :

print(named_imgs.abs().names)

访问与压缩(Accessors 和 Reduction)

你可以使用维度名称而不是位置来引用维度。 这些操作也支持名称传播. 索引当前还未实现但已经在规划实现了。 使用 named_imgs 张量, 我们可以做下列操作:

output = named_imgs.sum('C') # 沿着channel维度执行求和操作 print(output.names)

img0 = named_imgs.select('N', 0) # 获取一张图像 print(img0.names)

名称推理(name inference)

名称在两步操作之间传播的过程称之为: name inference:

检查名称: 操作算子可以在运行时执行自动检查，以检查某些维度名称是否必须匹配。
传播名称: 名称推理传播输出名称到输出张量。

我们先来体验两个很小的例子：adding 2 one-dim tensors with no broadcasting.

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)

传播名称: 通过返回两个名称中最精炼的那个名称来 统一(unify) 两个名称。 在 x + y 中, X 比 None 更精炼(refine).

print((x + y).names)

大多数名称推断规则都很简单，但其中一些规则可能具有意外的语义(unexpected semantics)。 让我们看看你可能会遇到的这些场景: 广播和矩阵乘法 . 广播(Broadcasting)

命名张量不会改变广播行为本身：任然按照位置进行广播. 然而, 当检查两个维度是否可以被广播的时候， PyTorch也会同时检查这些维度的名称是否匹配。

这将导致命名张量在广播操作期间防止意外对齐(preventing unintended alignment)。 在下面的例子中，我们将 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)

如果没有名称(names), 张量 per_batch_scale 会被对齐到 imgs 的最后一维，这不是我们希望的。 我们实际上想要执行的操作是把 per_batch_scale 和 imgs 的 batch 维对齐。 请查看 “通过名称显式广播” 功能 来实现通过名称对齐张量操作维度。 矩阵相乘

torch.mm(A, B) 在 A 的第二维和B的第一维执行点积(product)操作， 返回张量的第一维和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'))

实行一次状态转移过程

new_state = markov_states @ transition_matrix print(new_state.names)

如您所见，矩阵乘法不检查缩减维度(contracted dimensions)是否具有相同的名称。

接下来，我们将介绍两个由命名张量赋予的新行为：通过名称进行显式广播 和 通过名称展平/收缩维度 新行为一: 通过名称进行显式广播(Explicit broadcasting by names)

使用多维度的一个主要抱怨是需要对 “伪(dummy)” 维度进行 “unsqueze” ，以便某些操作可以成功执行。 比如, 在我们上面的 per-batch-scale 案例中, 在张量不命名的情况下，我们需要这样做：

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 中的特定顺序, adding one-sized dimensions where appropriate (tensor.align_to(names) works as well):

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)

注意: named tensors do not yet work with allclose

assert torch.allclose(named_result.rename(None), correct_result)

新行为二: 通过名称展平/收缩维度

一个常见操作是展平/收缩维度： flattening and unflattening dimensions. 目前，用户执行这一过程使用的是 view, reshape , 或 flatten ; 常见用法包括：将批处理维度展平以将张量发送到必须接受具有特定维度数的输入的运算符中 (i.e., conv2d 接受 4D 输入).

为了使这些操作比 view, reshape更有语义意义，我们 介绍一个新的方法tensor.unflatten(dim, namedshape) 方法 并更新 flatten 使其可以在命名张量中工作: tensor.flatten(dims, new_dim).

flatten can only flatten adjacent dimensions but also works on non-contiguous dims. One must pass into unflatten a named shape, which is a list of (dim, size) tuples, to specify how to unflatten the dim. It is possible to save the sizes during a flatten for unflatten but we do not yet do that.

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

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

自动微分的支持

自动微分(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) # 现在还是未命名的. 未来的版本会实现这一点

weight.grad.zero_() grad_loss = grad_loss.refine_names('C') loss = (x - weight).abs()

理想情况下，我们会检查loss和grad_loss的名称是否匹配，但我们还没有实现这一点

loss.backward(grad_loss)

print(weight.grad) # 仍然是未命名的 assert torch.allclose(weight.grad, correct_grad)

其他一些已支持的和未支持的特色

有关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))

作为权宜之计， 在使用任何尚不支持命名张量的名称之前，请通过tensor=tensor.rename（None）删除名称。 一个比较长的案例: Multi-head attention

现在，我们将通过一个完整的示例来实现一个 PyTorch 的 “nn.Module”： multi-head attention. 我们假设读者已经熟悉： multi-head attention; 如果你是新手, 请看 这个解释 或 这个解释.

我们采用了来自 ParlAI 的实现：multi-head attention; 尤其是，这里的代码. 阅读该示例中的代码；然后，与下面的代码进行比较， 请注意，代码中有四个地方加了注释 (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) 改善细化(refine)输入张量的维度

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

query=query.refine_names（…，'T'，'D'） 用作可强制执行的文档[serves as enforcable documentation]， 并将输入的未命名维度提升为命名维度。它检查最后两个维度是否可以细化[refine]为[‘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， 并将嵌入的dim拆分为多个head，最后将dim的顺序重新排列为 […，'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

我们的命名张量所实现的 prepare_head 函数变体使用的操作虽然更详细，但比 view 和 transpose 实现的 prepare_head 版本具有更多的语义意义，并且包含以名称形式存在的可执行文档[enforcable documentation]。

(III) 通过名称显式广播

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

mask 通常具有维度 [N, T] (在self attention中) 或者 [N, T, T_key] (在encoder attention中) 而 dot_prod 具有维度 [N, H, T, T_key]. To make mask broadcast correctly with dot_prod, we would usually unsqueeze dims 1 and -1 in the case of self attention or unsqueeze dim 1 in the case of encoder attention. Using named tensors, we simply align attn_mask to dot_prod using align_as and stop worrying about where to unsqueeze dims.

*(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)

以上工作如期望地那样进行。此外，请注意，在代码中我们根本没有提到批处理维度(batch dimension)的名称。 事实上，我们的MultiHeadAttention 模块是不知道 批处理维度(batch dimension)的 存在的。

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