如何无循环将张量b赋值给张量a的多切片以得到目标张量c？

当然可以做到！不用写循环，我们可以借助PyTorch的高级索引特性，一次性把张量b赋值到a的对应位置，直接得到目标张量c。

具体实现步骤 #

首先先初始化我们的初始张量：

import torch

# 初始化3x4的零张量a
a = torch.zeros((3, 4), dtype=torch.int)
# 定义张量b
b = torch.tensor([1, 2])

接下来，我们需要生成所有需要赋值的位置的行索引和列索引：

• 行索引：每个行需要被赋值2次，也就是[0, 0, 1, 1, 2, 2]
• 列索引：对应每个行的起始列依次后移，也就是[0, 1, 1, 2, 2, 3]

用PyTorch的张量操作可以轻松生成这些索引：

# 生成行索引：每个行索引重复2次
row_indices = torch.arange(3).repeat_interleave(2)
# 生成列索引：先把行索引重复成[0,0,1,1,2,2]，再加上[0,1,0,1,0,1]得到目标列位置
col_indices = torch.arange(3).repeat(2) + torch.tensor([0, 1]).repeat(3)

最后直接通过索引赋值，同时把b重复3次来匹配索引的长度：

a[row_indices, col_indices] = b.repeat(3)

现在打印a就能得到目标张量：

print(a)
# 输出：
# tensor([[1, 2, 0, 0],
#         [0, 1, 2, 0],
#         [0, 0, 1, 2]])

另一种更简洁的方式 #

如果你觉得生成索引有点麻烦，还可以利用torch.Tensor.scatter_方法来实现，原理类似：

a = torch.zeros((3, 4), dtype=torch.int)
# 构造要填充的数值（b重复3次）
values = b.repeat(3)
# 构造列索引矩阵
col_idx = torch.tensor([[0,1], [1,2], [2,3]])
# 使用scatter_赋值，dim=1表示按列维度填充
a.scatter_(dim=1, index=col_idx, src=values.reshape(3,2))

运行后同样能得到目标张量c。

这两种方法都完全不需要写循环，都是利用PyTorch的向量化操作来实现，效率也更高哦。

内容的提问来源于stack exchange，提问作者user76284