How Tensorflow set device for each Operation ?

In Tensorflow, we only need to use snippet below to assign a device to a Operation:

How dose it implement? Let’s take a look.

There is a mechanism called ‘context manager’ in Python. For example, we can use it to add a wrapper for a few codes:

The result of running this script is:

Function ‘tag()’ works like a decorator. It will do something before and after those codes laying under its ‘context’.

Tensorflow uses the same principle.

This will call class Graph’s function ‘device()’. Its implementation:

The key line is ‘self._add_device_to_stack()’. Context of ‘device’ will add device name into stack of python, and when developer create an Operation it will fetch device name from stack and set it to this Operation.
Let’s check the code routine of creating Operation:

‘self._device_function_stack.peek_objs’ is where it peek the device name from stack.

Leave a Reply

This site uses Akismet to reduce spam. Learn how your comment data is processed.