torch.rot90 的理解与使用

官方定义

torch.rot90(input, k, dims)

• input (Tensor) – the input tensor.
• k (int) – number of times to rotate
• dims (a list or tuple) – axis to rotate

Rotate a n-D tensor by 90 degrees in the plane specified by dims axis. Rotation direction is from the first towards the second axis if k > 0, and from the second towards the first for k < 0.

沿着dims轴90旋转张量input，当k>0旋转方向由dims[0]指向dims[1]，当k<0时 旋转方向由dims[1]指向dims[0]，旋转角度为：90$\times$k.

解读

通用的轴方向定义为： axis=0指向下方，axis=1指向右方，如下图

 ------------> axis=1
 | 
 |
 |
 \/
 axis=0

所以，如果固定dims=[0, 1]，则k>0表示逆时针旋转，k<0表示顺时针旋转。具体例子如下：

>>> x =torch.arange(6).view(3,2)
>>> x
tensor([[0, 1],
		[2, 3],
		[4, 5]])
		
>>> torch.rot90(x, 1, [0, 1])   # 由轴0转向轴1（逆时针）
tensor([[1, 3, 5],
		[0, 2, 4]])
		
>>> torch.rot90(x, -1, [0, 1])  # 由轴1转向轴0（顺时针）
tensor([[4, 2, 0],
		[5, 3, 1]])