Retrieve tensorflow device placement logging in a variable

slai-nick · April 28, 2022, 2:26pm

Hi,

I would like to retrieve tensorflow device placement as data I can manipulate from within python.
I use tf.debugging.set_log_device_placement(True) to log the device placement to stderr, however I haven’t found a way to catch it into a variable.

I tried both using contextlib.redirect_stderr and adding a StreamHandler to tf.get_logger(), none of these resulted in the logging output being caught.

See my original question from stack overflow.

Bhack · April 29, 2022, 3:54pm

I think that we don’t have a python API to do this as I suppose it is done at c++ level:

github.com

tensorflow/tensorflow/blob/master/tensorflow/core/common_runtime/placer.cc#L109-L128

      
        
            void LogDeviceAssignment(const Node* node, bool log_device_placement) {
              // Log placement if log_device_placement is set.
              if (log_device_placement) {
                printf("%s: (%s): %s\n", node->name().c_str(), node->type_string().c_str(),
                       node->assigned_device_name().c_str());
                LOG(INFO) << node->name() << ": "
                          << "(" << node->type_string()
                          << "): " << node->assigned_device_name();
              }
              if (VLOG_IS_ON(1)) {
                if (VLOG_IS_ON(4)) {
                  VLOG(4) << "\nNode:\n"
                          << node->def().DebugString()
                          << "placed on: " << node->assigned_device_name();
                } else {
                  VLOG(1) << node->name() << "(" << node->type_string()
                          << ") placed on: " << node->assigned_device_name();
                }
              }
            }

slai-nick · May 11, 2022, 1:12pm

Alright fair enough.
I wonder though why python’s contextlib.redirect_stderr isn’t able to catch it… But that’s more of a python question than a tensorflow. But if anyone knows answer I take it

Bhack · May 11, 2022, 1:27pm

Yes, you could try to take a look at:

https://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/

slai-nick · May 11, 2022, 1:49pm

Awesome, thank you for the link!