PyTorch中带梯度张量广播[B,1]张量做通道偏移时in-place报错

问题描述 #

我有一个形状为(B, 3, N)的张量p（示例为2个batch、3个通道（x/y/z）、5个点），且requires_grad=True：

# 2 batches, 3 channels (x, y, z), 5 points
p = torch.rand(2, 3, 5, requires_grad=True)
"""
p: tensor([[[0.8365, 0.0505, 0.4208, 0.7465, 0.6843],
         [0.9922, 0.2684, 0.6898, 0.3983, 0.4227],
         [0.3188, 0.2471, 0.9552, 0.5181, 0.6877]],

        [[0.1079, 0.7694, 0.2194, 0.7801, 0.8043],
         [0.8554, 0.3505, 0.4622, 0.0339, 0.7909],
         [0.5806, 0.7593, 0.0193, 0.5191, 0.1589]]], requires_grad=True)
"""

另有形状为[B, 1]的偏移张量z_shift：

z_shift = torch.tensor([[1.0], [10.0]], requires_grad=True)
"""
z_shift: tensor([[1.],
        [10.]], requires_grad=True)
"""
需求：为每个batch内所有点的z通道（索引2）加上对应偏移，x/y通道保持不变，预期结果如下：
```python
"""
p: tensor([[[0.8365, 0.0505, 0.4208, 0.7465, 0.6843],
         [0.9922, 0.2684, 0.6898, 0.3983, 0.4227],
         [1.3188, 1.2471, 1.9552, 1.5181, 1.6877]],

        [[0.1079, 0.7694, 0.2194, 0.7801, 0.8043],
         [0.8554, 0.3505, 0.4622, 0.0339, 0.7909],
         [10.5806, 10.7593, 10.0193, 10.5191, 10.1589]]])
"""

当requires_grad=False时，p[:, 2, :] += z_shift可实现需求，但在nn.Module的forward方法中（requires_grad=True）执行会报错：

RuntimeError: a view of a leaf Variable that requires grad is being used in an in-place operation.

解决方案 #

方法1：拼接新张量（推荐） #

构造新张量，保留原x/y通道，修改z通道后拼接，完全避免原地操作：

# 利用广播机制直接相加，无需手动扩展维度
p_new = torch.cat([
    p[:, :2, :],  # 保留x、y通道
    p[:, 2:3, :] + z_shift  # z通道加上对应偏移
], dim=1)

p_new是全新张量，支持正常梯度反向传播，不会触发错误。

方法2：克隆后原地修改 #

先克隆原张量（克隆后的张量非叶子节点），再对克隆版进行原地修改：

p_new = p.clone()
p_new[:, 2, :] += z_shift  # 广播机制自动匹配维度

此方法简单直观，适合需要保留原张量的场景。

方法3：逐通道掩码修改 #

通过掩码指定仅修改z通道，逻辑清晰但稍显繁琐：

# 创建掩码标记z通道位置
mask = torch.zeros_like(p, dtype=torch.bool)
mask[:, 2, :] = True

# 构造仅z通道有值的偏移张量
shift_tensor = torch.zeros_like(p)
shift_tensor[:, 2, :] = z_shift

# 应用偏移
p_new = torch.where(mask, p + shift_tensor, p)

报错原因 #

报错源于原地修改叶子节点的视图会破坏PyTorch的计算图追踪逻辑：当requires_grad=True时，p是叶子节点，p[:,2,:]是它的视图，原地操作会导致梯度无法正确反向传播。所有方案的核心都是避免直接修改原叶子节点的视图，转而创建新张量或克隆后修改。

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