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]])
Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐