使用 TPU 时抑制 Tensorflow“执行操作”消息

Question

我知道对于 Tensorflow 2 及以后的版本，Tensorflow 打印的输出可以根据级别相当容易地抑制：

import tensorflow as tf
tf.get_logger().setLevel('ERROR')

但是，当我在 Google Colab 上使用 TPU 时，使用这种方法根本不会过滤掉以“Executing op”开头的消息。如何过滤掉这些消息？

Answer 1

我最终不得不使用来自 logging 的自定义过滤器来解析 sys.stdout 和 sys.stderr：

# Suppress TPU messages which start with "Executing op"

import sys, re, logging

class Filter(object):
    def __init__(self, stream):
        self.stream = stream

    def __getattr__(self, attr_name):
        return getattr(self.stream, attr_name)

    def write(self, data):
        if not data.startswith("Executing op"):
            self.stream.write(data)
            self.stream.flush()

    def flush(self):
        self.stream.flush()
  
sys.stdout = Filter(sys.stdout)
sys.stderr = Filter(sys.stderr)

logger = logging.getLogger(__name__)

我不太确定这是否优雅，但我找不到更好的解决方案。希望这可以帮助遇到同样问题的其他人。