TensorFlow保存模型时出现引用未追踪资源的AssertionError如何解决

报错原因 #

Keras 自定义层的子层、可训练变量如果存放在普通 Python 列表、元组等原生容器中，Keras 的资源追踪机制不会自动识别这些对象，模型导出时就会判定这些变量属于未追踪资源，触发你遇到的报错。
你的代码里有两处不符合Keras追踪规则的实现：

• InputsEmbedding 层使用普通 list 存储 transformations 子层列表
• ConvolutionStack 层使用普通 list 存储 conv_blocks，且每个卷积块是普通 Python 元组，元组内的 Conv2D、LayerNormalization 等层也不会被自动追踪

你单独测试 InputsEmbedding 时未触发报错，是因为简单模型场景下 Keras 的追踪逻辑刚好能覆盖到相关变量，层级更深的完整复杂模型中，普通容器内的层就会丢失追踪。你无法通过变量名匹配到张量ID，是因为报错信息里的数字ID是运行时动态生成的张量实例标识，不是变量定义的静态名称，自然无法匹配。

解决方法 #

将所有存储子层的原生 Python 容器替换为 Keras 专门的可追踪容器：

1. 存储子层列表使用 tf.keras.layers.ModuleList 替换普通 list
2. 存储子层的固定结构集合也使用 ModuleList 替换普通 tuple

具体修改点 #

1. 调整 InputsEmbedding 层的子层存储逻辑：

# 原代码
self.transformations = list()
# 修改为
self.transformations = layers.ModuleList()

2. 调整 ConvolutionStack 层的子层存储逻辑：
   首先修改 conv_blocks 的容器类型：

# 原代码
self.conv_blocks = list()
# 修改为
self.conv_blocks = layers.ModuleList()

然后修改 _new_convolution_block 函数的返回值，将元组改为 ModuleList：

# 原返回代码
return (
    layers.Conv2D(
        filters=filters,
        kernel_size=kernel_size,
        strides=strides,
        use_bias=use_bias,
        kernel_regularizer=kernel_regularizer,
        bias_regularizer=bias_regularizer,
    ),
    norm_layer,
    activation_layer,
)
# 修改为
return layers.ModuleList([
    layers.Conv2D(
        filters=filters,
        kernel_size=kernel_size,
        strides=strides,
        use_bias=use_bias,
        kernel_regularizer=kernel_regularizer,
        bias_regularizer=bias_regularizer,
    ),
    norm_layer,
    activation_layer,
])

修改完成后重新运行测试，Keras 就能正常追踪所有子层的变量，模型保存报错就会消失。

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