使用PyTorch自定义自动求导特性时反向传播出现异常张量步长

问题成因 #

你遇到的步长为(0, 0)的张量，是PyTorch反向传播时的内存优化手段：

• 当计算sum()这类操作的梯度时，理论上需要生成和输入同形状的全1张量，但PyTorch不会分配新内存存储完整的全1张量，而是创建一个标量广播视图。
• 这个视图的形状为(3,3)，但步长设为(0,0)——访问任何位置的元素时，都会读取内存中同一个标量值（此处为1），以此实现广播效果，同时大幅节省内存。

在复杂自定义Autograd函数中直接调用contiguous()会出问题，原因在于：

• contiguous()会强制将视图转换为连续内存的张量，对于步长(0,0)的广播视图，转换后的张量会把内存中唯一的元素复制到所有位置，看似结果正确，但在多分支梯度叠加、in-place操作等场景下，会破坏PyTorch的自动求导追踪逻辑，或引发内存依赖冲突，最终导致错误的梯度结果。
• 此外，非连续视图的contiguous()转换可能会丢失原张量的梯度传播关联，导致后续梯度计算异常。

解决方法 #

1. 避免不必要的contiguous()转换 #

如果不需要修改张量内容，也没有操作强制要求连续内存，直接返回grad_output即可。PyTorch的视图张量在反向传播中是安全的，不会影响梯度计算的正确性：

@staticmethod
def backward(ctx, grad_output):
    print(grad_output.shape, grad_output.stride())
    return grad_output  # 直接返回，无需转换

2. 用clone()替代contiguous()（推荐） #

如果确实需要连续内存的张量（比如要进行in-place修改、或某些操作依赖连续张量），使用clone()会更稳妥。clone()会创建一个独立的连续张量，同时完整保留梯度传播信息：

@staticmethod
def backward(ctx, grad_output):
    print(grad_output.shape, grad_output.stride())
    return grad_output.clone()  # 生成独立的连续张量，确保梯度正确

3. 按需判断后转换 #

如果一定要用contiguous()，可以先判断张量是否连续，避免不必要的转换：

@staticmethod
def backward(ctx, grad_output):
    print(grad_output.shape, grad_output.stride())
    if not grad_output.is_contiguous():
        grad_output = grad_output.contiguous()
    return grad_output

不过这种方法在广播视图场景下的效果和clone()类似，但clone()的语义更清晰，也更不容易出错。

验证示例 #

修改后的代码运行后，x.grad会正确输出3x3的全1张量，且在复杂场景下不会出现梯度错误。

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