如何禁用警告"tensorflow:Method (on_train_batch_end) is slow compared to the batch update (). Check your ca

Question

我正试图用keras风格的tensorflow 2.0实现随机权重平均(SWA)，所以我需要每一步都更新SWA模型的权重。我已经编写了一个自定义回调来实现这一点，但我每走一步都会收到警告。以下是一些详细信息：

我的自定义回调：

class CustomCallback(tf.keras.callbacks.Callback):
def __init__(self, valid_data, output_path, swa_alpha=0.99, eval_every=500, eval_batch=16, fold=None):
self.valid_inputs = valid_data[0]
self.valid_outputs = valid_data[1]
self.eval_batch = eval_batch
self.swa_alpha = swa_alpha
self.fold = fold
self.output_path = output_path
self.rho_value = -1  # record the best rho for report
self.eval_every = eval_every
def on_train_begin(self, logs={}):
self.swa_weights = self.model.get_weights()
def on_batch_end(self, batch, logs={}):
# update swa parameters
alpha = min(1 - 1 / (batch + 1), self.swa_alpha)
current_weights = self.model.get_weights()
for i, layer in enumerate(self.model.layers):
self.swa_weights[i] = alpha * self.swa_weights[i] + (1 - alpha) * current_weights[i]
# validation
if batch > 0 and batch % self.eval_every == 0:
# do validation
val_pred = self.model.predict(self.valid_inputs, batch_size=self.eval_batch)
rho_val = compute_spearmanr(self.valid_outputs, val_pred)  # the metric
# set the swa parameters and do validation
self.model.set_weights(self.swa_weights)
swa_val_pred = self.model.predict(self.valid_inputs, batch_size=self.eval_batch)
swa_rho_val = compute_spearmanr(self.valid_outputs, swa_val_pred)
# reset the original parameters
self.model.set_weights(current_weights)
# check whether to save model and update best rho value
if rho_val > self.rho_value:
self.rho_value = rho_val
self.model.save_weights(f'{self.output_path}/fold-{fold}-best.h5')
del current_weights
gc.collect()

输出是这样的：

WARNING:tensorflow:Method (on_train_batch_end) is slow compared to the batch update (11.428264). Check your callbacks.
WARNING:tensorflow:Method (on_train_batch_end) is slow compared to the batch update (11.464315). Check your callbacks.
WARNING:tensorflow:Method (on_train_batch_end) is slow compared to the batch update (11.502968). Check your callbacks.
WARNING:tensorflow:Method (on_train_batch_end) is slow compared to the batch update (11.518413). Check your callbacks.

我每一步都会收到警告，这意味着在不运行验证代码的情况下，更新SWA参数(self.model.get_weights()和下面的for循环)的代码足够慢。

我知道更新参数非常慢，因为model.get_weights()和model.set_weights()都会对params进行深度复制(根据我的实验，新numpy ndarray的新列表)。

我认为我的SWA实现没有错(如果有任何错误，请告诉我)，所以我只想禁用警告。

我尝试过的：

1. 添加代码os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"和os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"以禁用WARNING
2. 在model.fit()中将verbose设置为2和0，即model.fit(..., verbose=2, ...)和model.fit(..., verbose=0, ...)

两者都不起作用。

知道吗？感谢您提前提供的帮助！

Answer 1

这不是一个非常令人满意的答案，但TF_CPP_MIN_LOG_LEVEL不工作是一个已知问题：TF_CPP_MIN_LOG-LEVEL不适用于TF2.0 dev20190820。

我在tensorflow==2.1.0-rc1上用一个玩具例子重现了你的问题：

import os
import time
os.environ['TF_CPP_MIN_LOG_LEVEL'] = "2"
import tensorflow as tf
tf.get_logger().setLevel("WARNING")
tf.autograph.set_verbosity(2)
print(tf.__version__)
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
model = tf.keras.models.Sequential([
tf.keras.layers.Flatten(input_shape=(28, 28)),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dropout(0.2),
tf.keras.layers.Dense(10, activation='softmax')
])
class CustomCallback(tf.keras.callbacks.Callback):
def on_train_batch_end(self, batch, logs=None):
time.sleep(3)
model.compile(optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
model.fit(x_train, y_train, epochs=1, callbacks=[CustomCallback()])

2.1.0
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11493376/11490434 [==============================] - 31s 3us/step
Train on 60000 samples
WARNING:tensorflow:Method (on_train_batch_end) is slow compared to the batch update (3.002797). Check your callbacks.
32/60000 [..............................] - ETA: 1:57:38 - loss: 2.4674 - accuracy: 0.0938WARNING:tensorflow:Method (on_train_batch_end) is slow compared to the batch update (3.002938). Check your callbacks.
...

标准建议(os.environ['TF_CPP_MIN_LOG_LEVEL']、tf.get_logger().setLevel("WARNING")或tf.autograph.set_verbosity(2))都不起作用，我怀疑您将不得不等到上述问题得到解决。