TensorFlow 2 compatibility: draft

[albertz]

I want to list all possible options we have to implement this.

First some references:

• Official TF1 to TF2 migration guide.
• My StackOverflow question: Best practice to write code compatible to both TensorFlow 1 and 2
  
• StackOverflow: Implicitly enable TensorFlow v2 behaviour in TensorFlow v1. In TF 1, use enable_v2_behavior and tensorflow.compat.v2 (since TF 1.14).
• Lingvo compat.py. The code then mostly uses import lingvo.compat as tf.

In any case, also for TF 2, we would use graph-mode, not eager mode.

The current goal:

• A transition phase for a while (a year or so):
  • The code should run fine on both TF 1 (including earlier versions, e.g. 1.8) and TF 2.
  • We should be able to optionally use some new TF 2 features.
• We slowly could drop the TF 1 support later.

I don't really see this question answered in the migration guide. Maybe the answer is that there simply is no "best practice" for this.

We could maybe do sth like this:

try:
  # https://www.tensorflow.org/guide/migrate
  import tensorflow.compat.v1 as tf
  tf.disable_v2_behavior()
except ImportError:
  import tensorflow as tf

However, that has some drawbacks:

• We would need to place this code snippet in every module where I use TF.
  • We could put this in some own module, like TFUtil.py, and then simply do from TFUtil import tf everywhere. But this is somewhat ugly.
• This might make code complicated which actually wants to use some new TF 2 feature.

We could maybe do sth like this:

try:
  # https://www.tensorflow.org/guide/migrate
  import tensorflow.compat.v1 as tf1
  tf1.disable_v2_behavior()
  import tensorflow as tf2
except ImportError:
  import tensorflow as tf1
  tf2 = None

This might be more clean. But I'm not sure if this is the way to go. Or how other people do it.

Also, instead of calling disable_v2_behavior, we could maybe just call disable_eager_execution. But that might make some other code more difficult to make compatible for both TF 1 and TF 2.

Or we could go the other way around, and use enable_v2_behavior in TF 1, and also tensorflow.compat.v2. But this works only since TF 1.14.

How do other bigger projects solve this?

Or maybe the answer is that this would be way too much effort, and so they simply migrate directly to TF 2?

Another aspect, but this might be less important for now: If it would be possible to also support JAX, or make the transition easy, that would be nice.

---

The solution now, as implemented in TFCompat.py, which is already in master:

import tensorflow as tf

if not getattr(tf, "compat", None) or not getattr(tf.compat, "v1", None):
  v1 = tf
  v2 = None
else:
  # PyCharm type-inference will take the latest reference,
  # so this `else` branch should lead us to a valid reference for "modern" TF versions (TF >=1.14, or TF 2).
  v1 = tf.compat.v1
  v2 = tf.compat.v2

if v2:
  tf.compat.v1.disable_eager_execution()
  tf.compat.v1.disable_v2_tensorshape()
  tf.compat.v1.disable_control_flow_v2()
  # tf.compat.v1.disable_v2_behavior()  -- not sure on this

try:
  import tensorflow.contrib
  have_contrib = True
except ImportError:
  have_contrib = False

So instead of using tf.compat.v1.XXX, you would use TFCompat.v1.XXX in your code (when needed). This works even on very old TF versions where tf.compat.v1 does not exist.
(In new code for optional new features, you could directly use tf.XXX or tf.compat.v1.XXX or tf.compat.v2.XXX (or whatever you want) (see contributing).)

The process of supporting all features of RETURNN in TF 2 is now work-in-progress. Many parts already work correctly.

At some later point, once we drop all TF <1.14 support or so, we could simply replace all TFCompat.v1 usage by tf.compat.v1.
Once we drop TF 1 support completely, we could slowly (gradually) replace all tf.compat.v1 (or TFCompat.v1) usage by new TF 2 variants.

[albertz]

Some update here: I pushed some initial TF 2 support already to master. The main logic is in TFCompat.py. See code usage. This is mostly safe for all existing code, i.e. we can work directly on the master branch.
The current TF 2 status: It already almost works. My goal is to make (almost) all tests pass. It looks quite good already. A lot of setups probably work already.

[albertz]

I consider this as complete now. All tests should pass now. Please open a new issue if anything does not work.

[albertz]

Small final update: Since commit 5aabe21, all relevant TF 2 test cases run through now, and are part of Travis now.

[jiangj-dc]

How to make SpecAug work with TF 2? The changes are tf.random_uniform to tf.random.uniform and tf.log to tf.math.log.

[albertz]

Yes, of course all your custom TF code in your config needs to be adopted for TF 2. Most code will work as is. And otherwise the changes are very minor. In all cases, it should always work by just using tf.compat.v1.XXX instead of tf.XXX.