（原型）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. **傳播名稱**：名稱推斷將輸出名稱傳播到輸出張量。

讓我們透過一個非常小的範例來瞭解，在沒有廣播的情況下添加 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')

自動微分支援

自動微分目前會忽略所有張量上的名稱，只將它們視為常規張量。梯度計算是正確的，但我們失去了名稱給予我們的安全性。將名稱處理引入自動微分的路線圖正在制定中。

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 使用 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 操作，或者在編碼器注意力的情況下對維度 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')

結論

感謝您的閱讀！命名張量仍在開發中；如果您有任何意見反饋和/或改進建議，請通過建立問題告知我們。