如何从现有模型中图层的入站节点中删除 training=True？

Question

假设有一个模型以h5 文件的形式给出，也就是说，我无法更改构建模型架构的代码：

from tensorflow.keras.layers import Input, BatchNormalization
from tensorflow.keras.models import Model

inputs = Input(shape=(4,))
outputs = BatchNormalization()(inputs, training=True)
model = Model(inputs=inputs, outputs=outputs)
model.save('model.h5', include_optimizer=False)

现在我想删除training=True 部分，即我想将BatchNormalization 视为附加到没有此标志的模型。

我目前的尝试如下：

import numpy as np
from tensorflow.keras.models import load_model

model = load_model('model.h5')

for layer in model.layers:
    for node in layer.inbound_nodes:
        if "training" in node.call_kwargs:
            del node.call_kwargs["training"]

model.predict(np.asarray([[1, 2, 3, 4]]))

但 model.predict 调用失败并出现以下错误（我使用的是 TensorFlow 2.5.0）：

ValueError: Could not pack sequence. Structure had 1 elements, but flat_sequence had 2 elements.  Structure: ((<KerasTensor: shape=(None, 4) dtype=float32 (created by layer 'input_1')>,), {}), flat_sequence: [<tf.Tensor 'model/Cast:0' shape=(None, 4) dtype=float32>, True].

如何解决/解决这个问题？

（当使用node.call_kwargs["training"] = False 而不是del node.call_kwargs["training"] 时model.predict 不会崩溃，但它只是表现得好像什么都没有改变，即，修改后的标志被忽略。）

Answer 1

我发现，修改call_kwargs帮助后再次保存并重新加载模型。

import numpy as np
from tensorflow.keras.models import load_model

model = load_model('model.h5')

# Removing training=True
for layer in model.layers:
    for node in layer.inbound_nodes:
        if "training" in node.call_kwargs:
            del node.call_kwargs["training"]

# The two following lines are the solution.
model.save('model_modified.h5', include_optimizer=False)
model = load_model('model_modified.h5')

model.predict(np.asarray([[1, 2, 3, 4]]))

一切都很好。 :)

Answer 2

你试过了吗

for layer in model.layers:
    layer.trainable=False