PyTorch 中的命名張量簡(jiǎn)介(實(shí)驗(yàn)性）

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

作者： Richard Zou

命名張量旨在通過(guò)允許用戶將顯式名稱與張量維度相關(guān)聯(lián)來(lái)使張量更易于使用。 在大多數(shù)情況下，采用尺寸參數(shù)的操作將接受尺寸名稱，從而無(wú)需按位置跟蹤尺寸。 此外，命名張量使用名稱來(lái)自動(dòng)檢查運(yùn)行時(shí)是否正確使用了 API，從而提供了額外的安全性。 名稱也可以用于重新排列尺寸，例如，支持“按名稱廣播”而不是“按位置廣播”。

本教程旨在作為 1.3 啟動(dòng)中將包含的功能的指南。 到最后，您將能夠：

• 創(chuàng)建具有命名尺寸的張量，以及刪除或重命名這些尺寸
• 了解操作如何傳播維度名稱的基礎(chǔ)

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

1.         廣播業(yè)務(wù)
2.         展平和展平尺寸

最后，我們將通過(guò)使用命名張量編寫多頭注意力模塊來(lái)將其付諸實(shí)踐。

PyTorch 中的命名張量受 Sasha Rush 的啟發(fā)并與之合作。 Sasha 在他的 2019 年 1 月博客文章中提出了最初的想法和概念證明。

基礎(chǔ)知識(shí)：命名尺寸

PyTorch 現(xiàn)在允許張量具有命名尺寸； 工廠函數(shù)采用新的<cite>名稱</cite>參數(shù)，該參數(shù)將名稱與每個(gè)維度相關(guān)聯(lián)。 這適用于大多數(shù)工廠功能，例如

• <cite>張量</cite>
• <cite>為空</cite>
• <cite>個(gè)</cite>
• <cite>零</cite>
• <cite>蘭恩</cite>
• <cite>蘭特</cite>

在這里，我們使用名稱構(gòu)造張量：

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個(gè)維度的名稱。

重命名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')

刪除名稱的首選方法是調(diào)用tensor.rename(None)：

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

得出:

(None, None, None, None)

未命名的張量(沒(méi)有命名尺寸的張量）仍然可以正常工作，并且在其repr中沒(méi)有名稱。

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

得出:

tensor([[[-0.1654,  0.5440, -0.1883]],


        [[ 1.8791, -0.1997,  2.4531]]])
(None, None, None)

命名張量不需要命名所有尺寸。

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

得出:

('N', None, None, None)

因?yàn)槊麖埩靠梢耘c未命名張量共存，所以我們需要一種不錯(cuò)的方式來(lái)編寫可識(shí)別命名張量的代碼，該代碼可用于命名張量和未命名張量。 使用tensor.refine_names(*names)優(yōu)化尺寸并將未命名的暗淡提升為已命名的暗淡。 細(xì)化維度定義為“重命名”，并具有以下限制：

• 可以將None暗號(hào)細(xì)化為任何名稱
• 命名的 dim 只能精簡(jiǎn)為具有相同的名稱。

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

大多數(shù)簡(jiǎn)單的操作都會(huì)傳播名稱。 命名張量的最終目標(biāo)是所有操作以合理，直觀的方式傳播名稱。 在 1.3 發(fā)行版中已添加了對(duì)許多常用操作的支持。 例如，這里是.abs()：

print(named_imgs.abs().names)

得出:

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

存取器和減少

可以使用尺寸名稱來(lái)引用尺寸，而不是位置尺寸。 這些操作還傳播名稱。 索引(基本索引和高級(jí)索引）尚未實(shí)施，但仍在規(guī)劃中。 使用上方的named_imgs張量，我們可以執(zhí)行以下操作：

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

名稱推斷

名稱在稱為名稱推斷的兩步過(guò)程中在操作上傳播：

1. 檢查名稱：操作員可以在運(yùn)行時(shí)執(zhí)行自動(dòng)檢查，以檢查某些尺寸名稱是否必須匹配。
2. 傳播名稱：名稱推斷將輸出名稱傳播到輸出張量。

讓我們看一個(gè)非常小的例子，添加 2 個(gè)一維張量，不進(jìn)行廣播。

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

檢查名稱：首先，我們將檢查這兩個(gè)張量的名稱是否與相匹配。 當(dāng)且僅當(dāng)兩個(gè)名稱相等(字符串相等）或至少一個(gè)為None(None本質(zhì)上是一個(gè)特殊的通配符名稱）時(shí)，兩個(gè)名稱才匹配。 因此，這三者中唯一會(huì)出錯(cuò)的是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.

傳播名稱：通過(guò)返回兩個(gè)名稱中最精確的名稱來(lái)統(tǒng)一這兩個(gè)名稱。 使用x + y，X比None更精細(xì)。

print((x + y).names)

得出:

('X',)

大多數(shù)名稱推斷規(guī)則都很簡(jiǎn)單明了，但其中一些可能具有意料之外的語(yǔ)義。 讓我們來(lái)看看您可能會(huì)遇到的一對(duì)：廣播和矩陣乘法。

廣播

命名張量不會(huì)改變廣播行為； 他們?nèi)匀话次恢脧V播。 但是，在檢查兩個(gè)尺寸是否可以廣播時(shí)，PyTorch 還會(huì)檢查這些尺寸的名稱是否匹配。

這導(dǎo)致命名張量防止廣播操作期間意外對(duì)齊。 在下面的示例中，我們將per_batch_scale應(yīng)用于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.

如果沒(méi)有names，則per_batch_scale張量與imgs的最后一個(gè)尺寸對(duì)齊，這不是我們想要的。 我們確實(shí)想通過(guò)將per_batch_scale與imgs的批次尺寸對(duì)齊來(lái)執(zhí)行操作。 有關(guān)如何按名稱對(duì)齊張量的信息，請(qǐng)參見(jiàn)新的“按名稱顯式廣播”功能，如下所述。

矩陣乘法

torch.mm(A, B)在A的第二個(gè)暗角和B的第一個(gè)暗角之間執(zhí)行點(diǎn)積，返回具有A的第一個(gè)暗角和B的第二個(gè)暗角的張量。 (其他 matmul 函數(shù)，例如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')

如您所見(jiàn)，矩陣乘法不會(huì)檢查收縮尺寸是否具有相同的名稱。

接下來(lái)，我們將介紹命名張量啟用的兩個(gè)新行為：按名稱的顯式廣播以及按名稱的展平和展平尺寸

新行為：按名稱明確廣播

有關(guān)使用多個(gè)維度的主要抱怨之一是需要unsqueeze“虛擬”維度，以便可以進(jìn)行操作。 例如，在之前的每批比例示例中，使用未命名的張量，我們將執(zhí)行以下操作：

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)

通過(guò)使用名稱，我們可以使這些操作更安全(并且易于與維數(shù)無(wú)關(guān)）。 我們提供了一個(gè)新的tensor.align_as(other)操作，可以對(duì)張量的尺寸進(jìn)行排列以匹配other.names中指定的順序，并在適當(dāng)?shù)牡胤教砑右粋€(gè)尺寸的尺寸(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)

新行為：按名稱展平和展平尺寸

一種常見(jiàn)的操作是展平和展平尺寸。 現(xiàn)在，用戶可以使用view，reshape或flatten來(lái)執(zhí)行此操作； 用例包括展平批處理尺寸以將張量發(fā)送到必須采用一定數(shù)量尺寸的輸入的運(yùn)算符(即 conv2d 采用 4D 輸入）。

為了使這些操作比查看或整形更具語(yǔ)義意義，我們引入了一種新的tensor.unflatten(dim, namedshape)方法并更新flatten以使用名稱：tensor.flatten(dims, new_dim)。

flatten只能展平相鄰的尺寸，但也可以用于不連續(xù)的暗光。 必須將名稱為的傳遞到unflatten中，該形狀是(dim, size)元組的列表，以指定如何展開(kāi)暗淡。 可以在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 當(dāng)前會(huì)忽略所有張量上的名稱，只是將它們視為常規(guī)張量。 梯度計(jì)算是正確的，但是我們失去了名稱賦予我們的安全性。 在路線圖上引入名稱以自動(dòng)分級(jí)的處理。

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([-1.4406,  0.6030,  0.1825])
tensor([-1.4406,  0.6030,  0.1825])

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

有關(guān) 1.3 發(fā)行版支持的功能的詳細(xì)分類，請(qǐng)參見(jiàn)此處。

特別是，我們要指出當(dāng)前不支持的三個(gè)重要功能：

• 通過(guò)torch.save或torch.load保存或加載命名張量
• 通過(guò)torch.multiprocessing進(jìn)行多重處理
• JIT 支持； 例如，以下將錯(cuò)誤

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)`.

解決方法是，在使用尚不支持命名張量的任何東西之前，請(qǐng)通過(guò)tensor = tensor.rename(None)刪除名稱。

更長(zhǎng)的例子：多頭注意力

現(xiàn)在，我們將通過(guò)一個(gè)完整的示例來(lái)實(shí)現(xiàn)一個(gè)常見(jiàn)的 PyTorch nn.Module：多頭注意。 我們假設(shè)讀者已經(jīng)熟悉多頭注意； 要進(jìn)行復(fù)習(xí)，請(qǐng)查看此說(shuō)明或此說(shuō)明。

我們采用 ParlAI 來(lái)實(shí)現(xiàn)多頭注意力的實(shí)現(xiàn)； 特別是此處。 閱讀該示例中的代碼； 然后，與下面的代碼進(jìn)行比較，注意有四個(gè)標(biāo)記為(I），(II），(III）和(IV）的位置，使用命名張量可以使代碼更易讀； 在代碼塊之后，我們將深入探討其中的每一個(gè)。

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）細(xì)化輸入張量調(diào)光

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

query = query.refine_names(..., 'T', 'D')用作可執(zhí)行的文檔，并將輸入尺寸提升為名稱。 它檢查是否可以將最后兩個(gè)維度調(diào)整為['T', 'D']，以防止在以后出現(xiàn)潛在的靜默或混亂的尺寸不匹配錯(cuò)誤。

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

首先要注意的是代碼如何清楚地說(shuō)明輸入和輸出尺寸：輸入張量必須以T和D變暗結(jié)束，輸出張量應(yīng)以H，T和D_head結(jié)束 昏暗。

要注意的第二件事是代碼清楚地描述了正在發(fā)生的事情。 prepare_head 獲取鍵，查詢和值，并將嵌入的 dim 拆分為多個(gè) head，最后將 dim 順序重新排列為[..., 'H', 'T', 'D_head']。 ParlAI 使用view和transpose操作實(shí)現(xiàn)以下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

我們命名的張量變量使用的操作雖然較為冗長(zhǎng)，但比view和transpose具有更多的語(yǔ)義含義，并包含名稱形式的可執(zhí)行文檔。

(III）按名稱明確廣播

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

mask通常具有暗淡[N, T](在自我關(guān)注的情況下）或[N, T, T_key](對(duì)于編碼器注意的情況），而dot_prod具有暗淡[N, H, T, T_key]。 為了使mask與dot_prod正確廣播，我們通常會(huì)在自注意的情況下將<cite>的</cite>調(diào)暗1和-1壓下，在編碼器的情況下，將unsqueeze調(diào)暗unsqueeze調(diào)暗 。 使用命名張量，我們只需使用align_as將attn_mask與dot_prod對(duì)齊，而不必?fù)?dān)心unsqueeze變暗的位置。

(IV）使用 align_to 和展平進(jìn)行更多尺寸操作

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在語(yǔ)義上比view和transpose更有意義(盡管更冗長(zhǎng)）。

運(yùn)行示例

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

以上工作正常。 此外，請(qǐng)注意，在代碼中我們根本沒(méi)有提到批處理維度的名稱。 實(shí)際上，我們的MultiHeadAttention模塊與批次尺寸的存在無(wú)關(guān)。

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

得出:

('T', 'D')

結(jié)論

感謝您的閱讀！ 命名張量仍在發(fā)展中。 如果您有反饋和/或改進(jìn)建議，請(qǐng)通過(guò)創(chuàng)建問(wèn)題來(lái)通知我們。

腳本的總運(yùn)行時(shí)間：(0 分鐘 0.074 秒）

Download Python source code: named_tensor_tutorial.py Download Jupyter notebook: named_tensor_tutorial.ipynb

由獅身人面像畫廊生成的畫廊