torch.stack()

As shown in the figure, there are three 3x3 Tensors a, b, c,

If I want to superimpose the elements of the last dimension of these three tensors to form a new tensor

Enter d = torch.stack ((a, b, c), dim = 2)

You will find that they are superimposed in the second dimension, specifically the changes are shown in the figure below

The position of d [0] [0] is composed of [1] of a [0] [0], [10] of b [0] [0], [100] of c [0] [0], these three A new element [1,10,100] of size 3 is formed by the superposition of elements, so the dimension of d is 3 * 3 * 3. Pay attention to the dim parameter here! ! ! ! ! ! ! Although they are superimposed on the second dimension, (counting from 0 in pytorch, so that the first dimension is dim = 0, the second dimension is dim = 1), but I wrote dim in the stack function But it is dim = 2, which is the third dimension, which means that our final generation result is located in the third dimension, not that the original element is the third dimension. So there is another way of writing, write dim = -1, no matter what you are, I will specify the last dimension.

The result is the same as when I first wrote dim = 2.

 

Add another example to let you feel the meaning of this dim parameter, it specifies the element where the result of the stack is done.

dim is set to 0, 1, 2 respectively, pay attention to the observation results

When c, dim = 0, c = [a, b]

When d, dim = 1, d = [[a [0], b [0]], [a [1], b [1]]]

e, dim = 2 o'clock, e = [[[a [0] [0], b [0] [0]], [a [0] [1], b [0] [1]], [a [ 0] [2], b [0] [2]]],                                    

 [   [ a[1][0], b[1][0] ]  , [ a[1][1], b[0][1] ]  ,  [ a[1][2],b[1][2] ] ]      ]
 

 

Published 943 original articles · Like 136 · Visit 330,000+

Guess you like

Origin blog.csdn.net/weixin_36670529/article/details/105349995