PyTorch-20 命名張量(Named Tensors)的介紹

請到這里去查看圖文教程:http://studyai.com/pytorch-1.4/intermediate/named_tensor_tutorial.html

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

本教程旨在作為1.3發(fā)布中包含的功能的指南。最后,您將能夠:

創(chuàng)建具有命名維度的張量,并移除或重命名這些維度。
了解操作/算子(operations)如何傳播維度名稱的基礎(chǔ)知識

請參見命名維度如何在兩個關(guān)鍵領(lǐng)域?qū)崿F(xiàn)更清晰的代碼:
        廣播操作(Broadcasting operations)
        展平和收縮維度(Flattening and unflattening dimensions)

最后,我們將通過使用named-tensors構(gòu)建多頭注意模塊(multi-head attention module)來學(xué)習(xí)實踐命名張量的這些知識點。

PyTorch中的命名張量是由Sasha Rush啟發(fā)并與 Sasha Rush 合作完成的。 Sasha在他的博客 January 2019 blog post 中提出了最初的想法和概念證明。
基礎(chǔ): 命名維度(named dimensions)

PyTorch現(xiàn)在允許張量具有命名維度;工廠函數(shù)(factory functions)采用一個新的“names”參數(shù),該參數(shù)將把每個維度與一個名稱相關(guān)聯(lián)。 這一方法在很多工廠函數(shù)中都可以使用:

tensor
empty
ones
zeros
randn
rand

現(xiàn)在我們構(gòu)造一個伴有名稱的張量:

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: 設(shè)置 .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)

刪除名稱的首選方法是調(diào)用 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.(已經(jīng)有名稱的維度 的 名稱 保持不變)

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

嘗試將一個 已經(jīng)存在名稱的維度 的 名稱 修改成 別的名稱

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

大多數(shù)簡單的操作可以傳播名稱. 命名張量的終極目標(biāo)是所有操作都可以以合理的、直觀的方式傳播名稱(propagate names). 在1.3版本中增加了對許多常見操作的支持, 比如: .abs() :

print(named_imgs.abs().names)

訪問與壓縮(Accessors 和 Reduction)

你可以使用維度名稱而不是位置來引用維度。 這些操作也支持名稱傳播. 索引當(dāng)前還未實現(xiàn)但已經(jīng)在規(guī)劃實現(xiàn)了。 使用 named_imgs 張量, 我們可以做下列操作:

output = named_imgs.sum('C') # 沿著channel維度執(zhí)行求和操作
print(output.names)

img0 = named_imgs.select('N', 0) # 獲取一張圖像
print(img0.names)

名稱推理(name inference)

名稱在兩步操作之間傳播的過程稱之為: name inference:

檢查名稱: 操作算子可以在運行時執(zhí)行自動檢查,以檢查某些維度名稱是否必須匹配。
傳播名稱: 名稱推理傳播輸出名稱到輸出張量。

我們先來體驗兩個很小的例子:adding 2 one-dim tensors with no broadcasting.

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

檢查名稱: 首先, 我們將檢查這兩個張量的名稱是否匹配,兩個名稱要匹配只要名稱對應(yīng)的字符串 相等即可,或者至少有一個是“None”(這里的 “None”可以理解為通配符式的名稱)。 按照這一規(guī)則,上面的三個量的相互加法中, 只有一個會失敗,即 x + z:

catch_error(lambda: x + z)

傳播名稱: 通過返回兩個名稱中最精煉的那個名稱來 統(tǒng)一(unify) 兩個名稱。 在 x + y 中, X 比 None 更精煉(refine).

print((x + y).names)

大多數(shù)名稱推斷規(guī)則都很簡單,但其中一些規(guī)則可能具有意外的語義(unexpected semantics)。 讓我們看看你可能會遇到的這些場景: 廣播和矩陣乘法 .
廣播(Broadcasting)

命名張量不會改變廣播行為本身:任然按照位置進行廣播. 然而, 當(dāng)檢查兩個維度是否可以被廣播的時候, PyTorch也會同時檢查這些維度的名稱是否匹配。

這將導(dǎo)致命名張量在廣播操作期間防止意外對齊(preventing unintended alignment)。 在下面的例子中,我們將 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)

如果沒有名稱(names), 張量 per_batch_scale 會被對齊到 imgs 的最后一維,這不是我們希望的。 我們實際上想要執(zhí)行的操作是把 per_batch_scale 和 imgs 的 batch 維對齊。 請查看 “通過名稱顯式廣播” 功能 來實現(xiàn)通過名稱對齊張量操作維度。
矩陣相乘

torch.mm(A, B) 在 A 的第二維和B的第一維執(zhí)行點積(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'))

實行一次狀態(tài)轉(zhuǎn)移過程

new_state = markov_states @ transition_matrix
print(new_state.names)

如您所見,矩陣乘法不檢查縮減維度(contracted dimensions)是否具有相同的名稱。

接下來,我們將介紹兩個由命名張量賦予的新行為:通過名稱進行顯式廣播 和 通過名稱展平/收縮維度
新行為一: 通過名稱進行顯式廣播(Explicit broadcasting by names)

使用多維度的一個主要抱怨是需要對 “偽(dummy)” 維度進行 “unsqueze” ,以便某些操作可以成功執(zhí)行。 比如, 在我們上面的 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)

通過使用命名張量,我們可以使這些操作更安全(而且在不確定維度的數(shù)量時也很容易執(zhí)行操作)。
我們提供一個新的 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. 目前,用戶執(zhí)行這一過程使用的是 view, reshape , 或 flatten ; 常見用法包括:將批處理維度展平以將張量發(fā)送到必須接受具有特定維度數(shù)的輸入的運算符中 (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) 目前會忽略所有張量上的名稱并將其視為常規(guī)張量進行計算。 雖然梯度的計算仍然是正確的但是卻損失了張量命名帶來的安全性。 對自動微分的支持也在開發(fā)路線圖中。

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) # 現(xiàn)在還是未命名的. 未來的版本會實現(xiàn)這一點

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

理想情況下,我們會檢查loss和grad_loss的名稱是否匹配,但我們還沒有實現(xiàn)這一點

loss.backward(grad_loss)

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

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

有關(guān)1.3版本支持的內(nèi)容的詳細(xì)分解, 請看這兒: 。

特別是,我們要指出目前不支持的三個重要功能:

通過 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))

作為權(quán)宜之計, 在使用任何尚不支持命名張量的名稱之前,請通過tensor=tensor.rename(None)刪除名稱。
一個比較長的案例: Multi-head attention

現(xiàn)在,我們將通過一個完整的示例來實現(xiàn)一個 PyTorch 的 “nn.Module”: multi-head attention. 我們假設(shè)讀者已經(jīng)熟悉: multi-head attention; 如果你是新手, 請看 這個解釋 或 這個解釋.

我們采用了來自 ParlAI 的實現(xiàn):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) 改善細(xì)化(refine)輸入張量的維度

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

query=query.refine_names(…,'T','D') 用作可強制執(zhí)行的文檔[serves as enforcable documentation], 并將輸入的未命名維度提升為命名維度。它檢查最后兩個維度是否可以細(xì)化[refine]為[‘T’,’D’], 以防止以后可能出現(xiàn)的無提示或混淆大小不匹配錯誤。

**(II) 操控 prepare_head 函數(shù)中張量的維度 **

(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 維度結(jié)束, 輸出張量以 H 維度、T 維度和 D_head 維度結(jié)束。

第二件要注意的事情是代碼如何清楚地描述了正在發(fā)生的事情。 prepare_head 獲取key、query和value, 并將嵌入的dim拆分為多個head,最后將dim的順序重新排列為 […,'H','T','D_head'] 。 ParlAI 實現(xiàn)的 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

我們的命名張量所實現(xiàn)的 prepare_head 函數(shù)變體使用的操作雖然更詳細(xì),但比 view 和 transpose 實現(xiàn)的 prepare_head 版本具有更多的語義意義,并且包含以名稱形式存在的可執(zhí)行文檔[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)

最后編輯于
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成,瀏覽時請結(jié)合常識與多方信息審慎甄別。
平臺聲明:文章內(nèi)容(如有圖片或視頻亦包括在內(nèi))由作者上傳并發(fā)布,文章內(nèi)容僅代表作者本人觀點,簡書系信息發(fā)布平臺,僅提供信息存儲服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容