torch as_strided函数详解

假设有一个矩阵X

$\left[\begin{array}{ccc}1& 2& 3\\ 4& 5& 6\\ 7& 8& 9\end{array}\right]$
如果卷积核大小是2*2,stride为1,那么就需要把矩阵转换为如下4个小矩阵的新矩阵
$\left[\begin{array}{c}\left[\begin{array}{cc}1& 2\\ 4& 5\end{array}\right]\left[\begin{array}{cc}2& 3\\ 5& 6\end{array}\right]\\ \left[\begin{array}{cc}4& 5\\ 6& 7\end{array}\right]\left[\begin{array}{cc}5& 6\\ 7& 8\end{array}\right]\end{array}\right]$
而as_strided就是进行这种操作，查看一下对于as_strided参数的描述

numpy.lib.stride_tricks.as_strided(x, shape=None, strides=None, subok=False, writeable=True)

x就是我们要分割的矩阵，可以当做是一个蓝图，shape，strides都是新矩阵的属性，也就是说这个函数按照给定的shape和strides来划分x，返回一个新的矩阵，最后两个参数不讨论，可以看numpy官方手册的描述
shape很好理解，用的也比较多
比如矩阵X的shape：(3, 3)；矩阵A的shape：(2,2,2,2)
至于strides属性，先看下面的代码：

a = np.arange(9, dtype=np.int32).reshape(3,3)
print(a)
'''
[[0 1 2]
 [3 4 5]
 [6 7 8]]
 '''
print(a.strides)
'''
(12, 4)
'''

这里12为从a[0][0]到a[1][0]的维度，4为从a[0][0]到a[0][1]的维度
因为np.int32为4B，从a[0][0]到a[1][0]即从数值0到数值3跨越了3个数值，所以strides=$3\ast 4B=12B$。a[0][0]到a[0][1]即从数值0到数值1跨越了1个数值，所以strides=$1\ast 4B=4B$。
注意：strides的单位始终为字节。
同类的例子如下：

b = np.arange(16,dtype=np.float).reshape(4,4)
print(b.strides)
#(32,8)
c = np.arange(27,dtype=np.float).reshape(3,3,3)
print(c.strides)
#(36,12,4)

当数值为0~15的时候
$\left[\begin{array}{cccc}1& 2& 3& 4\\ 5& 6& 7& 8\\ 9& 10& 11& 12\\ 13& 14& 15& 16\end{array}\right]$
由于np.float一次占据8B，所以b.strides[0]为由[0][0] = 1到[1][0] = 5，跨越了$4\ast 8B=32B$，b.strides[1]为由[0][0] = 1到[0][1] = 2，跨越了$1\ast 8B=8B$，所以b.strides = (32,8)
同理矩阵c的内容为
$\left[\begin{array}{c}\left[\begin{array}{ccc}1& 2& 3\\ 4& 5& 6\\ 7& 8& 9\end{array}\right]\\ \left[\begin{array}{ccc}10& 11& 12\\ 13& 14& 15\\ 16& 17& 18\end{array}\right]\\ \left[\begin{array}{ccc}19& 20& 21\\ 22& 23& 24\\ 25& 26& 27\end{array}\right]\end{array}\right]$
这里如果从a[0][0][0]到a[1][0][0]，即从1到10的位置，一个float类型的数据占据4B，所以$c.strides[0]=4B\ast 9=36B$。如果从a[0][0][0]到a[0][1][0]，即从1到4的位置，这里$c.strides[1]=4B\ast 3=12B$。如果从a[0][0][0]到a[0][0][1]，即从1到2的过程，这里$c.strides[2]=4B$，综上c.strides = (36,12,4)。
接下来看从X得到A的as_strided函数

A = as_strided(X,shape=(2,2,2,2),strides)

原向量X的内容为
$\left[\begin{array}{ccc}1& 2& 3\\ 4& 5& 6\\ 7& 8& 9\end{array}\right]$
结果为
$\left[\begin{array}{c}\left[\begin{array}{cc}1& 2\\ 4& 5\end{array}\right]\left[\begin{array}{cc}2& 3\\ 5& 6\end{array}\right]\\ \left[\begin{array}{cc}4& 5\\ 7& 8\end{array}\right]\left[\begin{array}{cc}5& 6\\ 8& 9\end{array}\right]\end{array}\right]$
这里的matrix[0][0][0][0]->matrix[1][0][0][0]:1->4,
search original matrix
$\left[\begin{array}{ccc}1& 2& 3\\ 4& 5& 6\\ 7& 8& 9\end{array}\right]$
$strides[0]=4B\ast 3=12B$
matrix[0][0][0][0]->matrix[0][1][0][0]:1->2,
search original matrix
$\left[\begin{array}{ccc}1& 2& 3\\ 4& 5& 6\\ 7& 8& 9\end{array}\right]$
$strides[1]=4B\ast 1=4B$
matrix[0][0][0][0]->matrix[0][0][1][0]:1->4
search original matrix
$\left[\begin{array}{ccc}1& 2& 3\\ 4& 5& 6\\ 7& 8& 9\end{array}\right]$
$strides[2]=4B\ast 3=12B$
matrix[0][0][0][0]->matrix[0][0][0][1]:1->2
search original matrix
$\left[\begin{array}{ccc}1& 2& 3\\ 4& 5& 6\\ 7& 8& 9\end{array}\right]$
$strides[2]=4B\ast 1=4B$
final output

X = np.arange(9, dtype=np.int32).reshape(3,3)
print(X)
'''
[[0 1 2]
 [3 4 5]
 [6 7 8]]
'''
A = np.lib.stride_tricks.as_strided(X, shape=(2,2,2,2), strides=(12,4,12,4))
print(A)
'''
[[[[0 1]
   [3 4]]

  [[1 2]
   [4 5]]]


 [[[3 4]
   [6 7]]

  [[4 5]
   [7 8]]]]
'''