Add an option to turn on XLA_JIT.

NVIDIA / OpenSeq2Seq

Toolkit for efficient experimentation with Speech Recognition, Text2Speech and NLP

https://nvidia.github.io/OpenSeq2Seq

Apache License 2.0

1.54k stars 371 forks source link

Add an option to turn on XLA_JIT. #417

Closed trentlo closed 5 years ago

trentlo commented 5 years ago

This PR adds an option to enable XLA Just-In-Time (JIT) compilation mode.

trentlo commented 5 years ago

@okuchaiev please help to take a look.

Note that I mentioned to you that I see performance gain (~20-30%) by enabling XLA for Transformer (after disabling padding_removal). I refer to a TF revision newer than Apr11. So, you will not see performance gain if you use a public container with TF older than Apr11. However, this added argument should still work fine to enable the XLA JIT.

okuchaiev commented 5 years ago

Thanks!