Add --compute_lr_on_cpu option.
Add --compute_lr_on_cpu option.

Performance results on a machine with 8 V100s, CUDA 10, Tensorflow commit 8fbed5f43ccac490088679c86d6989afae7d9e1a with 53faa313b7628cd8c9fbb836544cc6482cafb7a4 patched out, tf_cnn_benchmarks commit 2058fffb46d1a40c375bdfee7f18165bbd9b2ce7. Each number was run for 5 iterations and the average was taken. The command used was:

python tf_cnn_benchmarks.py --batch_size=256 --model=resnet50_v1.5 --data_dir=/raid/imagenet-2012-tfrecord  --variable_update=replicated --all_reduce_spec=nccl --use_fp16=true  --datasets_use_prefetch=$use_prefetch  --datasets_repeat_cached_sample=true --num_gpus=8 --xla_compile=$xla

XLA   use_prefetch   im/s LR on GPU  im/s LR on CPU  % improvement
On    On             5704.624        7527.928        31.96%
On    Off            7167.35         7252.132        1.18%
Off   On             4209.73         4124.76         -2.02%
Off   Off            4404.674        4985.15         13.17%

Performance improves in every case except when XLA is off and datasets_use_prefetch is on, which is the default. Therefore, --compute_lr_on_cpu defaults to False, despite it helping performance in most cases.

I am unsure why performance is harmed when XLA is off and datasets_use_prefetch is on. I would have expected performance to always be improved, since the LR computations are on scalar values, which CPUs are better at. And, even when LR computations on done on the GPU, TensorFlow still places some ops on the CPU, since certain ops are always done on the CPU when its inputs are int32.

PiperOrigin-RevId: 217183629

tensorflow/benchmarks,tensorflow/benchmarks,annarev/benchmarks,annarev/benchmarks