A better implementation of device placement? about sonnet HOT 2 CLOSED

@mingyr Hi, thanks for your detailed comment.

I've looked into this before and unfortunately there are no publicly exposed APIs from TF which allow us to have more control over this. I would like to be able to capture the surrounding device context when constructing a module, as your middle example shows. However unlike the various kinds of scopes, tf.device doesn't return anything you can hang onto and reenter, eg:

with tf.variable_scope('some_name') as scope:
  print(scope)  # Prints a VariableScope object we can query, reenter, etc

with tf.device('/cpu:0') as device_information:
  print(device_information)  # prints `None`

Without being able to capture the context, we cannot respect device placement directives around where we construct modules. If you read the TF source you see that this works by accessing various private members of the Graph object and others, and it would be a huge maintenance burden for us to start relying on these private APIs.

We've had some conversations with the TF team about creating an external more advanced use of device placement (and other things(), and those discussions are ongoing. I hope to have some improvement in this area in the public version soon.

It's very useful to hear from users about the pain points of the library, so once again thankyou.

from sonnet.

@malcolmreynolds I am sincerely appreciating the detailed explanation and fully understanding the situation. The important thing is not we invent something perfect but we invent something based on which we can advance our daily job. I am glad that Sonnet is such a library which helps me a great deal. Thanks for the contribution made by you engineers at Deepmind which benefits the deep learning community a lot.

from sonnet.