Vitis AI - 如何利用張量提升內(nèi)存使用效率達(dá)到內(nèi)存優(yōu)化效果

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

張量 （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） 不同，也沒有必要。

輸入 ［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訪問數(shù)據(jù)，或者也可以從原始存儲(chǔ)空間直接訪問數(shù)據(jù)。

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

輸入 ［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)存角度來看，張量即經(jīng)過排序的存儲(chǔ)空間表示法。

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

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

子集

通過原始數(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)算符即無需創(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í)，通過在b中分配來自存儲(chǔ)空間的非連續(xù)矩陣值，生成了新的張量。

輸入 ［13］：

a.is_contiguous()

輸出 ［13］：

True

輸出 ［14］：

b.is_contiguous()

輸出 ［14］：

False

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

輸入 ［15］：

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

輸出 ［15］：

tensor(0)

輸出 ［16］：

b.is_contiguous()

輸出 ［16］：

True

輸出 ［16］：

a.is_contiguous()

True