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

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

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

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

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

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

最后蜈垮,我們將通過使用named-tensors構建多頭注意模塊(multi-head attention module)來學習實踐命名張量的這些知識點菊碟。

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

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

tensor
empty
ones
zeros
randn
rand

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

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)

刪除名稱的首選方法是調(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ù)簡單的操作可以傳播名稱. 命名張量的終極目標是所有操作都可以以合理的沛慢、直觀的方式傳播名稱(propagate names). 在1.3版本中增加了對許多常見操作的支持, 比如: .abs() :

print(named_imgs.abs().names)

訪問與壓縮(Accessors 和 Reduction)

你可以使用維度名稱而不是位置來引用維度。 這些操作也支持名稱傳播. 索引當前還未實現(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',))

檢查名稱: 首先, 我們將檢查這兩個張量的名稱是否匹配集乔,兩個名稱要匹配只要名稱對應的字符串 相等即可去件,或者至少有一個是“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)

命名張量不會改變廣播行為本身:任然按照位置進行廣播. 然而, 當檢查兩個維度是否可以被廣播的時候丈攒, 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 的最后一維,這不是我們希望的景殷。 我們實際上想要執(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)轉移過程

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)

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

有關1.3版本支持的內(nèi)容的詳細分解存和, 請看這兒: 奕剃。

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

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

現(xiàn)在,我們將通過一個完整的示例來實現(xiàn)一個 PyTorch 的 “nn.Module”: multi-head attention. 我們假設讀者已經(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) 改善細化(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]术吗, 并將輸入的未命名維度提升為命名維度尉辑。它檢查最后兩個維度是否可以細化[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 維度結束, 輸出張量以 H 維度隘蝎、T 維度和 D_head 維度結束购啄。

第二件要注意的事情是代碼如何清楚地描述了正在發(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ù)變體使用的操作雖然更詳細映胁,但比 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)

最后編輯于
?著作權歸作者所有,轉載或內(nèi)容合作請聯(lián)系作者
  • 序言:七十年代末蚪黑,一起剝皮案震驚了整個濱河市盅惜,隨后出現(xiàn)的幾起案子中剩,更是在濱河造成了極大的恐慌,老刑警劉巖抒寂,帶你破解...
    沈念sama閱讀 216,843評論 6 502
  • 序言:濱河連續(xù)發(fā)生了三起死亡事件结啼,死亡現(xiàn)場離奇詭異,居然都是意外死亡屈芜,警方通過查閱死者的電腦和手機郊愧,發(fā)現(xiàn)死者居然都...
    沈念sama閱讀 92,538評論 3 392
  • 文/潘曉璐 我一進店門,熙熙樓的掌柜王于貴愁眉苦臉地迎上來井佑,“玉大人属铁,你說我怎么就攤上這事」蹋” “怎么了焦蘑?”我有些...
    開封第一講書人閱讀 163,187評論 0 353
  • 文/不壞的土叔 我叫張陵,是天一觀的道長盒发。 經(jīng)常有香客問我,道長宁舰,這世上最難降的妖魔是什么拼卵? 我笑而不...
    開封第一講書人閱讀 58,264評論 1 292
  • 正文 為了忘掉前任,我火速辦了婚禮蛮艰,結果婚禮上间学,老公的妹妹穿的比我還像新娘。我一直安慰自己印荔,他們只是感情好低葫,可當我...
    茶點故事閱讀 67,289評論 6 390
  • 文/花漫 我一把揭開白布。 她就那樣靜靜地躺著仍律,像睡著了一般嘿悬。 火紅的嫁衣襯著肌膚如雪。 梳的紋絲不亂的頭發(fā)上水泉,一...
    開封第一講書人閱讀 51,231評論 1 299
  • 那天善涨,我揣著相機與錄音,去河邊找鬼草则。 笑死钢拧,一個胖子當著我的面吹牛,可吹牛的內(nèi)容都是我干的炕横。 我是一名探鬼主播源内,決...
    沈念sama閱讀 40,116評論 3 418
  • 文/蒼蘭香墨 我猛地睜開眼,長吁一口氣:“原來是場噩夢啊……” “哼份殿!你這毒婦竟也來了膜钓?” 一聲冷哼從身側響起嗽交,我...
    開封第一講書人閱讀 38,945評論 0 275
  • 序言:老撾萬榮一對情侶失蹤,失蹤者是張志新(化名)和其女友劉穎颂斜,沒想到半個月后夫壁,有當?shù)厝嗽跇淞掷锇l(fā)現(xiàn)了一具尸體,經(jīng)...
    沈念sama閱讀 45,367評論 1 313
  • 正文 獨居荒郊野嶺守林人離奇死亡沃疮,尸身上長有42處帶血的膿包…… 初始之章·張勛 以下內(nèi)容為張勛視角 年9月15日...
    茶點故事閱讀 37,581評論 2 333
  • 正文 我和宋清朗相戀三年盒让,在試婚紗的時候發(fā)現(xiàn)自己被綠了。 大學時的朋友給我發(fā)了我未婚夫和他白月光在一起吃飯的照片司蔬。...
    茶點故事閱讀 39,754評論 1 348
  • 序言:一個原本活蹦亂跳的男人離奇死亡邑茄,死狀恐怖,靈堂內(nèi)的尸體忽然破棺而出葱她,到底是詐尸還是另有隱情撩扒,我是刑警寧澤,帶...
    沈念sama閱讀 35,458評論 5 344
  • 正文 年R本政府宣布吨些,位于F島的核電站搓谆,受9級特大地震影響,放射性物質(zhì)發(fā)生泄漏豪墅。R本人自食惡果不足惜泉手,卻給世界環(huán)境...
    茶點故事閱讀 41,068評論 3 327
  • 文/蒙蒙 一、第九天 我趴在偏房一處隱蔽的房頂上張望偶器。 院中可真熱鬧斩萌,春花似錦、人聲如沸屏轰。這莊子的主人今日做“春日...
    開封第一講書人閱讀 31,692評論 0 22
  • 文/蒼蘭香墨 我抬頭看了看天上的太陽霎苗。三九已至姆吭,卻和暖如春,著一層夾襖步出監(jiān)牢的瞬間唁盏,已是汗流浹背内狸。 一陣腳步聲響...
    開封第一講書人閱讀 32,842評論 1 269
  • 我被黑心中介騙來泰國打工, 沒想到剛下飛機就差點兒被人妖公主榨干…… 1. 我叫王不留厘擂,地道東北人昆淡。 一個月前我還...
    沈念sama閱讀 47,797評論 2 369
  • 正文 我出身青樓,卻偏偏與公主長得像刽严,于是被迫代替她去往敵國和親昂灵。 傳聞我的和親對象是個殘疾皇子,可洞房花燭夜當晚...
    茶點故事閱讀 44,654評論 2 354

推薦閱讀更多精彩內(nèi)容