如何利用張量提升內(nèi)存使用效率

注：本文轉(zhuǎn)自賽靈思中文社區(qū)論壇，源文鏈接在此。本文原作者為XILINX工程師。

以下為個(gè)人譯文，僅供個(gè)人學(xué)習(xí)記錄參考之用，如有疏漏之處，還請(qǐng)不吝賜教。

在數(shù)據(jù)處理中，對(duì)原始數(shù)據(jù)進(jìn)行重塑或重新排序并創(chuàng)建多個(gè)副本是很常見(jiàn)的行為。無(wú)論執(zhí)行任何新步驟，都會(huì)創(chuàng)建新副本。隨著程序的增大，占用的內(nèi)存也會(huì)增大，我?guī)缀鯊奈纯紤]過(guò)這個(gè)問(wèn)題，直到遇到了“內(nèi)存不足”錯(cuò)誤。

張量 (tensor) 的神奇之處在于多個(gè)張量可引用同一存儲(chǔ)空間，從而顯著提升內(nèi)存使用效率。

在下一篇的文章中，我將聊一聊張量所具有的更神奇的屬性，即跟蹤上級(jí)操作，但在本文中，我將主要介紹內(nèi)存優(yōu)化方面的內(nèi)容。

張量 (tensor) 的神奇之處在于多個(gè)張量可以引用同一存儲(chǔ)空間（即包含給定類型的數(shù)字的連續(xù)內(nèi)存區(qū)塊）。此行為由 torch.storage 進(jìn)行管理。

每個(gè)張量都包含 .storage 屬性，用于顯示內(nèi)存中存儲(chǔ)的張量?jī)?nèi)容。

在下一篇的文章中，我將聊一聊張量所具有的更神奇的屬性，即跟蹤上級(jí)操作，但在本文中，我將主要介紹內(nèi)存優(yōu)化方面的內(nèi)容。

注：全新 Vitis AI 1.2 發(fā)行版將首次為 PyTorch 提供支持。本文對(duì)于新增對(duì)此熱門框架的支持表示祝賀，并提供了 1 個(gè) PyTorch 專用的 Jupyter Notebook 格式示例。

輸入 [1]：

import torch
a = torch.randint(0, 9, (5,3))
a

輸出 [1]：

tensor([[4, 1, 6],

[0, 8, 8],

[1, 2, 1],

[0, 5, 7],

[0, 0, 7]])

輸入 [2]：

a.storage()

輸出 [2]：

4

1

6

0

8

8

1

2

1

0

5

7

0

0

7

[torch.LongStorage of size 15]

輸入 [3]：

a.shape

輸出 [3]：

torch.Size([5, 3])

我們可能需要對(duì)原始“a”張量進(jìn)行轉(zhuǎn)置 (transpose) 和平展 (flatten) 處理。

何必為了相同數(shù)據(jù)浪費(fèi)雙倍內(nèi)存？哪怕數(shù)據(jù)只是形狀 (shape) 不同，也沒(méi)有必要。

輸入 [4]：

b = torch.transpose(a, 0, 1)
b

輸出 [4]：

tensor([[4, 0, 1, 0, 0],

[1, 8, 2, 5, 0],

[6, 8, 1, 7, 7]])

a 和 b 確實(shí)是指向相同存儲(chǔ)空間的張量。

兩者表現(xiàn)方式不同，原因在于我們使用 stride 函數(shù)指令其按不同順序讀取該存儲(chǔ)空間。

b 的 stride 值為 (1,3)，即讀取存儲(chǔ)空間時(shí)，每隔 1 個(gè)元素都必須跳至下一行，并且每隔 3 個(gè)元素必須跳至下一列。

輸入 [5]：

b.stride(), a.stride()

輸出 [5]：

((1, 3), (3, 1))

我們可以從 a 或 b 訪問(wèn)數(shù)據(jù)，或者也可以從原始存儲(chǔ)空間直接訪問(wèn)數(shù)據(jù)。

但如果從存儲(chǔ)空間訪問(wèn)，則讀取的值將不再是張量。

輸入 [6]：

a[1,2], b[2,1], a.storage()[5], b.storage()[5]

輸出 [6]：

(tensor(8), tensor(8), 8, 8)

現(xiàn)在，令我感到疑惑不解的是，我發(fā)現(xiàn)這些張量的值神奇般地自行發(fā)生了改變：

更改 a 時(shí)，b 也變了。

輸入 [7]：

a[0,0] = 10
b[0,0]

輸出 [7]：

tensor(10)

發(fā)生這種狀況的原因是因?yàn)椋瑥膬?nèi)存角度來(lái)看，張量即經(jīng)過(guò)排序的存儲(chǔ)空間表示法。

從同一存儲(chǔ)空間生成的 2 個(gè)張量并非獨(dú)立張量，而且我必須牢記的是，當(dāng)我每次更改 1 個(gè)張量后，指向相同存儲(chǔ)空間的所有其它張量也都會(huì)被修改。

可見(jiàn)，即使高效的內(nèi)存利用方式也難免有其缺點(diǎn)！

子集

通過(guò)原始數(shù)據(jù)的子集仍然能夠有效利用內(nèi)存。

新的張量仍然指向原始存儲(chǔ)空間的子集。

輸入 [8]：

c = a[0:2, 0:2]
c

輸出 [8]：

tensor([[10, 1],

[ 0, 8]])

輸入 [9]：

c[0,0]=77
a

輸出 [9]：

tensor([[77, 1, 6],

[ 0, 8, 8],

[ 1, 2, 1],

[ 0, 5, 7],

[ 0, 0, 7]])

inplace 運(yùn)算符

inplace 運(yùn)算符即無(wú)需創(chuàng)建張量副本就可以直接對(duì)存儲(chǔ)空間進(jìn)行操作的函數(shù)。這些運(yùn)算符通常具有易于識(shí)別的名稱且以下劃線結(jié)尾。

輸入 [10]：

a.zero_()
b

輸出 [10]：

tensor([[0, 0, 0, 0, 0],

[0, 0, 0, 0, 0],

[0, 0, 0, 0, 0]])

張量克隆

如果確實(shí)需要 1 個(gè)獨(dú)立的新張量，可以對(duì)其進(jìn)行克隆。

這樣也會(huì)創(chuàng)建新的存儲(chǔ)空間。

輸入 [11]：

a_clone = a.clone()
a_clone[0,0] = 55
a_clone

輸出 [11]：

tensor([[55, 0, 0],

[ 0, 0, 0],

[ 0, 0, 0],

[ 0, 0, 0],

[ 0, 0, 0]])

輸入 [12]：

a

輸出 [12]：

tensor([[0, 0, 0],

[0, 0, 0],

[0, 0, 0],

[0, 0, 0],

[0, 0, 0]])

為連續(xù)張量重組存儲(chǔ)空間

部分函數(shù)僅適用于連續(xù)張量。

對(duì) a 進(jìn)行轉(zhuǎn)置時(shí)，通過(guò)在 b 中分配來(lái)自存儲(chǔ)空間的非連續(xù)矩陣值，生成了新的張量。

輸入 [13]：

a.is_contiguous()

輸出 [13]：

True

輸入 [14]：

b.is_contiguous()

輸出 [14]：

False

我們可將 b 設(shè)為連續(xù)張量，但這將導(dǎo)致 b 生成經(jīng)過(guò)重組的新存儲(chǔ)空間，從而導(dǎo)致 a 和 b 永遠(yuǎn)無(wú)法成為獨(dú)立張量：

輸入 [15]：

b = b.contiguous()
b[0,0] = 18
a[0,0]

輸出 [15]：

tensor(0)

輸入 [16]：

b.is_contiguous()

輸出 [16]：

True

輸入 [17]：

a.is_contiguous()

輸出 [17]：

True

審核編輯：湯梓紅