How to use K.get_session in Tensorflow 2.0 or how to migrate it?

In TensorFlow 2.0, the usage of K.get_session() has changed because TensorFlow 2.0 defaults to eager execution mode, which eliminates the need for a session to execute operations immediately. In TensorFlow 1.x, we often used K.get_session() to obtain the TensorFlow session for performing low-level operations such as initializing all variables, saving or loading models, etc.

If you need functionality similar to using K.get_session() in TensorFlow 1.x, there are several migration strategies:

1. Directly use TensorFlow 2.0's API

Since TensorFlow 2.0 defaults to eager execution, most operations can be executed directly without explicitly creating a session. For tasks like model training, evaluation, or other operations, you can directly leverage TensorFlow 2.0's high-level APIs, such as tf.keras. For example:

python
import tensorflow as tf

# Create a simple model
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(10, activation='relu', input_shape=(32,)),
    tf.keras.layers.Dense(1)
])

# Compile the model
model.compile(optimizer='adam', loss='mean_squared_error')

# Generate some random data for training
import numpy as np
x_train = np.random.random((1000, 32))
y_train = np.random.random((1000, 1))

# Train the model
model.fit(x_train, y_train, epochs=10)

2. Use tf.compat.v1.Session()

If your code depends on TensorFlow 1.x session functionality, you can continue using sessions via the tf.compat.v1 module. For instance, to explicitly initialize all variables, you can do the following:

python
import tensorflow as tf

# Disable eager execution
tf.compat.v1.disable_eager_execution()

# Create a variable
v = tf.Variable(1.0)

# Use tf.compat.v1.Session() to manage the session
with tf.compat.v1.Session() as sess:
    # Initialize all variables
    sess.run(tf.compat.v1.global_variables_initializer())

    # Use the session
    print(sess.run(v))  # Output: 1.0

3. Use tf.function to wrap functions

To retain the flexibility of eager execution while achieving graph execution efficiency in specific functions, you can use tf.function to decorate these functions. This enables you to obtain similar effects to building a static graph in TensorFlow 2.0:

python
import tensorflow as tf

@tf.function
def compute_area(side):
    return side * side

side = tf.constant(5)
print(compute_area(side))  # Output: tf.Tensor(25, shape=(), dtype=int32)

In summary, TensorFlow 2.0 provides a more concise and efficient approach to replace K.get_session() in TensorFlow 1.x. In most cases, you can directly use TensorFlow 2.0's API, or employ tf.compat.v1.Session() to maintain compatibility with legacy code where necessary.