深度學習框架中的「張量」不好用？也許我們需要重新定義Tensor了

科技 02-14

選自

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.shape

torch.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

]).shape

torch.Size([

96

, 3
, 96
])

但更令人擔憂的是，這個函數實際上可能會錯誤地用到 batch 維度，還會把不同圖像的屬性混到一起。如果在代碼中隱藏了這個維度，可能會產生一些本來很容易避免的、討厭的 bug。

陷阱 2：通過對齊進行廣播

張量最有用的地方是它們可以在不直接需要 for 循環的情況下快速執行數組運算。為此，要直接對齊維度，以便廣播張量。同樣，這是按照慣例和代碼文檔實現的，這使排列維度變得「容易」。例如，假設我們想對上圖應用掩碼。

# height x width


 
mask = torch.randint(
0
, 2
, [96
, 96
]).byte()
mask

try

:
    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.shape

OrderedDict([(

"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).shape

OrderedDict([(

"width"

, 96
), ("channels"
, 3
)])

# Runs torch.einsum(ijk,ijk->il, tensor1, tensor2)


im.dot("width"
, im2).shape

OrderedDict([(

"height"

, 96
), ("channels"
, 3
)])

# Runs torch.einsum(ijk,ijk->l, tensor1, tensor2)


im.dot(("height"
, "width"
), im2).shape

OrderedDict([(

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