Operator device placement for models from tensorflow hub

slai-nick · April 28, 2022, 3:16pm

Hi,

I would like to know the device placements of operators from a model retrieved from tensorflow hub.
I think the code will explain better:
test.py

import tensorflow as tf
import tensorflow_hub as hub

tf.debugging.set_log_device_placement(True)

layer = hub.KerasLayer('https://tfhub.dev/tensorflow/retinanet/resnet152_v1_fpn_1024x1024/1')
model = tf.keras.Sequential([layer])

When I run the script, it does print some device placement, but nothing related to any relevant operations.
I can get a list of all unique operators on a unix system using:

python test.py 2>&1|awk '{print $7,$10}'|sort|uniq

And I get the following:

AssignVariableOp /job:localhost/replica:0/task:0/device:CPU:0
DestroyResourceOp /job:localhost/replica:0/task:0/device:CPU:0
Identity /job:localhost/replica:0/task:0/device:CPU:0
RestoreV2 /job:localhost/replica:0/task:0/device:CPU:0
VarHandleOp /job:localhost/replica:0/task:0/device:CPU:0
_EagerConst /job:localhost/replica:0/task:0/device:CPU:0

Is it possible to retrieve operators placement when a model comes from tensorflow hub?

lgusm · May 3, 2022, 11:11am

I don’t understand your question

As far as I’m aware, you can place the model on any device you want (CPU, GPU and TPU) but I don’t think you can do per operator

slai-nick · May 11, 2022, 8:19am

Hello @lgusm, thank for taking the time to read my question.

So I am working on a tensorflow plugin for a device and I need real-world models to drive the development.
Here is my workflow has been looking like so far:

My first focus was on supporting Resnet50 operators, which is a widely available model and I have an implementation that doesn’t come from the hub.
I made a simple python script that loads the model and I activated tf.debugging.set_log_device_placement(True) option, this gives me a list of all the kernels I need to support for supporting Resnet50.

Now I would like to support other real-world tensorflow models, and tensorflow hub is a good place to find them.
However when I load a model from tensorflow hub like the RetinaNet in the example above, I don’t see any relevant kernel (I would like to see Conv2D for example).
I understand that a tensorflow hub model encapsulate the model as a KerasLayer, but at the lower level I assume that the Conv2D and other operators/kernels will be executed. I would like to retrieve what kernels/operators a model from tensorflow hub is going to execute.

Let me know if my question is better understandable now.

Mark_Daoust · May 11, 2022, 9:46pm

IDK how device placement works with SavedModel (hub just wraps SavedModel).

But try having a look at:

# load the saved_model
mod = hub.load('https://tfhub.dev/tensorflow/retinanet/resnet152_v1_fpn_1024x1024/1')
# Get the graph out of the default signature and convert it to a graph_def proto.
mod.signatures['serving_default'].graph.as_graph_def()

Bhack · May 11, 2022, 10:01pm

There are some limits related to hard-coding the device placement with tf.function

github.com

tensorflow/tensorflow/blob/v2.8.0/tensorflow/python/saved_model/save.py#L1285-L1295


      
          `tf.function` does not hard-code device annotations from outside the function
          body, instead of using the calling context's device. This means for example
          that exporting a model that runs on a GPU and serving it on a CPU will
          generally work, with some exceptions:
          
            * `tf.device` annotations inside the body of the function will be hard-coded
              in the exported model; this type of annotation is discouraged.
            * Device-specific operations, e.g. with "cuDNN" in the name or with
              device-specific layouts, may cause issues.
            * For `ConcreteFunctions`, active distribution strategies will cause device
              placements to be hard-coded in the function.

Something can be controlled at save time with:

But it is still not fully supported with load:

github.com/tensorflow/tensorflow

Python API of `saved_model.load` to support `tf.saved_model.experimental.VariablePolicy`

opened 07:31PM - 12 Jan 22 UTC

closed 06:41PM - 27 Jun 22 UTC

wenscarl

type:feature comp:apis TF 2.7 awaiting PR merge

**System information** - TensorFlow version (you are using): TF2.7 - CPU memor…y: 16Gb - GPU memory: 10Gb **Describe the feature and the current behavior/state.** Some models need to store large weights/tables on CPU and smaller ones on GPU. tf.saved_model.experimental.VariablePolicy is able to store device placement but not fully supported by tf.saved_model.load. Every variable is loaded to GPU if available. There is support on C++ level, but missing in python. A repo.py is as follows. Run as python repo.py --size 11, A OOM is observed. Set the number 11 to be any number b/w the memory size of your CPU and GPU in GB. ``` import argparse import tensorflow as tf class LN(tf.keras.layers.Layer): def __init__(self, rows, cols,trainable=True): super(LN, self).__init__(dtype=tf.float32) self.rows = rows self.cols = cols self.mat = None self.trainable=trainable def build(self, input_shape): self.mat = self.add_weight("mat",shape=[self.rows, self.cols], dtype=tf.float32, trainable=self.trainable) def call(self, indices): return tf.gather(params=self.mat, indices=indices) class TestModel(tf.keras.Model): def __init__(self, rows, cols): super().__init__() self.ln = LN(rows=rows, cols=cols) @tf.function def call(self, x): with tf.device('/cpu:0'): x = self.ln(x) with tf.device('/gpu:0'): x = tf.math.reduce_sum(x, axis=1) return x def main(): parser = argparse.ArgumentParser() parser.add_argument('--size', type=int, default=100, help='Table size in GiB') parser.add_argument('--load_cpu', type=bool, default=False, help='Load variables to cpu') args = parser.parse_args() ncols = 128 nrows = args.size * 2**30 // 512 path = '/tmp/saved_model_test' model = TestModel(rows=nrows, cols=ncols) indices = tf.zeros(shape=(65536, 1), dtype=tf.int32) outputs = model(indices) print(outputs) model.summary() print('Saving the model') tf.saved_model.save( model, path, options = tf.saved_model.SaveOptions( experimental_variable_policy=tf.saved_model.experimental.VariablePolicy.SAVE_VARIABLE_DEVICES)) print('Saved successfully') if args.load_cpu: with tf.device('/cpu:0'): loaded = tf.saved_model.load(path) # Good else: loaded = tf.saved_model.load(path) # OOM print('Load successfully') if __name__ == '__main__': main() ```

slai-nick · May 12, 2022, 9:36am

Awesome, I can extract the list of operators from this. Thank you!

By the way, is there any documentation about “graph_def”?
I am currently able to extract the list of op by converting the graph_def to string and extracting the info from it, but I am sure I could extract the relevant information directly from it.

slai-nick · May 12, 2022, 9:38am

Is it an explanation of why the device placement logging “doesn’t work” for operators within hub models? (if so then I am deducing that hub models are loaded inside tf.function)

Bhack · May 12, 2022, 9:40am

No, It was just related to the cases where a device could be hardcoded in the saved model.

slai-nick · May 12, 2022, 9:52am

Okay, thanks for clarifying.

Mark_Daoust · May 12, 2022, 11:12am

It’s a proto, defined here: tensorflow/tensorflow/core/framework/graph.proto at master · tensorflow/tensorflow · GitHub

Hub loads savedmodels. savedmodels contain a collection of tf.functions.

slai-nick · May 12, 2022, 11:55am

Thank you, I found it in the meantime

I am going to close the subject as I have all the information I needed. Thanks again to both of you for taking the time to answer!

Topic		Replies	Views
What is needed to run a hub model? General Discussion pluggable_device , tfhub , help_request	13	2452	June 7, 2022
TFX Object Detection General Discussion models , tfx , tflite , tfhub , help_request	13	8393	November 7, 2022
Retrieve tensorflow device placement logging in a variable General Discussion help_request	4	748	May 11, 2022
Parallelising custom function in tensorflow using graph execution General Discussion tfrt , gpu , xla , help_request , tfcore	13	2876	January 21, 2022
Tensorflow HUB - transfer learning for object detection General Discussion model_garden , tfhub , help_request	12	7118	April 25, 2022

Operator device placement for models from tensorflow hub

Related topics