如何在PyTorch中结合切片赋值、掩码赋值与广播机制？

问题场景 #

给定以下张量与掩码：

import torch

x = torch.zeros(2, 3, 4, 6)
mask = torch.tensor([[True, True, False], [True, False, True]])
y = torch.rand(2, 3, 1, 3)

需求是：

• 仅对mask中为True的位置（对应x的第0、1维度）进行赋值；
• y的第2维度（长度1）需要广播到x的第2维度（长度4）；
• 仅用y的第3维度前3个元素，覆盖x的第3维度前3个元素。

错误方法分析 #

1. 直接混合掩码与切片报错

   x[mask, :, :3] = y[mask]
   

   报错原因：布尔索引x[mask]会将x的前两维压缩为一维（长度为mask中True的数量，即4），得到形状(4,4,6)的张量。此时x[mask, :, :3]的实际形状是(4,4,3)，而y[mask]的形状是(4,1,3)，虽理论可广播，但PyTorch的索引解析逻辑可能导致形状匹配错误，且写法逻辑模糊。

2. 链式索引赋值不生效

   x[mask][:, :, :3] = y[mask]
   

   无报错但赋值无效的原因：x[mask]返回的是原张量的拷贝而非视图，对拷贝的修改不会同步到原张量x。

高效解决方案（原地赋值，无大中间张量） #

使用高级索引定位mask为True的位置，结合切片直接原地修改：

# 获取mask中True元素的第0、1维度索引
idx0, idx1 = torch.where(mask)

# 直接对目标位置赋值：
# x[idx0, idx1, :, :3] 定位到所有需要修改的区域，形状(4,4,3)
# y[idx0, idx1].expand(...) 将y的第2维度广播到与x的第2维度一致，形状(4,4,3)
x[idx0, idx1, :, :3] = y[idx0, idx1].expand(-1, x.shape[2], -1)

原理说明 #

• torch.where(mask)返回两个一维张量idx0和idx1，存储所有mask为True的位置坐标，内存占用极小；
• x[idx0, idx1, :, :3]精准定位到x中需要修改的区域：每个(idx0[k], idx1[k])对应x的第0、1维度位置，第2维度全选，第3维度取前3个元素；
• y[idx0, idx1].expand(-1, x.shape[2], -1)将y的第2维度从1扩展为4（与x的第2维度长度一致），实现广播，最终形状与目标区域完全匹配，直接完成原地赋值。

这种方法无需生成大尺寸中间张量，适合神经网络中大型张量的高效修改。

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