PyTorch中结合高级索引与赋值的高效实现问题

你遇到的核心问题是提取出的[batch_size, width]张量无法直接广播到[batch_size, x, y, width]，只需要给提取后的张量增加两个单维度，就能触发PyTorch的自动广播，无需显式循环或repeat操作。

优化代码实现 #

import torch

# 假设已定义变量：
# x: shape (batch, x_dim, y_dim, width)
# goals: shape (batch_size, 2)，存储每个batch对应的(y, x)坐标
batch_size = x.shape[0]
x_dim, y_dim = x.shape[1], x.shape[2]

# 提取每个batch目标位置的值
goal_y = goals[:, 1]
goal_x = goals[:, 0]
target_vals = x[torch.arange(batch_size), goal_y, goal_x, :]

# 增加两个单维度，将形状从(batch, width)转为(batch, 1, 1, width)
# 两种等价写法任选其一：
target_vals = target_vals.unsqueeze(1).unsqueeze(1)
# 或更简洁的索引写法：target_vals = target_vals[:, None, None, :]

# 直接赋值，广播会自动将1扩展为x_dim和y_dim
g = x.clone()
g[:, :, :, :] = target_vals

为什么这方法更优？ #

• 无需复制数据：unsqueeze只是改变张量的形状视图，不会额外占用内存；而repeat会实际复制数据，内存开销更大。
• 完全向量化：避免了Python循环的性能损耗，在大batch或大尺寸张量上速度提升明显。
• 符合PyTorch设计风格：利用原生广播机制，代码简洁易读。

验证形状匹配 #

• 处理后的target_vals形状是(batch, 1, 1, width)
• 赋值目标g[:, :, :, :]形状是(batch, x_dim, y_dim, width)
• 根据PyTorch广播规则，单维度会自动扩展为对应维度的大小，完美匹配赋值需求。

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