这是由于使用了XLA(Accelerated Linear Algebra)优化,该优化在选择要优化(或跟踪)的Op时会有一些限制。解决此问题的一种方法是在使用XLA时指定要跟踪的Op。例如,以下代码中使用tf.xla.experimental.compile
函数来编译模型并指定要跟踪的Op列表:
import tensorflow as tf
from absl import logging
# 定义模型
inputs = tf.keras.layers.Input(shape=(10,))
dense = tf.keras.layers.Dense(64, activation='relu')(inputs)
outputs = tf.keras.layers.Dense(1, activation='sigmoid')(dense)
model = tf.keras.Model(inputs=inputs, outputs=outputs)
# 编译模型并指定要跟踪的Op
with tf.xla.experimental.jit_scope(compile_ops=True):
model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
logging.info("XLA enabled")
model.summary()
# 训练模型
x_train = tf.random.normal(shape=(100, 10))
y_train = tf.random.uniform(shape=(100, 1), maxval=2, dtype=tf.int32)
model.fit(x_train, y_train, epochs=1)
# 保存模型
tf.saved_model.save(model, "saved_model")
在这个例子中,我们在使用XLA的作用域内编译模型,并将compile_ops
设置为True
以明确指定要跟踪的Op。这样,当我们保存模型时,就不会收到_update_step_xla
未跟踪的函数的警告了。