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


      
          
          // Returns true if the node has no inputs and produces outputs
          // that are consumed by a single node.
          //
          // TODO(vrv): Currently this handles only nodes with one output, but
          // this could be extended to handle the case where a node has many
          // outputs that are connected to nodes in the same colocation group.
          bool IsGeneratorNode(const Node* node) {
            return node->num_inputs() == 0 && node->num_outputs() == 1 &&
                   !IsRefType(node->output_type(0));
          }
          
          // If a node is an Identity op with input and output on the same device,
          // assign this Identity the same device. If the node already has a requested
          // or assigned device, don't touch it.
          bool MatchIdentityOperation(const Node* node) {
            if (!node) {
              return false;
            }

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!

Topic		Replies	Views
Parallelising custom function in tensorflow using graph execution General Discussion tfrt , gpu , xla , help_request , tfcore	13	2876	January 21, 2022
Operator device placement for models from tensorflow hub General Discussion models , tfhub , help_request	10	2205	May 12, 2022
Trying to understand Tensorflow Internals, would appreciate any insight on Tensor data General Discussion help_dev , education , tfcore	4	2346	July 5, 2021
Can't find libdevice directory ${CUDA_DIR}/nvvm/libdevice General Discussion build	10	10723	April 10, 2023
How does AssignVarOp work? General Discussion pluggable_device , help_request	11	2243	September 6, 2022

Retrieve tensorflow device placement logging in a variable

Related topics