PyTorch的slice究竟做了什麼事 stride和stride_offset的妙用Tensor ViewsTensor的本質切片張量的步長切片張量的offset直接上代碼首先建立一個形狀為3x1x1x2的張量xtorch.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_slicex_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(dtypetorch.int64,devicecpu)of size6]x_slice.storage()55512345[torch.storage.TypedStorage(dtypetorch.int64,devicecpu)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 offseti_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的張量xtorch.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中的值相符