Loading multiple tensorflow models at once in tf1.x causes problems because the operations in each model may have same name, and thus cause collision. While this can be avoided by manually giving unique names to tensors when developing a model, this is a cumbersome strategy.

Another approach would be to save each model in a separate graph, and assign separate session to each of these graphs. By assigning a graph for each model, there is no need to worry about name collision.

I wondered if assigning multiple graphs to a single session would be possible. But I have only ran into errors and it looks like this is not possible.

These findings can be used in code like this.

import tensorflow as tf, numpy as np, os


class Model1:
    def __init__(self, model_arch, model_weight):
        self.graph = tf.Graph()

        with self.graph.as_default():
            self.sess = tf.Session(graph=self.graph)

            with self.sess.as_default():
                with open(model_arch, 'r') as fd:
                    self.model = tf.keras.models.model_from_json(fd.read())
                self.model.load_weights(model_weight)
    
    def predict(self, x):
        with self.graph.as_default():
            with self.sess.as_default():
                out = self.model.predict(x)
        
        return out


model_arch = '/model_arch_json_path'
model_weight = '/model_weight_path'



os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1"


m1 = Model1(model_arch, model_weight)
m2 = Model1(model_arch, model_weight)

testinput = np.random.rand(1, 224,224,3).astype('float32')

out = m1.predict(testinput)

print(out.shape)

out = m2.predict(testinput)

print(out.shape)

the above code works well.

One thing I found out is that first I though the with tf.graph and with self.sess.as_default() context block would cause a problem since normally at the end of context block, the context is closed automatically. But according to the tensorflow docs, this is not the case for tf.Session.


0 Comments

Leave a Reply

Your email address will not be published.