reuse sess in backend

onnx_tf/backend_rep.py

@@ -15,6 +15,7 @@ def __init__(self, graph=None, inputs=None, outputs=None, tensor_dict=None):
     self._graph = graph
     self._inputs = inputs or []
     self._outputs = outputs or []
+    self._sess = None
     self._tensor_dict = tensor_dict or {}
 
   @property
@@ -60,30 +61,30 @@ def run(self, inputs, **kwargs):
 
     # TODO: handle name scope if necessary
     with self.graph.as_default():
-      with tf.Session() as sess:
-        if isinstance(inputs, dict):
-          feed_dict = inputs
-        elif isinstance(inputs, list) or isinstance(inputs, tuple):
-          if len(self.inputs) != len(inputs):
-            raise RuntimeError('Expected {} values for uninitialized '
-                               'graph inputs ({}), but got {}.'.format(
-                                   len(self.inputs), ', '.join(self.inputs),
-                                   len(inputs)))
-          feed_dict = dict(zip(self.inputs, inputs))
-        else:
-          # single input
-          feed_dict = dict([(self.inputs[0], inputs)])
-
-        feed_dict = {
-            self.tensor_dict[key]: feed_dict[key]
-            for key in self.inputs
-        }
-
-        sess.run(tf.global_variables_initializer())
-        outputs = [self.tensor_dict[output] for output in self.outputs]
-
-        output_values = sess.run(outputs, feed_dict=feed_dict)
-        return namedtupledict('Outputs', self.outputs)(*output_values)
+      sess = self._sess or tf.Session()
+      self._sess = sess
+
+      if isinstance(inputs, dict):
+        feed_dict = inputs
+      elif isinstance(inputs, list) or isinstance(inputs, tuple):
+        if len(self.inputs) != len(inputs):
+          raise RuntimeError('Expected {} values for uninitialized '
+                             'graph inputs ({}), but got {}.'.format(
+                                 len(self.inputs), ', '.join(self.inputs),
+                                 len(inputs)))
+        feed_dict = dict(zip(self.inputs, inputs))
+      else:
+        # single input
+        feed_dict = dict([(self.inputs[0], inputs)])
+
+      feed_dict = {self.tensor_dict[key]: feed_dict[key] for key in self.inputs}
+
+      sess.run(tf.global_variables_initializer())
+      outputs = [self.tensor_dict[output] for output in self.outputs]
+
+      output_values = sess.run(outputs, feed_dict=feed_dict)
+
+      return namedtupledict('Outputs', self.outputs)(*output_values)
 
   def export_graph(self, path):
     """Export backend representation to a Tensorflow proto file.