PyTorch中的memory format - NCHW和channels last

PyTorch中的memory format - NCHW和channels last

參考Channels Last Memory Format in PyTorch，在PyTorch中，張量的記憶體格式分為NCHW（非channels last）和channels last兩種。這兩種記憶體格式的差別在於維度的排列順序不同；前者按batch size, channel, height, width的順序排列；後者則把channel維度移到最後，其它維度的順序則保持不變。

NCHW張量和channels last張量

我們直接從程式碼層面來觀察兩種記憶體格式的不同之處。首先創造一個NCHW張量：

x=torch.arange(8).reshape(1,2,2,2)

查看其內容：

# tensor viewprint(x)# flattenedprint(x.flatten())

tensor([[[[0, 1], [2, 3]], [[4, 5], [6, 7]]]]) tensor([0, 1, 2, 3, 4, 5, 6, 7])

它的形狀：

print(x.shape)

torch.Size([1, 2, 2, 2])

它的步長：

# strideprint(x.stride())

(8, 4, 2, 1)

如果我們把它轉成channels last記憶體格式：

x_cl=x.to(memory_format=torch.channels_last)# tensor viewprint(x_cl)# flattenedprint(x_cl.flatten())

tensor([[[[0, 1], [2, 3]], [[4, 5], [6, 7]]]]) tensor([0, 1, 2, 3, 4, 5, 6, 7])

可以發現print出來的元素排列竟然都沒變！

查看channels last張量的形狀：

print(x_cl.shape)

torch.Size([1, 2, 2, 2])

仍然維持不變。

再繼續查看channels last張量的步長：

# stride print(x_cl.stride())

為：

(8, 1, 4, 2)

總算跟NCHW張量不一樣了。

張量的底層記憶體

為什麼兩種記憶體格式的張量之形狀和print的結果都一樣，只有步長不一樣呢？這得從張量的底層記憶體說起。

張量（torch.tensor）的底層為（torch.storage.TypedStorage），我們可以透過以下程式碼查看張量中各元素在記憶體中的實際排列。

首先是NCHW張量：

# memory layoutprint(torch.tensor(list(x.storage())))

tensor([0, 1, 2, 3, 4, 5, 6, 7])

接著是channels last張量：

# memory layoutprint(torch.tensor(list(x_cl.storage())))

tensor([0, 4, 1, 5, 2, 6, 3, 7])

可以看到張量被轉成channels last後，各元素在記憶體中的排列順序確實是有變化的。

至於為什麼順序會變成這樣呢？因為變為channels last格式後，張量最內層的維度變為C，各維度由內往外的順序則變為CWHN。

如果我們想要按照CWHN的順序由內往外存取張量，應該要依照這個順序：

• N = 0, C = 0, H = 0, W = 0
• N = 0, C = 1, H = 0, W = 0
• N = 0, C = 0, H = 0, W = 1
• N = 0, C = 1, H = 0, W = 1
• N = 0, C = 0, H = 1, W = 0
• N = 0, C = 1, H = 1, W = 0
• N = 0, C = 0, H = 1, W = 1
• N = 0, C = 1, H = 1, W = 1

用它們去存取四維張量：

tensor([[[[0, 1], [2, 3]], [[4, 5], [6, 7]]]])

得到的結果正是剛剛看到的tensor([0, 4, 1, 5, 2, 6, 3, 7])！

邏輯結構與底層記憶體間的橋樑

那麼為何兩種記憶體格式的張量的形狀一樣，print的結果也一致呢？

這是因為PyTorch中的張量（torch.tensor）只是一個「邏輯層面」的抽象的概念，它實際上是對實體記憶體（torch.storage.TypedStorage）的包裝。

張量的shape描述的是張量的「邏輯結構」，固定按NCHW的順序排列。因為to(memory_format=torch.channels_last)只改變張量的「底層」，不改變頂層的「邏輯結構」，所以兩種記憶體格式的張量的shape一致。

storage是底層記憶體，張量則是一個邏輯結構，它們之間需要有一座做轉換的橋樑。stride便是一個將張量的「邏輯座標」對應到「底層記憶體索引」的映射規則。

來看看x，它是一個NCHW張量，有四個維度。其形狀為[1, 2, 2, 2]，內容則為：

tensor([[[[0, 1], [2, 3]], [[4, 5], [6, 7]]]])

我們用[0, 0, 0, 0]到[0, 1, 1, 1]的索引去存取張量，實際上對應到底層記憶體的哪裡呢？這裡就得用到stride了：

它的stride為[8, 4, 2, 1]，假設邏輯層面的四維索引為[i_n, i_c, i_h, i_w]，則經過換算得到的底層記憶體的index為i_n * 8 + i_c * 4 + i_h * 2 + i_w * 1。

代入實際的四維索引看看：

• [0, 0, 0, 0]：0 * 8, 0 * 4, 0 * 2, 0 * 1 = 0
• [0, 0, 0, 1]：0 * 8, 0 * 4, 0 * 2, 1 * 1 = 1
• [0, 0, 1, 0]：0 * 8, 0 * 4, 1 * 2, 0 * 1 = 2
• [0, 0, 1, 1]：0 * 8, 0 * 4, 1 * 2, 1 * 1 = 3
• [0, 1, 0, 0]：0 * 8, 1 * 4, 0 * 2, 0 * 1 = 4
• [0, 1, 0, 1]：0 * 8, 1 * 4, 0 * 2, 1 * 1 = 5
• [0, 1, 1, 0]：0 * 8, 1 * 4, 1 * 2, 0 * 1 = 6
• [0, 1, 1, 1]：0 * 8, 1 * 4, 1 * 2, 1 * 1 = 7

所以從[0, 0, 0, 0]到[0, 1, 1, 1]，算出來的底層記憶體索引分別為由0到7，接著用算出來的底層記憶體索引去存取底層記憶體：

tensor([0, 1, 2, 3, 4, 5, 6, 7])

正好可以得到[0, 1, 2, 3, 4, 5, 6, 7]的值。

接著來看channels last張量，它的stride為[8, 1, 4, 2]，假設邏輯層面的四維索引為[i_n, i_c, i_h, i_w]，則經過換算得到的底層記憶體的index為i_n * 8 + i_c * 1 + i_h * 4 + i_w * 2。

代入實際的四維索引看看：

• [0, 0, 0, 0]：0 * 8, 0 * 1, 0 * 4, 0 * 2 = 0
• [0, 0, 0, 1]：0 * 8, 0 * 1, 0 * 4, 1 * 2 = 2
• [0, 0, 1, 0]：0 * 8, 0 * 1, 1 * 4, 0 * 2 = 4
• [0, 0, 1, 1]：0 * 8, 0 * 1, 1 * 4, 1 * 2 = 6
• [0, 1, 0, 0]：0 * 8, 1 * 1, 0 * 4, 0 * 2 = 1
• [0, 1, 0, 1]：0 * 8, 1 * 1, 0 * 4, 1 * 2 = 3
• [0, 1, 1, 0]：0 * 8, 1 * 1, 1 * 4, 0 * 2 = 5
• [0, 1, 1, 1]：0 * 8, 1 * 1, 1 * 4, 1 * 2 = 7

所以從[0, 0, 0, 0]到[0, 1, 1, 1]，算出來的底層記憶體索引分別為[0, 2, 4, 6, 1, 3, 5, 7]，用算出來的索引去存取底層記憶體：

tensor([0, 4, 1, 5, 2, 6, 3, 7])

很巧地，得到的結果也是[0, 1, 2, 3, 4, 5, 6, 7]。

這表示不管張量底層的記憶體格式如何，只要我們用邏輯層面的四維索引去存取張量，就會得到一樣的結果！

而print函數做的正是這件事：它會算出邏輯層面的NCHW四維索引對應到底層記憶體的index，取出對應的元素後，再將這些元素整理成 N,C,H,W 四個維度排版印出來。