如何在 TensorFlow 2 中重置初始化

Question

If I try to change parallelism in TensorFlow 2 after initializing a tf.Variable ,

import tensorflow as tf
_ = tf.Variable([1])
tf.config.threading.set_inter_op_parallelism_threads(1)

I get an error

RuntimeError: Inter op parallelism cannot be modified after initialization.

I understand why that could be, but it (and possibly other factors) are causing my tests to interfere with each other. For example

def test_model():  # this test
   v = tf.Variable([1])
   ...

def test_threading():  # is breaking this test
   tf.config.threading.set_inter_op_parallelism_threads(1)
   ...

How do I reset the TensorFlow state so that I can set the threading?

Answer 1

This is achievable in a "hacky" way. But I'd recommend doing this the right way (ie by setting up config at the beginning).

import tensorflow as tf
from tensorflow.python.eager import context

_ = tf.Variable([1])

context._context = None
context._create_context()

tf.config.threading.set_inter_op_parallelism_threads(1)

Edit : What is meant by setting up config at the beginning,

import tensorflow as tf
from tensorflow.python.eager import context

tf.config.threading.set_inter_op_parallelism_threads(1)
_ = tf.Variable([1])

But there could be circumstances where you cannot always do this. Merely pointing out the conventional way of setting up config in tf . So if your circumstances don't allow you to fix tf.config at the beginning you have to reset your tf.eager.context as shown in the solution above.