Batch Prediction using GPUs with local runner

[jlewi]

Using GPUs for batch prediction could be really valuable.

There's a couple of different ways we could support this.

1. If we use a framework like Spark or Flink running on a K8s cluster with GPU nodes then the workers should be able to directly use the GPUs.

2. If we run Spark/Flink/Dataflow external to K8s such that the workers don't have direct access to GPUs then we could deploy TFServing on a K8s cluster with GPUs and then the workers could send batches of requests to the model to do inference.

[jlewi]

Some things to investigate

• TFServing currently doesn't support multi GPUs but you can potentially work around this by running 1 container per GPU.

• NVIDIA has TensorRT to optimize nets for GPUs

  • But this only supports a subset of TF graphs.

• NVIDIA has GRE

[yixinshi]

/assign yixinshi

[ankushagarwal]

With the new integration between TensorRT and TensorFlow 1.7, TensorRT optimizes compatible sub-graphs and let's TensorFlow execute the rest.

[jlewi]

@yixinshi How's this going?

[yixinshi]

I think I will start with a container image of Tf-serving as the base image for batch prediction. Then optimize it to use TensorFlow RT and/or GRE.

[jlewi]

P1 because we'd like to have this in our 0.2 release. However, we will probably only support local beam runner for GPUs in 0.2.

[jlewi]

@yixinshi What is the likelihood this will make 0.2?

[bhack]

Is tensorflow serving compilable with TensorRT? tensorflow/serving#864

[jlewi]

@bhack Don't know.

[jlewi]

@yixinshi Can we close this issue? I think batch prediction with GPUs and local runner is working now?