深度學習框架中的「張量」不好用？也許我們需要重新定義Tensor了
選自harvardnlp
作者：Alexander Rush
機器之心編譯
參與：李詩萌、路雪
本文介紹了張量的陷阱和一種可以閃避陷阱的替代方法named tensor，並進行了概念驗證。
儘管張量在深度學習的世界中無處不在，但它是有破綻的。它催生出了一些壞習慣，比如公開專用維度、基於絕對位置進行廣播，以及在文檔中保存類型信息。這篇文章介紹了一種具有命名維度的替代方法named tensor，並對其進行了概念驗證。這一改變消除了對索引、維度參數、einsum 式解壓縮以及基於文檔的編碼的需求。這篇文章附帶的原型PyTorch 庫可以作為namedtensor 使用。
PyTorch 庫參見：https://github.com/harvardnlp/NamedTensor
實現：
- Jon Malmaud 指出xarray 項目（http://xarray.pydata.org/en/stable/）的目標與namedtensor 非常相似，xarray 項目還增加了大量Pandas 和科學計算的支持。
- Tongfei Chen 的Nexus 項目在Scala 中提出了靜態類型安全的張量。
- Stephan Hoyer 和Eric Christiansen 為TensorFlow 建立了標註張量庫，Labed Tensor，和本文的方法是一樣的。
- Nishant Sinha 有TSA 庫，它使用類型註釋來定義維度名稱。
#@title Setup #!rm -fr NamedTensor/; git clone -q https://github.com/harvardnlp/NamedTensor.git #!cd NamedTensor; pip install -q .; pip install -q torch numpy opt_einsum
import numpy
import torch
from namedtensor import NamedTensor, ntorch
from namedtensor import _im_init
_im_init()
張量陷阱
這篇文章是關於張量類的。張量類是多維數組對象，是Torch、TensorFlow、Chainer 以及NumPy 等深度學習框架的核心對象。張量具備大量存儲空間，還可以向用戶公開維度信息。
ims = torch.tensor(numpy.load('test_images.npy'))
ims.shapetorch.Size([6, 96, 96, 3])
該示例中有4 個維度，對應的是batch_size、height、width 和channels。大多數情況下，你可以通過代碼註釋弄明白維度的信息，如下所示：
# batch_size x height x width x channels
ims[0]
這種方法簡明扼要，但從編程角度看來，這不是構建複雜軟件的好方法。
陷阱1：按慣例對待專用維度
代碼通過元組中的維度標識符操縱張量。如果要旋轉圖像，閱讀註釋，確定並更改需要改變的維度。
def rotate(ims):
# batch_size x height x width x channels
rotated = ims.transpose(1, 2)
# batch_size x width x height x channels
return rotated
rotate(ims)[0]
這段代碼很簡單，而且從理論上講記錄詳盡。但它並沒有反映目標函數的語義。旋轉的性質與batch 或channel 都無關。在確定要改變的維度時，函數不需要考慮這些維度。
這就產生了兩個問題。首先，令人非常擔心的是如果我們傳入單例圖像，函數可以正常運行但是卻不起作用。
rotate(ims[0]).shapetorch.Size([96, 3, 96])
但更令人擔憂的是，這個函數實際上可能會錯誤地用到batch 維度，還會把不同圖像的屬性混到一起。如果在代碼中隱藏了這個維度，可能會產生一些本來很容易避免的、討厭的bug。
陷阱2：通過對齊進行廣播
張量最有用的地方是它們可以在不直接需要for 循環的情況下快速執行數組運算。為此，要直接對齊維度，以便廣播張量。同樣，這是按照慣例和代碼文檔實現的，這使排列維度變得「容易」。例如，假設我們想對上圖應用掩碼。
# height x width
mask = torch.randint(0, 2, [96, 96]).byte()
masktry:
ims.masked_fill(mask, 0)
except RuntimeError:
error = "Broadcasting fail %s %s"%(mask.shape, ims.shape)
error'Broadcasting fail torch.Size([96, 96]) torch.Size([6, 96, 96, 3])'
這裡的失敗的原因是：即便我們知道要建立掩碼的形狀，廣播的規則也沒有正確的語義。為了讓它起作用，你需要使用view 或squeeze 這些我最不喜歡的函數。
# either
mask = mask.unsqueeze(-1)
# or
mask = mask.view(96, 96, 1)
# height x width x channels
ims.masked_fill(mask, 1)[0]
注意，最左邊的維度不需要進行這樣的運算，所以這裡有些抽象。但閱讀真正的代碼後會發現，右邊大量的view 和squeeze 變得完全不可讀。
陷阱3：通過註釋訪問
看過上面兩個問題後，你可能會認為只要足夠小心，運行時就會捕捉到這些問題。但是即使很好地使用了廣播和索引的組合，也可能會造成很難捕捉的問題。
a = ims[1].mean(2, keepdim=True)
# height x width x 1
# (Lots of code in between)
# .......................
# Code comment explaining what should be happening.
dim = 1
b = a + ims.mean(dim, keepdim=True)[0]
# (Or maybe should be a 2? or a 0?)
index = 2
b = a + ims.mean(dim, keepdim=True)[0]
b
我們在此假設編碼器試著用歸約運算和維度索引將兩個張量結合在一起。（說實話這會兒我已經忘了維度代表什麼。）
重點在於無論給定的維度值是多少，代碼都會正常運行。這裡的註釋描述的是在發生什麼，但是代碼本身在運行時不會報錯。
Named Tensor：原型
根據這些問題，我認為深度學習代碼應該轉向更好的核心對象。為了好玩，我會開發一個新的原型。目標如下：
- 維度應該有人類可讀的名字。
- 函數中不應該有維度參數。
- 廣播應該通過名稱匹配。
- 轉換應該是顯式的。
- 禁止基於維度的索引。
- 應該保護專用維度。
為了試驗這些想法，我建立了一個叫做NamedTensor 的庫。目前它只用於PyTorch，但從理論上講類似的想法也適用於其他框架。
建議1：分配名稱
庫的核心是封裝了張量的對象，並給每個維度提供了名稱。我們在此用維度名稱簡單地包裝了給定的torch 張量。
named_ims = NamedTensor(ims, ("batch", "height", "width", "channels"))
named_ims.shapeOrderedDict([('batch', 6), ('height', 96), ('width', 96), ('channels', 3)])
此外，該庫有針對PyTorch 構造函數的封裝器，可以將它們轉換為命名張量。
ex = ntorch.randn(dict(height=96, width=96, channels=3))
ex
大多數簡單的運算只是簡單地保留了命名張量的屬性。
ex.log() # or ntorch.log(ex) None
建議2：訪問器和歸約
名字的第一個好處是可以完全替換掉維度參數和軸樣式參數。例如，假設我們要對每列進行排序。
sortex, _ = ex.sort("width")
sortex
另一個常見的操作是在匯集了一個或多個維度的地方進行歸約。
named_ims.mean("batch")named_ims.mean(("batch", "channels"))
建議3：廣播和縮並
提供的張量名稱也為廣播操作提供了基礎。當兩個命名張量間存在二進制運算時，它們首先要保證所有維度都和名稱匹配，然後再應用標準的廣播。為了演示，我們回到上面的掩碼示例。在此我們簡單地聲明了一下掩碼維度的名稱，然後讓庫進行廣播。
im = NamedTensor(ims[0], ("height", "width", "channels"))
im2 = NamedTensor(ims[1], ("height", "width", "channels"))
mask = NamedTensor(torch.randint(0, 2, [96, 96]).byte(), ("height", "width"))
im.masked_fill(mask, 1)
加和乘等簡單運算可用於標準矩陣。
im * mask.double()
在命名向量間進行張量縮併的更普遍的特徵是dot 方法。張量縮並是einsum 背後的機制，是一種思考點積、矩陣-向量乘積、矩陣-矩陣乘積等泛化的優雅方式。
# Runs torch.einsum(ijk,ijk->jk, tensor1, tensor2)
im.dot("height", im2).shapeOrderedDict([('width', 96), ('channels', 3)])# Runs torch.einsum(ijk,ijk->il, tensor1, tensor2)
im.dot("width", im2).shapeOrderedDict([('height', 96), ('channels', 3)])# Runs torch.einsum(ijk,ijk->l, tensor1, tensor2)
im.dot(("height", "width"), im2).shapeOrderedDict([('channels', 3)])
類似的註釋也可用於稀疏索引（受einindex 庫的啟發）。這在嵌入查找和其他稀疏運算中很有用。
pick, _ = NamedTensor(torch.randint(0, 96, [50]).long(), ("lookups",)) \
.sort("lookups")
# Select 50 random rows.
im.index_select("height", pick)
建議4：維度轉換
在後台計算中，所有命名張量都是張量對象，因此維度順序和步幅這樣的事情就尤為重要。transpose 和view 等運算對於保持維度的順序和步幅至關重要，但不幸的是它們很容易出錯。
那麼，我們來考慮領域特定語言shift，它大量借鑒了Alex Rogozhnikov 優秀的einops 包（https://github.com/arogozhnikov/einops）。
tensor = NamedTensor(ims[0], ("h", "w", "c"))
tensor
維度轉換的標準調用。
tensor.transpose("w", "h", "c")
拆分和疊加維度。
tensor = NamedTensor(ims[0], ("h", "w", "c"))
tensor.split(h=("height", "q"), height=8).shape
OrderedDict([('height', 8), ('q', 12), ('w', 96), ('c', 3)])
tensor = NamedTensor(ims, ('b', 'h', 'w', 'c'))
tensor.stack(bh = ('b', 'h')).shape
OrderedDict([('bh', 576), ('w', 96), ('c', 3)])
鏈接Ops。
tensor.stack(bw=('b', 'w')).transpose('h', 'bw', 'c')
這裡還有一些einops 包中有趣的例子。
tensor.split(b=('b1', 'b2'), b1=2).stack(a=('b2', 'h'), d=('b1', 'w'))\
.transpose('a', 'd', 'c')
建議5：禁止索引
一般在命名張量範式中不建議用索引，而是用上面的index_select 這樣的函數。
在torch 中還有一些有用的命名替代函數。例如unbind 將維度分解為元組。
tensor = NamedTensor(ims, ('b', 'h', 'w', 'c'))
# Returns a tuple
images = tensor.unbind("b")
images[3]
get 函數直接從命名維度中選擇了一個切片。
# Returns a tuple
images = tensor.get("b", 0).unbind("c")
images[1]
最後，可以用narrow 代替花哨的索引。但是你一定要提供一個新的維度名稱（因為它不能再廣播了）。
tensor.narrow( 30, 50, h='narowedheight').get("b", 0)
建議6：專用維度
最後，命名張量嘗試直接隱藏不應該被內部函數訪問的維度。mask_to 函數會保留左邊的掩碼，它可以使任何早期的維度不受函數運算的影響。最簡單的使用掩碼的例子是用來刪除batch 維度的。
def bad_function(x, y):
# Accesses the private batch dimension
return x.mean("batch")
x = ntorch.randn(dict(batch=10, height=100, width=100))
y = ntorch.randn(dict(batch=10, height=100, width=100))
try:
bad_function(x.mask_to("batch"), y)
except RuntimeError as e:
error = "Error received: " + str(e)
error'Error received: Dimension batch is masked'
這是弱動態檢查，可以通過內部函數關閉。在將來的版本中，也許我們會添加函數註釋來lift 未命名函數，來保留這些屬性。
示例：神經註意力
為了說明為什麼這些選擇會帶來更好的封裝屬性，我們來思考一個真實世界中的深度學習例子。這個例子是我的同事Tim Rocktashel 在一篇介紹einsum 的博客文章中提出來的。和原始的PyTorch 相比，Tim 的代碼是更好的替代品。雖然我同意enisum 是一個進步，但它還是存在很多上述陷阱。
下面來看神經註意力的問題，它需要計算，
首先我們要配置參數。
def random_ntensors(names, num=1, requires_grad=False):
tensors = [ntorch.randn(names, requires_grad=requires_grad)
for i in range(0, num)]
return tensors[0] if num == 1 else tensors
class Param:
def __init__(self, in_hid, out_hid):
torch.manual_seed(0)
self.WY, self.Wh, self.Wr, self.Wt = \
random_ntensors(dict(inhid=in_hid, outhid=out_hid),
num=4, requires_grad=True)
self.bM, self.br, self.w = \
random_ntensors(dict(outhid=out_hid),
num=3,
requires_grad=True)
現在考慮這個函數基於張量的enisum 實現。
# Einsum Implementation
import torch.nn.functional as F
def einsum_attn(params, Y, ht, rt1):
# -- [batch_size x hidden_dimension]
tmp = torch.einsum("ik,kl->il", [ht, params.Wh.values]) + \
torch.einsum("ik,kl->il", [rt1, params.Wr.values])
Mt = torch.tanh(torch.einsum("ijk,kl->ijl", [Y, params.WY.values]) + \
tmp.unsqueeze(1).expand_as(Y) + params.bM.values)
# -- [batch_size x sequence_length]
at = F.softmax(torch.einsum("ijk,k->ij", [Mt, params.w.values]), dim=-1)
# -- [batch_size x hidden_dimension]
rt = torch.einsum("ijk,ij->ik", [Y, at]) + \
torch.tanh(torch.einsum("ij,jk->ik", [rt1, params.Wt.values]) +
params.br.values)
# -- [batch_size x hidden_dimension], [batch_size x sequence_dimension]
return rt, at
該實現是對原版PyTorch 實現的改進。它刪除了這項工作必需的一些view 和transpose。但它仍用了squeeze，引用了private batch dim，使用了非強制的註釋。
接下來來看namedtensor 版本：
def namedtensor_attn(params, Y, ht, rt1):
tmp = ht.dot("inhid", params.Wh) + rt1.dot("inhid", params.Wr)
at = ntorch.tanh(Y.dot("inhid", params.WY) + tmp + params.bM) \
.dot("outhid", params.w) \
.softmax("seqlen")
rt = Y.dot("seqlen", at).stack(inhid=('outhid',)) + \
ntorch.tanh(rt1.dot("inhid", params.Wt) + params.br)
return rt, at
該代碼避免了三個陷阱：
- （陷阱1）該代碼從未提及batch 維度。
- （陷阱2）所有廣播都是直接用縮並完成的，沒有views。
- （陷阱3）跨維度的運算是顯式的。例如，softmax 明顯超過了seqlen。
# Run Einsum
in_hid = 7; out_hid = 7
Y = torch.randn(3, 5, in_hid)
ht, rt1 = torch.randn(3, in_hid), torch.randn(3, in_hid)
params = Param(in_hid, out_hid)
r, a = einsum_attn(params, Y, ht, rt1)# Run Named Tensor (hiding batch)
Y = NamedTensor(Y, ("batch", "seqlen", "inhid"), mask=1)
ht = NamedTensor(ht, ("batch", "inhid"), mask=1)
rt1 = NamedTensor(rt1, ("batch", "inhid"), mask=1)
nr, na = namedtensor_attn(params, Y, ht, rt1)
結論/請求幫助
深度學習工具可以幫助研究人員實現標準模型，但它們也影響了研究人員的嘗試。我們可以用現有工具很好地構建模型，但編程實踐無法擴展到新模型。（例如，我們最近研究的是離散隱變量模型，它通常有許多針對特定問題的變量，每個變量都有自己的變量維度。這個設置幾乎可以立即打破當前的張量範式。）
這篇博文只是這種方法的原型。如果你感興趣，我很願意為構建這個庫作出貢獻。還有一些想法：
擴展到PyTorch 之外：我們是否可以擴展這種方法，使它支持NumPy 和TensorFlow？
與PyTorch 模塊交互：我們是否可以通過類型註釋「lift」PyTorch 模塊，從而了解它們是如何改變輸入的？
錯誤檢查：我們是否可以給提供前置條件和後置條件的函數添加註釋，從而自動檢查維度？
原文鏈接：http://nlp.seas.harvard.edu/NamedTensor?fbclid=IwAR2FusFxf-c24whTSiF8B3R2EKz_-zRfF32jpU8D-F5G7rreEn9JiCfMl48
本文為機器之心編譯，轉載請聯繫本公眾號獲得授權。
✄------------------------------------------------
加入機器之心（全職記者/ 實習生）：hr@jiqizhixin.com
投稿或尋求報導：content @jiqizhixin.com
廣告& 商務合作：bd@jiqizhixin.com
留言
張貼留言