PyTorch張量拼接、切分、索引的實現(xiàn)

更新時間：2024年03月21日 10:21:40 作者：timerring

在學習深度學習的過程中,遇到的第一個概念就是張量,張量在pytorch中的計算十分重要,本文主要介紹了PyTorch張量拼接、切分、索引的實現(xiàn),具有一定的參考價值,感興趣的可以了解一下

一、張量拼接與切分

1.1 torch.cat

功能：將張量按維度dim 進行拼接

tensors : 張量序列
dim：要拼接的維度

 t = torch.ones((2, 3))

    t_0 = torch.cat([t, t], dim=0)
    t_1 = torch.cat([t, t, t], dim=1)

    print("t_0:{} shape:{}\nt_1:{} shape:{}".format(t_0, t_0.shape, t_1, t_1.shape))

t_0:tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]]) shape:torch.Size([4, 3])
t_1:tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1., 1.]]) shape:torch.Size([2, 9])

（2，3） -> （2，6）

這里的dim維度與axis相同，0代表列，1代表行。

1.2 torch.stack

功能：在新創(chuàng)建的維度 dim 上進行拼接（會拓寬原有的張量維度）

tensors：張量序列
dim：要拼接的維度

t = torch.ones((2, 3))
t_stack = torch.stack([t, t, t], dim=2)
print("\nt_stack:{} shape:{}".format(t_stack, t_stack.shape))

可見，它在新的維度上進行了拼接。

參數(shù)[t, t, t]的意思就是在第n個維度上拼接成這個樣子。

t_stack:tensor([[[1., 1., 1.],
         [1., 1., 1.],
         [1., 1., 1.]],

        [[1., 1., 1.],
         [1., 1., 1.],
         [1., 1., 1.]]]) shape:torch.Size([2, 3, 3])
# 在第二維度上進行了拼接
Process finished with exit code 0

1.3 torch.chunk

功能：將張量按維度 dim 進行平均切分

返回值：張量列表

注意事項：若不能整除，最后一份張量小于其他張量。

input : 要切分的張量
chunks 要切分的份數(shù)
dim 要切分的維度

    # cut into 3
    a = torch.ones((2, 7))  # 7
    list_of_tensors = torch.chunk(a, dim=1, chunks=3)   # 3

    for idx, t in enumerate(list_of_tensors):
        print("第{}個張量：{}, shape is {}".format(idx+1, t, t.shape))

可知，切分是7/3向上取整，每份是3，最后剩下的維度直接輸出即可。

第1個張量：tensor([[1., 1., 1.],
        [1., 1., 1.]]), shape is torch.Size([2, 3])
第2個張量：tensor([[1., 1., 1.],
        [1., 1., 1.]]), shape is torch.Size([2, 3])
第3個張量：tensor([[1.],
        [1.]]), shape is torch.Size([2, 1])

1.4 torch.split

torch.split(Tensor, split_size_or_sections, dim)

功能：將張量按維度 dim 進行切分

返回值：張量列表

tensor : 要切分的張量
split_size_or_sections 為 int 時，表示
每一份的長度；為 list 時，按 list 元素切分
dim 要切分的維度

    t = torch.ones((2, 5))

    list_of_tensors = torch.split(t, [2, 1, 1], dim=1)  # [2 , 1, 2]
    for idx, t in enumerate(list_of_tensors):
        print("第{}個張量：{}, shape is {}".format(idx+1, t, t.shape))

是按照指定長度list進行切分的。注意list中長度總和必須為原張量在改維度的大小，不然會報錯。

第1個張量：tensor([[1., 1., 1.],
        [1., 1., 1.]]), shape is torch.Size([2, 3])
第2個張量：tensor([[1., 1., 1.],
        [1., 1., 1.]]), shape is torch.Size([2, 3])
第3個張量：tensor([[1.],
        [1.]]), shape is torch.Size([2, 1])

二、張量索引

2.1 torch.index_select

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

功能：在維度dim 上，按 index 索引數(shù)據(jù)

返回值：依index 索引數(shù)據(jù)拼接的張量

input : 要索引的張量
dim 要索引的維度
index 要索引數(shù)據(jù)的序號

    t = torch.randint(0, 9, size=(3, 3))
    idx = torch.tensor([0, 2], dtype=torch.long)    # if float will report an error
    t_select = torch.index_select(t, dim=0, index=idx)
    print(idx)
    print("t:\n{}\nt_select:\n{}".format(t, t_select))

可見idx是一個存儲序號的張量，而torch.index_select通過該張量索引原tensor并且拼接返回。

tensor([0, 2])
t:
tensor([[4, 5, 0],
        [5, 7, 1],
        [2, 5, 8]])
t_select:
tensor([[4, 5, 0],
        [2, 5, 8]])

2.2 torch.masked_select

功能：按mask 中的 True 進行索引

返回值：一維張量(無法確定true的個數(shù)，因此也就無法顯示原來的形狀，因此這里返回一維張量)

input : 要索引的張量
mask 與 input 同形狀的布爾類型張量

    t = torch.randint(0, 9, size=(3, 3))
    mask = t.le(5)  # ge is mean greater than or equal/   gt: greater than  le  lt
    t_select = torch.masked_select(t, mask)
    print("t:\n{}\nmask:\n{}\nt_select:\n{} ".format(t, mask, t_select))

通過掩碼來索引。

tensor([[4, 5, 0],
        [5, 7, 1],
        [2, 5, 8]])
mask:
tensor([[ True,  True,  True],
        [ True, False,  True],
        [ True,  True, False]])
t_select:
tensor([4, 5, 0, 5, 1, 2, 5]) 

Process finished with exit code 0

到此這篇關(guān)于PyTorch張量拼接、切分、索引的實現(xiàn)的文章就介紹到這了,更多相關(guān)PyTorch張量拼接切分索引內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: