請到這里去查看圖文教程: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)