Python函数

1. torch.gather()

torch.gather()函数用于根据指定索引从输入的张量中收集元素。

调用形式为：

torch.gather(input, dim, index, out=None)

input：输入张量

dim：指定收集操作沿着哪个维度进行

index：一个包含了从input中收集元素的索引的张量

out：可选参数，指定输出张量的位置

其中，input与index的维度必须一致，shape可以不同；output与index的维度与shape均一致。

input_tensor = torch.tensor([[1, 2, 3],
                             [4, 5, 6],
                             [7, 8, 9]])


indices = torch.tensor([[0, 1],
                        [2, 1]])

result = torch.gather(input_tensor, dim=0, index=indices)

print(result)

代码输出为：

tensor([[1, 5],
        [7, 5]])

output[i][j] = input [index[i][j]] [j]  # if dim = 0

output[i][j] = input [i] [index[i][j]]  # if dim = 1

若在dim = 0的维度上进行收集，即在二维矩阵上行向量上进行收集，index矩阵中的值代表input行向量的索引，即用index[i][j]替换input[i][j]中的i值，j值不变。实例中，index[0][1] = 1，故        output [0][1] = input [index[0][1]] [1] = input[1][1] = 5.

dim = 1时同理。

扩展到三维：

input_tensor = torch.tensor([
    [[0, 1, 2],
     [3, 4, 5]],

    [[6, 7, 8],
     [9, 10, 11]]
])

indices = torch.tensor([
    [[0, 1],
     [1, 0]],

    [[1, 1],
     [0, 1]]
])

result_test = input_tensor.gather(dim=1, index=indices)

输出为：

tensor([[[ 0,  4],
         [ 3,  1]],

        [[ 9, 10],
         [ 6, 10]]])

在二维张量中，dim = 0是二维矩阵行向量，dim = 1是列向量；

扩展到三维张量，将一个二维矩阵看作是一个batch，dim = 0是每个batch的索引。

output[i][j][k] = input [index[i][j][k]] [j] [k]   # if dim = 0

output[i][j][k] = input [i] [index[i][j][k]] [k]   # if dim = 1

output[i][j][k] = input [i] [j] [index[i][j][k]]    # if dim = 2

在实例中，若dim = 1，index[0][1][0] = 1，故output[0][1][0] = input[0] [index[0][1][0]] [0] =             input[0][1][0] = 3

2. torch.nn.functional.pad()

调用形式：

torch.nn.functional.pad(input, pad, mode='constant', value=0)

pad：扩充维度，需要输入一个元组，元素个数 ≤ 2 * input维数，因为在每一个维度上，可以选择在该维度的两侧进行扩充。且扩充时，从input的最后一个维度开始，向前一个维度方向扩充。同维度上，先左后右，先上后下。

mode：扩充方式，constant、reflect、replicate，分别表示常量、反射、复制。

value：指定扩充操作时的补充值，只在mode = constant时有效。

3.torch.stack()

调用形式：

torch.stack(tensor, dims=0, out=None)

将tensor在dims的维度上进行堆叠。要求：输入的张量形状一致。

 举例当输入的两个张量都是二维时：

import math
import torch
from torch import nn
from torch.nn import init
from torch.nn import functional as F

alpha = torch.tensor([
            [1, 2, 3],
            [4, 5, 6]
                        ])
beta = torch.tensor([
            [7, 8, 9],
            [10, 11, 12]
                        ])
theta_0 = torch.stack([alpha, beta], dim=0)
print('0:', theta_0, theta_0.shape)
theta_1 = torch.stack([alpha, beta], dim=1)
print('1:', theta_1, theta_1.shape)
theta_2 = torch.stack([alpha, beta], dim=2)
print('2:', theta_2, theta_2.shape)

当dim=0时，是在channel的维度上进行堆叠；

当dim=1时，是在行维度上进行堆叠；

当dim=2时，是在列维度上进行堆叠。

输出为：

0: tensor([[[ 1,  2,  3],
         [ 4,  5,  6]],

        [[ 7,  8,  9],
         [10, 11, 12]]]) 
torch.Size([2, 2, 3])

1: tensor([[[ 1,  2,  3],
         [ 7,  8,  9]],

        [[ 4,  5,  6],
         [10, 11, 12]]]) 
torch.Size([2, 2, 3])

2: tensor([[[ 1,  7],
         [ 2,  8],
         [ 3,  9]],

        [[ 4, 10],
         [ 5, 11],
         [ 6, 12]]]) 
torch.Size([2, 3, 2])

4.torch.permute()

调用形式：

output_tensor = torch.permute(input_tensor, (1, 2, 0))

如果输入的张量有三个维度，该函数将原张量维度置换为（1， 2， 0）

alpha = torch.tensor([
            [1, 2, 3],
            [4, 5, 6]
                        ])
beta = torch.tensor([
            [7, 8, 9],
            [10, 11, 12]
                        ])
theta = torch.stack([alpha, beta], dim=0)
theta_new = torch.permute(theta, (1, 2, 0))
print(theta_new)

输出为：

tensor([[[ 1,  7],
         [ 2,  8],
         [ 3,  9]],

        [[ 4, 10],
         [ 5, 11],
         [ 6, 12]]])