PyTorch的slice究竟做了什麼事？－－ stride和stride_offset的妙用

直接上代碼，首先建立一個形狀為3x1x1x2的張量：

x=torch.arange(6).reshape(3,1,1,2)print(x)print(x.shape)print(x.stride())

其內容，形狀及步長如下：

# x tensor([[[[0, 1]]], [[[2, 3]]], [[[4, 5]]]]) # x.shape torch.Size([3, 1, 1, 2]) # x.stride (2, 2, 2, 1)

在最後一個維度做裁切，只取第0個元素，得到x_slice：

x_slice = x[:,:,:,0:1] print(x_slice) print(x_slice.shape) print(x_slice.stride())

其內容，形狀及步長如下：

# x_slice tensor([[[[0]]], [[[2]]], [[[4]]]]) # x_slice.shape torch.Size([3, 1, 1, 1]) # x_slice.stride (2, 2, 2, 1)

裁切得到的x_slice的內容和形狀都符合預期，但步長看起來卻有點奇怪：既然最後三個維度的尺寸都是1，為何步長不是(1, 1, 1, 1)呢？

Tensor Views

其實這跟一個叫Tensor Views的概念有關：

PyTorch allows a tensor to be a View of an existing tensor. View tensor shares the same underlying data with its base tensor. Supporting View avoids explicit data copy, thus allows us to do fast and memory efficient reshaping, slicing and element-wise operations.

PyTorch 允許一個張量（tensor）成為另一個既有張量的 View（視圖）。View 張量與其基礎張量（base tensor）共享相同的底層資料。
支援 View 可以避免顯式的資料複製，因此能夠以更快速、且更節省記憶體的方式進行 reshape、slice，以及逐元素（element-wise）運算。

當中提到了slice，也就是我們方才使用的[:,:,:,0:1]操作。

說明中還提到了Veiw 張量與其基礎張量共享底層資料，我們可以寫程式來驗證這一點。

untyped_storage().data_ptr()指令可獲取張量底層記憶體的指標，我們可以用它來查看x和x_slice的記憶體位址：

>>> x.untyped_storage().data_ptr() 101369920 >>> x_slice.untyped_storage().data_ptr() 101369920

發現兩者一樣！這也證實了兩者確實共享了底層記憶體。

我們可以再進一步驗證，把x_slice的第0個元素改成555試試：

>>> x_slice[0] = 555 >>> x_slice tensor([[[[555]]], [[[ 2]]], [[[ 4]]]]) >>> x tensor([[[[555, 1]]], [[[ 2, 3]]], [[[ 4, 5]]]])

可以發現x的第0個元素也確實被改動了！這更進一步證明了它們共享記憶體的事實。

Tensor的本質

從以上實驗可以發現，x和x_slice這兩個張量的形狀、步長不同，卻共享同一塊底層記憶體。PyTorch底層就是這麼實作的。

PyTorch 中Tensor的源碼位於aten/src/ATen/core/Tensor.h，但它只是一個wrapper。
張量真正的實作位於pytorch/c10/core/TensorImpl.h，來看看TensorImpl的定義：

structC10_APITensorImpl:publicc10::intrusive_ptr_target{protected:Storage storage_;// ...protected:// ...c10::impl::SizesAndStrides sizes_and_strides_;int64_tstorage_offset_=0;

TensorImpl中最核心的成員變數為storage_，它表示一塊連續記憶體，而TensorImpl本身代表的則是一個多維的矩陣，是一個抽象的概念。這兩者之間需要有一個橋樑做轉換，這座橋樑便是sizes_and_strides和storage_offset_成員變數。

當我們要存取抽象的張量中的元素時，我們會用到多維索引。sizes_and_strides這個成員變數記錄了張量的各維尺寸和步長，storage_offset_則記錄了抽象的張量相對於底層記憶體的offset，有了這些資訊，我們便能將多維索引換算成對應元素在底層記憶體中的memory offset。

這些元素可分別以以下成員函數存取：

• sizes成員函數

c10/core/TensorImpl.h

/** * Return a reference to the sizes of this tensor. This reference remains * valid as long as the tensor is live and not resized. */IntArrayRefsizes()const{if(C10_UNLIKELY(matches_policy(SizesStridesPolicy::CustomSizes))){returnsizes_custom();}returnsizes_and_strides_.sizes_arrayref();}

• strides成員函數

c10/core/TensorImpl.h

/** * Return a reference to the strides of this tensor. This reference remains * valid as long as the tensor is live and not restrided. */IntArrayRefstrides()const{if(C10_UNLIKELY(matches_policy(SizesStridesPolicy::CustomStrides))){returnstrides_custom();}returnsizes_and_strides_.strides_arrayref();}

• storage成員函數

c10/core/TensorImpl.h

/** * Return the underlying storage of a Tensor. Multiple tensors may share * a single storage. A Storage is an impoverished, Tensor-like class * which supports far less operations than Tensor. * * Avoid using this method if possible; try to use only Tensor APIs to perform * operations. */ TENSORIMPL_MAYBE_VIRTUAL const Storage& storage() const { if (C10_UNLIKELY(storage_access_should_throw_)) { throw_storage_access_error(); } return storage_; }

• storage_offset成員函數

c10/core/TensorImpl.h

/** * Return the offset in number of elements into the storage that this * tensor points to. Most tensors have storage_offset() == 0, but, * for example, an index into a tensor will have a non-zero storage_offset(). * * WARNING: This is NOT computed in bytes. */int64_tstorage_offset()const{// TODO: maybe this should be toggled by stridesif(C10_UNLIKELY(matches_policy(SizesStridesPolicy::CustomSizes))){returnstorage_offset_custom();}returnstorage_offset_;}

在Python中則可透過以下方法/成員變數存取：

x.shape x.stride() x.storage() x.storage_offset()

切片張量的步長

回到剛才的問題，為何x_slice的步長是(2, 2, 2, 1)而不是(1, 1, 1, 1)呢？我們得從張量的底層記憶體看起。

slice後的張量x_slice與base tensorx共享一塊記憶體：

>>>x.storage()55512345[torch.storage.TypedStorage(dtype=torch.int64,device=cpu)of size6]>>>x_slice.storage()55512345[torch.storage.TypedStorage(dtype=torch.int64,device=cpu)of size6]

但x_slice又只是x的一部份：

>>>x_slice tensor([[[[555]]],[[[2]]],[[[4]]]])

可想而知，x_slice中的各元素在底層記憶體中並不是連續的。

假設四維張量在邏輯層面的四維索引為[i_n, i_c, i_h, i_w]，strides為(s_n, s_c, s_h, s_w)，我們可以透過以下公式將張量的四維索引轉換成元素在底層記憶體裡的memory offset：i_n * s_n + i_c * s_c + i_h * s_h + i_w * s_w。

因為我們已知strides為(2, 2, 2, 1)，代入後公式可以化為：i_n * 2 + i_c * 2 + i_h * 2 + i_w * 1。

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

• [0, 0, 0, 0]：0 * 2, 0 * 2, 0 * 2, 0 * 1 = 0
• [1, 0, 0, 0]：1 * 2, 0 * 2, 0 * 2, 0 * 1 = 2
• [2, 0, 0, 0]：2 * 2, 0 * 2, 0 * 2, 0 * 1 = 4

也就是說：x_slice裡的三個元素分別對應底層記憶體中的第0, 2, 4個元素。

實際驗證一下：

>>>x_slice[0,0,0,0]tensor(555)>>>x_slice[1,0,0,0]tensor(2)>>>x_slice[2,0,0,0]tensor(4)

確實符合公式計算出來的結果。

再來看一下x_slice：

>>> x.stride() (2, 2, 2, 1)

第一個維度的步長為2，表示四維索引的第一個維度每加1，對應元素在底層記憶體裡的位置就會加2，這也跟我們觀察到的相符。

再來看另外一個例子，首先建立一個形狀為3x1x1x4的張量：

>>>x=torch.arange(12).reshape(3,1,1,4)>>>x tensor([[[[0,1,2,3]]],[[[4,5,6,7]]],[[[8,9,10,11]]]])>>>x.shape torch.Size([3,1,1,4])>>>x.stride()(4,4,4,1)

在最後一個維度做裁切，每間隔2取一個元素，得到新的x_slice：

>>> x_slice = x[:,:,:,0::2] >>> x_slice tensor([[[[ 0, 2]]], [[[ 4, 6]]], [[[ 8, 10]]]]) >>> x_slice.shape torch.Size([3, 1, 1, 2]) >>> x_slice.stride() (4, 4, 4, 2)

這次最後一個維度的步長變成2了，表示四維索引最後一個維度每加1，在底層記憶體裡的位置就得加2，來實際驗證一下：

• [0, 0, 0, 0]：0 * 4, 0 * 4, 0 * 4, 0 * 2 = 0
• [0, 0, 0, 1]：0 * 4, 0 * 4, 0 * 4, 1 * 2 = 2
• [1, 0, 0, 0]：1 * 4, 0 * 4, 0 * 4, 0 * 2 = 4
• [1, 0, 0, 1]：1 * 4, 0 * 4, 0 * 4, 1 * 2 = 6
• [2, 0, 0, 0]：1 * 4, 0 * 4, 0 * 4, 0 * 2 = 8
• [2, 0, 0, 1]：1 * 4, 0 * 4, 0 * 4, 1 * 2 = 10

這個結果確實符合我們的推論。

另外第一個維度的步長則變成4，表示第一個維度每加1，對應元素在底層記憶體裡的位置就得加4，這也與上面的結果相符。

切片張量的offset

剛剛在最後一維都是從第0個元素開始切片，如果我們改成slice最後一維的第1個元素會發生什麼事呢？

>>> x = torch.arange(6).reshape(3,1,1,2) >>> x_slice_1 = x[:,:,:,1:2] >>> x_slice_1 tensor([[[[1]]], [[[3]]], [[[5]]]])

看看storage_offset_成員變數為何？

>>>x_slice_1.storage_offset()1

跟剛剛的x_slice對照看看：

>>>x_slice.storage_offset()0

可以看到，切片張量x_slice_1的storage_offset_不再是0了。

在這種情況下，我們得在先前的換算公式上加上storage_offset_一項，變為i_n * s_n + i_c * s_c + i_h * s_h + i_w * s_w + storage_offset_。

代入實際的數字看看：

• [0, 0, 0, 0]：0 * 2, 0 * 2, 0 * 2, 0 * 1 + 1 = 1
• [1, 0, 0, 0]：1 * 2, 0 * 2, 0 * 2, 0 * 1 + 1 = 3
• [2, 0, 0, 0]：2 * 2, 0 * 2, 0 * 2, 0 * 1 + 1 = 5

接著用算出來的底層記憶體索引去存取底層記憶體：

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

正好可以得到[1, 3, 5]的值，與x_slice_1中的值相符！