使用TensorFlow剪枝模型时遇EagerTensor无assign属性错误求助

DevOps 小满

2026-8-11

权重剪枝时触发AttributeError: 'EagerTensor' object has no attribute 'assign'

在使用TensorFlow Model Optimization的prune_low_magnitude对加载的Keras模型进行权重剪枝时，执行model_for_pruning = prune_low_magnitude(model, **pruning_params)时触发以下错误：
AttributeError: 'tensorflow.python.framework.ops.EagerTensor' object has no attribute 'assign'

剪枝代码

# pruning

import tensorflow_model_optimization as tfmot
import numpy as np

origModelFile = 'modeltest.h5'
model = tf.keras.models.load_model(origModelFile)

prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude
epochs = 15
batch_size = 2048
validation_split = 0.1
num_images = x_train.shape[0] * (1 - validation_split)

end_step = np.ceil(num_images / batch_size).astype(np.int32) * epochs
pruning_params = {
      'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(initial_sparsity=.95,
                                                               final_sparsity=.8,
                                                               begin_step=0,
                                                               end_step=end_step)
}
model_for_pruning = prune_low_magnitude(model, **pruning_params) # 报错行

model_for_pruning.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

pruned_model = tfmot.sparsity.keras.strip_pruning(model_for_pruning)

pruned_model.summary()

完整报错栈

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
~\AppData\Local\Temp\ipykernel_33564\3560105584.py in <module>
     13                                                                end_step=end_step)
     14 }
---&gt; 15 model_for_pruning = prune_low_magnitude(model, **pruning_params)
     16 
     17 model_for_pruning.compile(optimizer='adam',

A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\keras\metrics.py in inner(*args, **kwargs)
     72       except Exception as error:
     73         self.bool_gauge.get_cell(MonitorBoolGauge._FAILURE_LABEL).set(True)
---&gt; 74         raise error
     75 
     76     if self.bool_gauge:

A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\keras\metrics.py in inner(*args, **kwargs)
     67     def inner(*args, **kwargs):
     68       try:
---&gt; 69         results = func(*args, **kwargs)
     70         self.bool_gauge.get_cell(MonitorBoolGauge._SUCCESS_LABEL).set(True)
     71         return results

A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\prune.py in prune_low_magnitude(to_prune, pruning_schedule, block_size, block_pooling_type, pruning_policy, sparsity_m_by_n, **kwargs)
    208     if pruning_policy:
    209       pruning_policy.ensure_model_supports_pruning(to_prune)
--&gt; 210     return _add_pruning_wrapper(to_prune)
    211   elif is_keras_layer:
    212     params.update(kwargs)

A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\prune.py in _add_pruning_wrapper(layer)
    179         raise ValueError('Subclassed models are not supported currently.')
    180 
--&gt; 181       return keras.models.clone_model(
    182           layer, input_tensors=None, clone_function=_add_pruning_wrapper)
    183     if isinstance(layer, pruning_wrapper.PruneLowMagnitude):

A:\Anaconda\lib\site-packages\keras\models\cloning.py in clone_model(model, input_tensors, clone_function)
    446           model, input_tensors=input_tensors, layer_fn=clone_function)
    447     else:
--&gt; 448       return _clone_functional_model(
    449           model, input_tensors=input_tensors, layer_fn=clone_function)
    450 

A:\Anaconda\lib\site-packages\keras\models\cloning.py in _clone_functional_model(model, input_tensors, layer_fn)
    187   # Reconstruct model from the config, using the cloned layers.
    188   input_tensors, output_tensors, created_layers = (
--&gt; 189       functional.reconstruct_from_config(model_configs,
    190                                          created_layers=created_layers))
    191   metrics_names = model.metrics_names

A:\Anaconda\lib\site-packages\keras\engine\functional.py in reconstruct_from_config(config, custom_objects, created_layers)
   1310         while layer_nodes:
   1311           node_data = layer_nodes[0]
-&gt; 1312           if process_node(layer, node_data):
   1313             layer_nodes.pop(0)
   1314           else:

A:\Anaconda\lib\site-packages\keras\engine\functional.py in process_node(layer, node_data)
   1254         input_tensors = (
   1255             base_layer_utils.unnest_if_single_tensor(input_tensors))
-&gt; 1256       output_tensors = layer(input_tensors, **kwargs)
   1257 
   1258       # Update node index map.

A:\Anaconda\lib\site-packages\keras\utils\traceback_utils.py in error_handler(*args, **kwargs)
     65     except Exception as e:  # pylint: disable=broad-except
     66       filtered_tb = _process_traceback_frames(e.__traceback__)
---&gt; 67       raise e.with_traceback(filtered_tb) from None
     68     finally:
     69       del filtered_tb

A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\pruning_wrapper.py in tf__call(self, inputs, training, **kwargs)
     71                 update_mask = ag__.converted_call(ag__.ld(utils).smart_cond, (ag__.ld(training), ag__.ld(add_update), ag__.ld(no_op)), None, fscope)
     72                 ag__.converted_call(ag__.ld(self).add_update, (ag__.ld(update_mask),), None, fscope)
---&gt; 73                 ag__.converted_call(ag__.ld(self).add_update, (ag__.converted_call(ag__.ld(self).pruning_obj.weight_mask_op, (), None, fscope),), None, fscope)
     74 
     75                 def get_state_1():

A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\pruning_impl.py in tf__weight_mask_op(self)
     11                 try:
     12                     do_return = True
---&gt; 13                     retval_ = ag__.converted_call(ag__.ld(tf).group, (ag__.converted_call(ag__.ld(self)._weight_assign_objs, (), None, fscope),), None, fscope)
     14                 except:
     15                     do_return = False

A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\pruning_impl.py in tf___weight_assign_objs(self)
    122                 _ = ag__.Undefined('_')
    123                 masked_weight = ag__.Undefined('masked_weight')
--&gt; 124                 ag__.if_stmt(ag__.converted_call(ag__.ld(tf).distribute.get_replica_context, (), None, fscope), if_body_1, else_body_1, get_state_4, set_state_4, (), 0)
    125                 try:
    126                     do_return = True

A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\pruning_impl.py in if_body_1()
     98                     def else_body():
     99                         pass
--&gt; 100                     ag__.if_stmt(ag__.ld(values_and_vars), if_body, else_body, get_state_2, set_state_2, (), 0)
    101 
    102                 def else_body_1():

A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\pruning_impl.py in if_body()
     94 
     95                     def if_body():
---&gt; 96                         ag__.converted_call(ag__.ld(assign_objs).append, (ag__.converted_call(ag__.converted_call(ag__.ld(tf).distribute.get_replica_context, (), None, fscope).merge_call, (ag__.ld(update_fn),), dict(args=(ag__.ld(values_and_vars),)), fscope),), None, fscope)
     97 
     98                     def else_body():

A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\pruning_impl.py in update_fn(distribution, values_and_vars)
     52                         var = ag__.Undefined('var')
     53                         value = ag__.Undefined('value')
---&gt; 54                         ag__.for_stmt(ag__.ld(values_and_vars), None, loop_body, get_state, set_state, (), {'iterate_names': '(value, var)'})
     55                         try:
     56                             do_return_1 = True

A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\pruning_impl.py in loop_body(itr)
     49                         def loop_body(itr):
     50                             (value, var) = itr
---&gt; 51                             ag__.converted_call(ag__.ld(update_objs).append, (ag__.converted_call(ag__.ld(distribution).extended.update, (ag__.ld(var), ag__.ld(update_var)), dict(args=(ag__.ld(value),)), fscope_1),), None, fscope_1)
     52                         var = ag__.Undefined('var')
     53                         value = ag__.Undefined('value')

A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\pruning_impl.py in update_var(variable, reduced_value)
     34                                 try:
     35                                     do_return_2 = True
---&gt; 36                                     retval__2 = ag__.converted_call(ag__.ld(tf_compat).assign, (ag__.ld(variable), ag__.ld(reduced_value)), None, fscope_2)
     37                                 except:
     38                                     do_return_2 = False

A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\keras\compat.py in tf__assign(ref, value, name)
     34                         do_return = False
     35                         raise
---&gt; 36                 ag__.if_stmt(ag__.converted_call(ag__.ld(hasattr), (ag__.ld(tf), 'assign'), None, fscope), if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)
     37                 return fscope.ret(retval_, do_return)
     38         return tf__assign

A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\keras\compat.py in else_body()
     30                     try:
     31                         do_return = True
---&gt; 32                         retval_ = ag__.converted_call(ag__.ld(ref).assign, (ag__.ld(value),), dict(name=ag__.ld(name)), fscope)
     33                     except:
     34                         do_return = False

AttributeError: Exception encountered when calling layer "prune_low_magnitude_conv2d" (type PruneLowMagnitude).

in user code:

    File "A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\pruning_wrapper.py", line 288, in call  *
        self.add_update(self.pruning_obj.weight_mask_op())
    File "A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\pruning_impl.py", line 254, in weight_mask_op  *
        return tf.group(self._weight_assign_objs())
    File "A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\pruning_impl.py", line 225, in update_var  *
        return tf_compat.assign(variable, reduced_value)
    File "A:\Anaconda\lib\site-packages\tensorflow_model_optimization\python\core\keras\compat.py", line 28, in assign  *
        return ref.assign(value, name=name)

    AttributeError: 'tensorflow.python.framework.ops.EagerTensor' object has no attribute 'assign'


Call arguments received by layer "prune_low_magnitude_conv2d" (type PruneLowMagnitude):
  • inputs=tf.Tensor(shape=(None, 14, 8, 8), dtype=float32)
  • training=False
  • kwargs=<class 'inspect._empty'>

解决方案

1. 匹配TensorFlow与TensorFlow Model Optimization版本

TFMO与TF版本必须严格匹配，否则会出现兼容性问题。例如：

TensorFlow 2.10 → TFMO 0.13.0
TensorFlow 2.15 → TFMO 0.18.0

执行以下命令安装对应版本：

pip install tensorflow-model-optimization==<对应版本号>

2. 修改模型加载方式，确保权重为可训练Variable

方式一：分开加载模型结构与权重

如果之前保存模型时分开保存了结构和权重：

from tensorflow.keras.models import model_from_json

# 加载模型结构
with open('model_config.json', 'r') as f:
    model_json = f.read()
model = model_from_json(model_json)

# 加载权重
model.load_weights('model_weights.h5')

如果之前没有分开保存，需要先重新保存：

# 假设orig_model是你的原始模型
orig_model.save_weights('model_weights.h5')
with open('model_config.json', 'w') as f:
    f.write(orig_model.to_json())

方式二：加载模型时禁用自动编译

model = tf.keras.models.load_model(origModelFile, compile=False)
# 手动编译模型
model.compile(
    optimizer='adam',
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy']
)

3. 确认模型类型为Functional/Sequential

TFMO剪枝不支持子类化（Subclassed）模型。如果你的模型是自定义子类模型，需要转换为Functional模型：

# 假设model是子类化模型，先通过输入示例得到输出
dummy_input = tf.random.normal(shape=(1, *x_train.shape[1:]))
model_output = model(dummy_input)
# 转换为Functional模型
functional_model = tf.keras.Model(inputs=dummy_input, outputs=model_output)

4. 剪枝前触发模型前向传播

在调用prune_low_magnitude前，让模型运行一次前向传播，确保权重被正确初始化为Variable：

# 生成匹配模型输入形状的dummy数据
dummy_input = tf.random.normal(shape=(1, *x_train.shape[1:]))
model(dummy_input)

5. 检查并修复权重类型

遍历模型层，检查权重是否为tf.Variable：

for layer in model.layers:
    for weight in layer.weights:
        print(f"Layer: {layer.name}, Weight: {weight.name}, Type: {type(weight)}")

如果发现EagerTensor类型的权重，需要重新保存模型，确保权重以Variable形式存储。

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

使用TensorFlow剪枝模型时遇EagerTensor无assign属性错误求助

剪枝代码 #

完整报错栈 #

解决方案 #

1. 匹配TensorFlow与TensorFlow Model Optimization版本 #

2. 修改模型加载方式，确保权重为可训练Variable #