View matrix as matrix of submatrices

1

For my last project i need to use a matrix as matrix of “vectorized” submatrices.
For example give the matrix
A = [[1 2 3]
[4 5 6]
[7 8 9]]

The matrix of submatrices with dimension 2x2 is

M = [[ 1 2 4 5]
[ 2 3 5 6]
[ 4 5 7 8]
[ 5 6 8 9]]

Where each row is one of the 2x2 matrix composing A.

For now i’m initializing M with zeros and using a for-loop slicing A to create M but it is time-comsuming ( O (n*m) ).
I was wondering if there exist a more pytorch-y way to do it

This is my code

>>> getSubMatrix = lambda m,i,j : m[i:i+2,j:j+2].contiguous().view(-1,2*2)[0]             
>>>
>>> A = torch.randn(3,3)
>>> M = torch.zeros(4, 2*2)
>>> k = 0
>>> for i in range(0, 2):
...     for j in range(0, 2):
...         M[k] = getSubMatrix(A, i, j)
...         k = k + 1
...
>>> A
tensor([[ 0.4257,  0.7940, -1.2986],
        [ 0.3243, -1.3812, -1.0442],
        [-0.3122,  1.2312,  0.9811]])
>>> M
tensor([[ 0.4257,  0.7940,  0.3243, -1.3812],                                             
        [ 0.7940, -1.2986, -1.3812, -1.0442],                                             
        [ 0.3243, -1.3812, -0.3122,  1.2312],                                             
        [-1.3812, -1.0442,  1.2312,  0.9811]])                                            
>>>

Thanks a lot :blush:

2

tensor.unfold might be what you’re looking for:

A = torch.tensor([[ 0.4257,  0.7940, -1.2986],
                  [ 0.3243, -1.3812, -1.0442],
                  [-0.3122,  1.2312,  0.9811]])

M = A.unfold(dimension=0, size=2, step=1).unfold(dimension=1, size=2, step=1)
M = M.contiguous().view(4, 4)
print(M)
> tensor([[ 0.4257,  0.7940,  0.3243, -1.3812],
          [ 0.7940, -1.2986, -1.3812, -1.0442],
          [ 0.3243, -1.3812, -0.3122,  1.2312],
          [-1.3812, -1.0442,  1.2312,  0.9811]])