MindSpore 2.3.1 + Ascend 910B：自定义算子在动态Shape场景下反向传播出现Shape推断不一致
我在使用MindSpore 2.3.1框架配合Ascend 910B硬件进行一个多模态模型的训练时，遇到了自定义算子的梯度传播问题。该自定义算子在前向推理阶段能够正常工作，但在反向传播（bprop）时，当输入数据具有动态维度（如batch size为None）时，会出现形状推断失败的错误。
环境信息：
错误信息：
核心代码片段：
已尝试的解决方案：
-
检查环境变量配置，确认ASCEND_HOME、LD_LIBRARY_PATH等路径正确
-
尝试固定所有输入维度，问题暂时消失，但无法满足实际业务需求
-
在bprop函数中显式处理动态维度，但依然出现形状不匹配
-
检查算子注册时的动态轴支持设置，确认已开启dynamic_shape=True
-
清理TBE编译缓存，重新编译算子
具体现象：
-
当batch size为具体数值（如32）时，训练正常
-
当batch size为动态维度（None）时，反向传播报错
-
错误发生在梯度计算阶段，框架期望输出形状为(None, 128, 64)，但实际得到(32, 128, 64)
-
CANN日志显示反向算子输入Shape被硬编码为具体数值
疑问点：
-
在动态Shape场景下，自定义算子的bprop函数应该如何正确返回具有动态属性的Tensor？
-
是否需要为反向算子单独注册infer_shape函数？具体应该如何实现？
点赞
收藏
回复
分享
举报
已解决
MindSpore 2.3.1 + Ascend 910B：自定义算子在动态Shape场景下反向传播出现Shape推断不一致
我在使用MindSpore 2.3.1框架配合Ascend 910B硬件进行一个多模态模型的训练时，遇到了自定义算子的梯度传播问题。该自定义算子在前向推理阶段能够正常工作，但在反向传播（bprop）时，当输入数据具有动态维度（如batch size为None）时，会出现形状推断失败的错误。
环境信息：
硬件环境：Ascend 910B
MindSpore版本：2.3.1
CANN版本：7.0.RC1
执行模式：PYNATIVE模式（训练时切换为GRAPH模式）
Python版本：3.9
操作系统：Ubuntu 20.04
错误信息：
核心代码片段：
import mindspore as ms from mindspore import ops, nn from mindspore.ops import CustomRegOp, DataType # 自定义算子前向实现 def my_custom_forward(x, y): # 实际计算逻辑 output = x * y + x.sum(axis=1, keepdims=True) return output # 自定义算子反向实现 def my_custom_backward(dout, x, y): # 梯度计算逻辑 dx = dout * y + dout.sum(axis=1, keepdims=True) dy = dout * x return dx, dy # 算子注册 custom_op = CustomRegOp("MyDynamicOp") \ .input("x", "required") \ .input("y", "required") \ .output("output") \ .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \ .attr("dynamic_shape", "bool", "true") # 在模型中使用 class MyModel(nn.Cell): def __init__(self): super().__init__() self.custom_op = ops.Custom(my_custom_forward, out_shape=lambda x, y: (None, x[1], x[2]), out_dtype=ms.float32, func_type="hybrid") def construct(self, x, y): return self.custom_op(x, y)已尝试的解决方案：
检查环境变量配置，确认
ASCEND_HOME、LD_LIBRARY_PATH等路径正确尝试固定所有输入维度，问题暂时消失，但无法满足实际业务需求
在bprop函数中显式处理动态维度，但依然出现形状不匹配
检查算子注册时的动态轴支持设置，确认已开启
dynamic_shape=True清理TBE编译缓存，重新编译算子
具体现象：
当batch size为具体数值（如32）时，训练正常
当batch size为动态维度（None）时，反向传播报错
错误发生在梯度计算阶段，框架期望输出形状为
(None, 128, 64)，但实际得到(32, 128, 64)CANN日志显示反向算子输入Shape被硬编码为具体数值
疑问点：
在动态Shape场景下，自定义算子的bprop函数应该如何正确返回具有动态属性的Tensor？
是否需要为反向算子单独注册infer_shape函数？具体应该如何实现？