在 GRAPH_MODE 下 construct 中调用 Parameter.set_data 引发 RuntimeError

[huadaox]

在 MindSpore 的 GRAPH_MODE 中实现自定义 nn.Cell 类时，若在 construct 方法中调用 Parameter 的 set_data 方法动态更新参数值，会导致编译失败，报错 RuntimeError: Failed to compile in GRAPH_MODE because the 'Parameter' object's method 'set_data' is not supported in 'construct' or function with @jit decorator.。此限制在 PYNATIVE_MODE 下不存在，表明这是 MindSpore 静态图编译的一个明确设计限制。
MindSpore 2.5.0
Python 3.9.20
windows11
仅CPU

这里是一个最小复现例子
import mindspore
import mindspore.nn as nn
import mindspore.ops as ops
from mindspore import Tensor, Parameter, context
import numpy as np

class TestCell(nn.Cell):
def init(self, num_features):
super(TestCell, self).init()
self.limit = Parameter(Tensor(1.0, dtype=mindspore.float32), requires_grad=False)
self.inc_step = Tensor(0.0001, dtype=mindspore.float32)
self.max_limit = Tensor(3.0, dtype=mindspore.float32)

def construct(self, x):
    clipped_x = ops.clip_by_value(x, -self.limit, self.limit)
    batch_size = Tensor(x.shape[0], dtype=mindspore.float32)
    limit_new = self.limit + self.inc_step * batch_size
    self.limit.set_data(ops.minimum(limit_new, self.max_limit))
    return clipped_x

测试设置

context.set_context(mode=context.GRAPH_MODE, device_target="CPU")
batch_size, channels, height, width = 4, 3, 32, 32
data = np.random.randn(batch_size, channels, height, width).astype(np.float32)
input_tensor = Tensor(data, dtype=mindspore.float32)

net = TestCell(num_features=channels)
output = net(input_tensor)
print("Output Shape:", output.shape)

=========================
报错如下：Traceback (most recent call last):
File "E:\Code\BatchRenormalization\BatchRenormalization\questionTarget.py", line 28, in
output = net(input_tensor)
File "E:\ProgramData\anaconda3\envs\Allign\lib\site-packages\mindspore\nn\cell.py", line 732, in call
out = self.compile_and_run(*args, **kwargs)
File "E:\ProgramData\anaconda3\envs\Allign\lib\site-packages\mindspore\nn\cell.py", line 1150, in compile_and_run
self.compile(*args, **kwargs)
File "E:\ProgramData\anaconda3\envs\Allign\lib\site-packages\mindspore\nn\cell.py", line 1133, in compile
_cell_graph_executor.compile(self, *self._compile_args, phase=self.phase,
File "E:\ProgramData\anaconda3\envs\Allign\lib\site-packages\mindspore\common\api.py", line 1897, in compile
result = self._graph_executor.compile(obj, args, kwargs, phase, self._use_vm_mode())
RuntimeError: Failed to compile in GRAPH_MODE because the 'Parameter' object's method 'set_data' is not supported in 'construct' or function with @jit decorator. Try to use the 'Parameter.set_data' externally such as initialized in the method 'init' before assigning.
For more details, please refer to https://www.mindspore.cn/docs/zh-CN/master/model_train/program_form/static_graph.html

---

• C++ Call Stack: (For framework developers)

---

mindspore\ccsrc\pipeline\jit\ps\static_analysis\prim.cc:1685 mindspore::abstract::`anonymous-namespace'::CheckObjAttrValid

---

• The Traceback of Net Construct Code:

---

0 In file E:\Code\BatchRenormalization\BatchRenormalization\questionTarget.py:18, 8~27

    self.limit.set_data(ops.minimum(limit_new, self.max_limit))
    ^~~~~~~~~~~~~~~~~~~

(See file 'E:\Code\BatchRenormalization\BatchRenormalization\rank_0\om/analyze_fail.ir' for more details. Get instructions about analyze_fail.ir at https://www.mindspore.cn/search?inputValue=analyze_fail.ir)

不含动态赋值的例子正常运行
class SimpleClip(nn.Cell):
def init(self, num_features):
super(SimpleClip, self).init()
self.limit = Tensor(1.0, dtype=mindspore.float32)
def construct(self, x):
return ops.clip_by_value(x, -self.limit, self.limit)

context.set_context(mode=context.GRAPH_MODE, device_target="CPU")
batch_size, channels, height, width = 4, 3, 32, 32
data = np.random.randn(batch_size, channels, height, width).astype(np.float32)
input_tensor = Tensor(data, dtype=mindspore.float32)
net = SimpleClip(num_features=3)
output = net(input_tensor)
print("Output Shape:", output.shape) # 输出 (4, 3, 32, 32)

影响 此限制阻止了在 GRAPH_MODE 下实现需要动态更新参数的算法（如 Batch Renormalization），限制了静态图优化的适用范围。例如，Batch Renormalization 要求标量 r_max 在每次迭代中更新，这一需求无法在当前 GRAPH_MODE 下实现。

[zhouyifeng888]

这个方式确实是不支持静态图操作的，不过应该也不需要这个操作，我用了很久了也没碰到过这个需要，通常训练过程中参数是自动更新的，如果需要手动在图里面去设置某些值的话，可以使用mindspore.ops.assign方法，参考文档：
https://www.mindspore.cn/docs/zh-CN/r2.5.0/api_python/ops/mindspore.ops.assign.html?highlight=assign#mindspore.ops.assign
还有一些index_addd、inplace_index_add等方法也可以修改Parameter中的值，我在llm的静态kvcache更新中就使用了这两个方法，具体用法可以文档上搜一下