pytorch框架常用函数总结[自用-持续更新]

torch.permute(dims)

调换维度，dims必须match原始tensor的维度

>>>x = torch.randn(2,3,4)
>>>x.size()
Out[5]: torch.Size([2, 3, 4])
>>>x.permute(2,0,1).size()
Out[6]: torch.Size([4, 2, 3])

torch.cat((A,B),axis)

对A, B两个tensor进行拼接，参数axis指定拼接的方式。axis=0为按行拼接；axis=1为按列拼接

>>>a = torch.ones(2,3)
Out[42]: 
tensor([[1., 1., 1.],
        [1., 1., 1.]])
>>>b = torch.Tensor([[1,2,3],[4,5,6]])
Out[44]: 
tensor([[1., 2., 3.],
        [4., 5., 6.]])
>>>torch.cat((a,b),dim=0)
Out[45]: 
tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 2., 3.],
        [4., 5., 6.]])
>>>torch.cat((a,b),dim=1)
Out[46]: 
tensor([[1., 1., 1., 1., 2., 3.],
        [1., 1., 1., 4., 5., 6.]])

torch.stack()

沿着新的维度对张量序列进行拼接，所有张量形状要相同，新增维度的大小为len(inputs)

>>>a
Out[9]: 
tensor([[1., 2., 3.],
        [4., 5., 6.],
        [7., 8., 9.]])
>>>b = a*10
Out[11]: 
tensor([[10., 20., 30.],
        [40., 50., 60.],
        [70., 80., 90.]])
>>>c = a*100
Out[13]: 
tensor([[100., 200., 300.],
        [400., 500., 600.],
        [700., 800., 900.]])
#a.size()
#Out[17]: torch.Size([3, 3])

>>>inputs = [a,b]
>>>torch.stack(inputs,dim=0)
Out[19]: 
tensor([[[ 1.,  2.,  3.],
         [ 4.,  5.,  6.],
         [ 7.,  8.,  9.]],
        [[10., 20., 30.],
         [40., 50., 60.],
         [70., 80., 90.]]])
>>>torch.stack(inputs,dim=0).shape
Out[20]: torch.Size([2, 3, 3])


>>>torch.stack(inputs,dim=1)
Out[21]: 
tensor([[[ 1.,  2.,  3.],
         [10., 20., 30.]],
        [[ 4.,  5.,  6.],
         [40., 50., 60.]],
        [[ 7.,  8.,  9.],
         [70., 80., 90.]]])
>>>torch.stack(inputs,dim=1).shape
Out[22]: torch.Size([3, 2, 3])


>>>torch.stack(inputs,dim=2)
Out[23]: 
tensor([[[ 1., 10.],
         [ 2., 20.],
         [ 3., 30.]],
        [[ 4., 40.],
         [ 5., 50.],
         [ 6., 60.]],
        [[ 7., 70.],
         [ 8., 80.],
         [ 9., 90.]]])
>>>torch.stack(inputs,dim=2).shape
Out[24]: torch.Size([3, 3, 2])

torch.multinomial(inputs, num_samples, replacement=False, out=None)

对inputs的每一行进行num_samples次采样，返回LongTensor采样inputs下标，replacement指的是取样时是否是有放回的取样，True是有放回，False是无放回

>>>weights = torch.Tensor([[0,10,3,0],[1,5,0,0]])
>>>torch.multinomial(weights,2)
Out[37]: 
tensor([[1, 2],
        [0, 1]])
>>>torch.multinomial(weights,1)
Out[38]: 
tensor([[1],
        [1]])
>>>torch.multinomial(weights,5,replacement=True)
Out[39]: 
tensor([[1, 1, 1, 2, 1],
        [0, 1, 1, 1, 1]])
>>>torch.multinomial(weights,5)
Traceback (most recent call last):
  File "J:\Anaconda3\lib\site-packages\IPython\core\interactiveshell.py", line 3331, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-40-a9bae4de2fc5>", line 1, in <module>
    torch.multinomial(weights,5)
RuntimeError: cannot sample n_sample > prob_dist.size(-1) samples without replacement
#replacement=False情况下不会取样到weight为0的样本index

torch.ge(input, other, out=None) -> Tensor

第二个参数other可以为一个数或与第一个参数input相同形状和类型的张量，逐元素比较input和other，if input>=other true else false

>>>torch.ge(torch.Tensor([[1, 2], [3, 4]]),\
			 torch.Tensor([[1, 1], [4, 4]]))
Out[47]: 
tensor([[ True,  True],
        [False,  True]])

torch.ne(input, other, out=Tensor) -> Tensor

第二个参数可以为一个数或与第一个参数相同形状和类型的张量，逐元素比较input和other， if input != other true else false

>>>torch.ne(torch.Tensor([[1, 2], [3, 4]]), torch.Tensor([[1, 1], [4, 4]]))
Out[48]: 
tensor([[False,  True],
        [ True, False]])

numpy.cumsum()

按照axis对不同维度进行累加求和

>>>x = np.array([[3, 1, 2],[5,6,7]])
Out[4]: 
array([[3, 1, 2],
       [5, 6, 7]])
>>>x.shape
Out[5]: (2, 3)
>>>x[:,::-1]
Out[6]: 
array([[2, 1, 3],
       [7, 6, 5]])
>>>x[:,::-1].cumsum(axis=1)
Out[7]: 
array([[ 2,  3,  6],
       [ 7, 13, 18]], dtype=int32)
>>>x[:,::-1].cumsum(axis=1)[:,::-1]
Out[8]: 
array([[ 6,  3,  2],
       [18, 13,  7]], dtype=int32)