【pytorch select】索引函数

select(dim, index)：第一个参数为索引的维度，第二个参数为索引的维度的序列号

import torch
a = torch.rand((3, 4))
print(a)
>>> tensor([[0.8664, 0.9759, 0.3063, 0.0686],
	        [0.6778, 0.0574, 0.3194, 0.4253],
	        [0.5045, 0.8318, 0.1745, 0.3150]])
print(a.select(dim=1, index=1)) # 取第1个维度中索引为1的值
>>> tensor([0.9759, 0.0574, 0.8318])

使用copy_赋值修改张量

import torch
a = torch.rand((3, 4)
print(a)
>>> tensor([[0.6115, 0.2551, 0.8714, 0.3236],
	        [0.3369, 0.4372, 0.2083, 0.4733],
	        [0.0046, 0.0981, 0.9148, 0.7852]])
y = torch.tensor([3, 2, 1])
a.select(1, 1).copy_(y.data) # 将第1个维度中索引为1的张量替换为y
print(a)
>>> tensor([[0.6115, 3.0000, 0.8714, 0.3236],
	        [0.3369, 2.0000, 0.2083, 0.4733],
	        [0.0046, 1.0000, 0.9148, 0.7852]])