TFTRT: Respect device placement requested by user

[meena-at-work]

Previously TRT engines would always default to running on GPU 0. On a multi GPU system, this does not make the best use of the available resources. This commit adds the ability to specify the GPU on which the TRT engine should run.

Signed-off-by: Meenakshi Venkataraman meenakshiv@nvidia.com

CC: @bixia1 @DEKHTIARJonathan @tfeher

[meena-at-work]

A TFTRT user will now be able to specify device placement with the tf.device() API before invoking convert.

model = create_model()
 input_saved_model_dir='./saved_model/my_model'
 model.save(input_saved_model_dir)

 with tf.device('gpu:1'):
     model_loaded = tf.saved_model.load(export_dir='./saved_model/my_model')
     print("Loaded Model device:", model_loaded.variables[0].device)

 with tf.device('gpu:2'):
     from tensorflow.python.compiler.tensorrt import trt_convert as trt
     converter = trt.TrtGraphConverterV2(input_saved_model_dir=input_saved_model_dir)
     converter.convert()
     converter.summary()

The output for the above with the feature is:

INFO:tensorflow:Assets written to: ./saved_model/my_model/assets
Loaded Model device: /job:localhost/replica:0/task:0/device:GPU:1
INFO:tensorflow:Linked TensorRT version: (8, 2, 4)
INFO:tensorflow:Loaded TensorRT version: (8, 2, 4)
INFO:tensorflow:Placing imported graph from `./saved_model/my_model` on device: /job:localhost/replica:0/task:0/device:GPU:2
TRTEngineOP Name                 Device        # Nodes # Inputs      # Outputs     Input DTypes       Output Dtypes      Input Shapes       Output Shapes     
================================================================================================================================================================
TRTEngineOp_000_000              device:GPU:2  5       1             1             ['float32']        ['float32']        [[-1, 784]]        [[-1, 10]]        

        - Const: 2x
        - MatMul: 2x
        - Relu: 1x

================================================================================================================================================================
[*] Total number of TensorRT engines: 1
[*] % of OPs Converted: 71.43% [5/7]