torch.index_select函数测试

最新推荐文章于 2025-07-21 00:28:47 发布

原创最新推荐文章于 2025-07-21 00:28:47 发布 · 237 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#矩阵 #pytorch #线性代数 #深度学习 #人工智能

文章目录

1. description
2. pytorch

1. description

index_select 函数的方法是可以按照序号的方式随机抽取指定维度的向量，比如按照行进行采样，按照列进行采样。
在这里插入图片描述

2. pytorch

code

import torch
import torch.nn as nn

torch.manual_seed(23224)

torch.set_printoptions(precision=3, sci_mode=False, threshold=torch.inf)

if __name__ == "__main__":
    run_code = 0
    batch_size = 4
    model_dim = 6
    a_total = batch_size * model_dim
    a_matrx = torch.arange(a_total).reshape(batch_size, model_dim)
    a_index = torch.randperm(model_dim)
    print(f"a_matrx=\n{a_matrx}")
    print(f"a_index=\n{a_index}")
    a_column = torch.index_select(input=a_matrx, dim=1, index=a_index)
    print(f"a_column=\n{a_column}")

result

a_matrx=
tensor([[ 0,  1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10, 11],
        [12, 13, 14, 15, 16, 17],
        [18, 19, 20, 21, 22, 23]])
a_index=
tensor([2, 3, 5, 4, 0, 1])
a_column=
tensor([[ 2,  3,  5,  4,  0,  1],
        [ 8,  9, 11, 10,  6,  7],
        [14, 15, 17, 16, 12, 13],
        [20, 21, 23, 22, 18, 19]])