How to replace tf.graph and tf.session from tensorflow 1.x with appropriate code in tensorflow 2?

Viewed 270

I am migrating code from tensorflow 1.x to tensorflow 2.7. I have retrained models and replaced the old ones with new models, however, they are very slow as compared to the old ones, and I think the reason is that models are now loaded for every input sent to server whereas in tf 1.x, sessions and graphs were used so model were loaded only once thus computations were reduced. Here is the code used for loading model in tf 1.

set_session(sess)
englishToHindiEncoder=load_model("models/englishToHindiEncoder.h5",compile=False)
englishToHindiEncoder._make_predict_function()
englishToHindiDecoder=load_model("models/englishToHindiDecoder.h5",compile=False)
englishToHindiDecoder._make_predict_function()
graph = tf.compat.v1.get_default_graph()

Further, below is the code where these models are being called and used.

def decode_sequence(vector):
  revhindiDict = dict((i, char) for char, i in hindiDict.items())
  global sess
  global graph
  states_value = np.zeros((1, 1, 82))
  with graph.as_default():
    set_session(sess)
    states_value = englishToHindiEncoder.predict(vector)
  target_seq = np.zeros((1, 1, 82))
  target_seq[0, 0, hindiDict['\t']] = 1

I want to know how can I remove these lines of code and replace them with corresponding tensorflow 2.0 code. The lines that need to be removed are as follows :

set_session(sess)
graph = tf.compat.v1.get_default_graph()

and

with graph.as_default():
   set_session(sess)

as these lines are now obsolete in tf 2. One way to do is to directly delete those lines but it significantly increases the processing time when requests are sent in bulk. So, I want to know the most opimized and proper way to replace those sessions and graphs code in tensorflow 2.0 and rewrite it.

1 Answers

In Tensorflow 2.X, to run the code in the graph mode you can wrap the function with @tf.function.

For Example:

@tf.function
def decode_sequence(vector):
  revhindiDict = dict((i, char) for char, i in hindiDict.items())
  #global sess
  #global graph
  states_value = np.zeros((1, 1, 82))
  #with graph.as_default():
    #set_session(sess)
  states_value = englishToHindiEncoder.predict(vector)
  target_seq = np.zeros((1, 1, 82))
  target_seq[0, 0, hindiDict['\t']] = 1

Please refer to this document for more details.

Related