Vectorizing assignment of a tensor to a slice in PyTorch

Question

I'm trying to vectorize a slice assignment of the form

for i in range(a.shape[1]):
    for j in range(a.shape[2]):
        a[:,i,j,:,i:i+b.shape[2],j:j+b.shape[3]] = b

where b itself is an array. This is because the nested Python loop is too inefficient and is taking up most of the runtime. Is there a way to do this?

For a simpler case, consider the following:

for i in range(a.shape[1]):
    a[:,i,:,i:i+b.shape[2]] = b

This is what b and a might look like:

You can see the diagonal, "sliding" structure of the resulting matrix.

@techytushar Forgot to mention `n` and `m` are just `a.shape[1]` and `a.shape[2]` (i.e. precisely the sizes of the dimensions they're indexing into). I've edited the question. — user76284, Sep 27 '19 at 20:17
Use one of the [`np.ndarray` methods](https://docs.scipy.org/doc/numpy/reference/arrays.ndarray.html), rather than doing a (Python) assign on a range of cells, which will be slower. For example, look into methods like [`np.put`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.put.html#numpy.put). It helps if you tell us more about the array you're trying to create: what structure does it have? triangular? sparse? semi-sparse? rectangular blocks of ones? How large is it? — smci, Sep 27 '19 at 20:21
@smci I've added an example for the simpler case of a single loop. — user76284, Sep 27 '19 at 20:24
How large are `b` and `a`, typically? Post us an actual code example using random data, please. Not a screenshot. — smci, Sep 27 '19 at 20:25
@smci The sizes are arbitrary, so long as `a.shape[0] == b.shape[0]` and `a.shape[2] == b.shape[1]` (for the 1D case). — user76284, Sep 27 '19 at 20:26
Ok your question title doesn't say it, but it's a (3D) tensor slice. `b` is a tensor, not a 2D array. Please read the PyTorch doc for the equivalent of TensorFlow's methods like [`tf.slice`](https://www.tensorflow.org/api_docs/python/tf/slice), rather than numpy ndarray — smci, Sep 27 '19 at 20:28
Have you read PyTorch's doc on assigning to tensor slices? What does it say? — smci, Sep 27 '19 at 20:40
I think you mean `a.shape[0]==b.shape[0]` and `a.shape[2]==b.shape[1]` and `a.shape[1]==a.shape[-1]-b.shape[-1]+1`. The `a.shape[-1]` and `b.shape[-1]` could be arbitrary size. — zihaozhihao, Sep 27 '19 at 21:23
Are `a` and `b` `np.ndarray` or something else like 'tensor`? Obviously numpy arrays can be 4 or 6d, but the arrays are something else there's little point in focusing on `numpy` methods. — hpaulj, Sep 27 '19 at 22:10
@smci Unfortunately I couldn't find anything in the docs about this. — user76284, Sep 28 '19 at 00:59

Divakar · Answer 1 · 2019-09-27T22:30:54.097

We can leverage np.lib.stride_tricks.as_strided based scikit-image's view_as_windows to get sliding windowed views into a 0s padded version of the input and being a view would be efficient on memory and performance. More info on use of as_strided based view_as_windows.

Hence, for the simpler case, it would be -

from skimage.util.shape import view_as_windows

def sliding_2D_windows(b, outshp_axis1):
    # outshp_axis1 is desired output's shape along axis=1
    n = outshp_axis1-1
    b1 = np.pad(b,((0,0),(0,0),(n,n)),'constant')
    w_shp = (1,b1.shape[1],b.shape[2]+n)
    return view_as_windows(b1,w_shp)[...,0,::-1,0,:,:]

Sample run -

In [192]: b
Out[192]: 
array([[[54, 57, 74, 77],
        [77, 19, 93, 31],
        [46, 97, 80, 98]],

       [[98, 22, 68, 75],
        [49, 97, 56, 98],
        [91, 47, 35, 87]]])

In [193]: sliding_2D_windows(b, outshp_axis1=3)
Out[193]: 
array([[[[54, 57, 74, 77,  0,  0],
         [77, 19, 93, 31,  0,  0],
         [46, 97, 80, 98,  0,  0]],

        [[ 0, 54, 57, 74, 77,  0],
         [ 0, 77, 19, 93, 31,  0],
         [ 0, 46, 97, 80, 98,  0]],

        [[ 0,  0, 54, 57, 74, 77],
         [ 0,  0, 77, 19, 93, 31],
         [ 0,  0, 46, 97, 80, 98]]],


       [[[98, 22, 68, 75,  0,  0],
         [49, 97, 56, 98,  0,  0],
         [91, 47, 35, 87,  0,  0]],

        ....

        [[ 0,  0, 98, 22, 68, 75],
         [ 0,  0, 49, 97, 56, 98],
         [ 0,  0, 91, 47, 35, 87]]]])

Is there a version of this that could work for PyTorch tensors? — user76284, Sep 28 '19 at 03:22
@user76284 Well I am not really familiar with PyTorch area. So, I am not sure. — Divakar, Sep 28 '19 at 06:13
@user76284 Does PyTorch have anything similar to NumPy's [as_strided](https://docs.scipy.org/doc/numpy-1.15.0/reference/generated/numpy.lib.stride_tricks.as_strided.html)? If so, should be straight-forward then. — Divakar, Sep 28 '19 at 16:43

zihaozhihao · Answer 2 · 2019-09-27T21:20:25.690

Assuming b has a shape (2,3,x1), and a has a shape (2,x2-x1+1,3,x2). In your screenshot, we can infer that x1=4, x2=6.

import numpy as np

b_shape = (2,3,4)
a_shape = (2,3,3,6)

b = np.arange(1,25).reshape(b_shape)
#array([[[ 1,  2,  3,  4],
#        [ 5,  6,  7,  8],
#        [ 9, 10, 11, 12]],
#
#       [[13, 14, 15, 16],
#        [17, 18, 19, 20],
#        [21, 22, 23, 24]]])

c = np.pad(b, (*[(0,0) for _ in range(len(b_shape[:-1]))], (0,a_shape[-1]-b_shape[-1])), 'constant')

#array([[[ 1,  2,  3,  4,  0,  0],
#        [ 5,  6,  7,  8,  0,  0],
#        [ 9, 10, 11, 12,  0,  0]],
#
#       [[13, 14, 15, 16,  0,  0],
#        [17, 18, 19, 20,  0,  0],
#        [21, 22, 23, 24,  0,  0]]])

a = np.stack([np.roll(c, shift=i) for i in range(a_shape[-1]-b_shape[-1]+1)], axis=1)

# array([[[[ 1,  2,  3,  4,  0,  0],
#          [ 5,  6,  7,  8,  0,  0],
#          [ 9, 10, 11, 12,  0,  0]],

#         [[ 0,  1,  2,  3,  4,  0],
#          [ 0,  5,  6,  7,  8,  0],
#          [ 0,  9, 10, 11, 12,  0]],

#         [[ 0,  0,  1,  2,  3,  4],
#          [ 0,  0,  5,  6,  7,  8],
#          [ 0,  0,  9, 10, 11, 12]]],


#        [[[13, 14, 15, 16,  0,  0],
#          [17, 18, 19, 20,  0,  0],
#          [21, 22, 23, 24,  0,  0]],

#         [[ 0, 13, 14, 15, 16,  0],
#          [ 0, 17, 18, 19, 20,  0],
#          [ 0, 21, 22, 23, 24,  0]],

#         [[ 0,  0, 13, 14, 15, 16],
#          [ 0,  0, 17, 18, 19, 20],
#          [ 0,  0, 21, 22, 23, 24]]]])

Yes, the only for-loop is for different shift of `np.roll`. Probably there are some other better methods. :) — zihaozhihao, Sep 27 '19 at 21:32

Vectorizing assignment of a tensor to a slice in PyTorch

2 Answers2