PyTorch如何沿轴对不同大小张量索引实现无循环批量赋值

可以实现，无需任何循环，仅通过PyTorch的广播机制和布尔索引即可完成，和你给出的循环实现效果完全一致。

实现思路 #

你的循环逻辑本质是对张量T的每个位置(i,j)，将最后两个维度中高度小于x[i]、宽度小于y[j]的区域赋值为-1，我们可以通过广播生成全局布尔掩码，直接批量赋值：

• 首先生成最后两个维度的索引序列
• 把x、y调整维度后和索引序列做广播比较，生成全局掩码
• 用掩码直接对张量赋值

完整代码示例 #

import torch

# 初始化示例数据，可替换为你自己的张量
T = torch.zeros(8, 8, 4, 4)
x = torch.tensor([1, 2, 3, 1, 2, 3, 1, 2])
y = torch.tensor([2, 3, 2, 3, 2, 3, 2, 3])

# 生成最后两个维度的索引，注意和T保持相同设备避免报错
idx_h = torch.arange(4, device=T.device)
idx_w = torch.arange(4, device=T.device)

# 广播生成全局布尔掩码
mask = (idx_h < x.view(8, 1, 1, 1)) & (idx_w < y.view(1, 8, 1, 1))

# 批量赋值
T[mask] = -1

效果验证 #

你可以直接对比循环实现和上述实现的结果，二者完全相等：

# 循环实现的对照结果T1
T1 = torch.zeros(8,8,4,4)
for i in range(8):
    for j in range(8):
        T1[i, j, :x[i], :y[j]] = -1

# 对比结果
print(torch.allclose(T, T1)) # 输出为True

内容的提问来源于stack exchange，提问作者Mustafa Süve