backend.prepare input_tensor_dict argument not useful because of new graph creation

[anishathalye]

It would be nice if the input_tensor_dict keyword argument to backend.prepare() was actually useful, but right now, because the code creates a new graph, it's not usable. (The caller can't create a tensor that's in the same graph that the code is going to use, because that graph isn't created until later.)

If we try to pass in a placeholder created outside/before the call to prepare(), we get the following kind of error:

ValueError: Tensor("split:0", shape=(3, 3, 3, 16), dtype=float32) must be from the same graph as Tensor("transpose:0", shape=(1, 360, 360, 3), dtype=float32) (graphs are <tensorflow.python.framework.ops.Graph object at 0x7f129c388748> and <tensorflow.python.framework.ops.Graph object at 0x7f129cc90048>).

I'm not familiar with the implementation details of onnx-tensorflow. Would it be reasonable to modify prepare() to take a graph= keyword argument, and then use the supplied graph instead of always creating a new graph?

[chudegao]

@anishathalye can you show some example code how to reproduce this issue and what's your use case. Currently, graph is only used for special case(training). Maybe you can try create a new tensor as value of input_tensor_dict to avoid graph conflict.

[anishathalye]

Yes, I was using this for training-like purposes; I needed to compute derivatives through the model for the purpose of producing adversarial examples for NeuralHash. More specifically, I had an ONNX model m(_), and I needed to compute derivatives through m(f(_)), where f was implemented elsewhere in my code. I'm not using this fancier functionality in the public repo that I linked, but here's essentially what I was trying to do that led to the error mentioned above:

a = tf.placeholder(...)
b = f(a)  # calls some tensorflow functions, creates some more graph nodes...
model = backend.prepare(input_tensor_dict={'input': b}, training_mode=True)
...
g, = tf.gradients(model.tensor_dict['output'], a)

I can share the full code example if that would be helpful, but I think this explanation should suffice. The issue is that backend.prepare() places the model in a new TF graph, and the a tensor (and whatever tensors were created by f(a)) are in the default graph, so I can't compute gradients through the whole thing with a single tf.gradients() application. (Probably possible to work around by splitting up the gradient computation, and using the grad_ys kwarg of tf.gradients(), but that is a little annoying.)

If backend.prepare() supports a graph kwarg, then I can do something like backend.prepare(..., graph=tf.compat.v1.get_default_graph()), and then everything works out.

I was using this modification to onnx-tensorflow in a local fork to achieve this:

diff --git i/onnx_tf/backend.py w/onnx_tf/backend.py
index 8f58002..193ffc0 100644
--- i/onnx_tf/backend.py
+++ w/onnx_tf/backend.py
@@ -94,13 +94,14 @@ class TensorflowBackend(Backend):
                                              **kwargs)

   @classmethod
-  def _onnx_graph_to_tensorflow_rep(cls, graph_def, opset, strict, **kwargs):
+  def _onnx_graph_to_tensorflow_rep(cls, graph_def, opset, strict, graph=None, **kwargs):
     """ Convert ONNX graph to TensorflowRep.

     :param graph_def: ONNX GraphProto object.
     :param opset: ONNX OperatorSetIdProto list.
     :param strict: whether to enforce semantic equivalence between the original model
       and the converted tensorflow model.
+    :param graph: The TensorFlow graph in which to place the model.
     :kwargs: additional arguements to generate tensor_dict for model debugging
     :return: TensorflowRep object.
     """
@@ -128,7 +129,7 @@ class TensorflowBackend(Backend):

     module = BackendTFModule(handlers, opset, strict, graph_def, cls)
     signatures = dict()
-    tf_rep_graph = tf.Graph()
+    tf_rep_graph = graph if graph is not None else tf.Graph()
     with tf_rep_graph.as_default():
       for value_info in graph_def.input:
         if value_info.name in initialized or not value_info.type.HasField(

[chudegao]

I think it's ok to add graph to kwargs. Do not need to update api _onnx_graph_to_tensorflow_rep. The arg can be get like :

 graph = kwargs['graph'] if 'graph' in kwargs else None 

Would you like to cretae a PR? Or I can create one.

[anishathalye]

It will probably be easier for you to make a proper PR, since you're more familiar with the code, so feel free to go ahead and do it.

[chudegao]

Ok. I will make a PR. Can you share some exmple code how to create the graph/input tensor. I will test following your case to make sure it can work.

[anishathalye]

After you setup the model.onnx following the instructions in AsuharietYgvar/AppleNeuralHash2ONNX, you can run the following code. Though you don't have to use that model, you could do this with any other ONNX model as well.

import numpy as np
import tensorflow as tf
import onnx
from onnx_tf.backend import prepare

tf.compat.v1.disable_eager_execution()

x = tf.compat.v1.placeholder(tf.float32, (1, 3, 360, 360))
onnx_model = onnx.load('model.onnx')
model = prepare(onnx_model, training_mode=True, input_tensor_dict={'image': x}, graph=tf.compat.v1.get_default_graph())
logits = model.tensor_dict['leaf/logits']
g, = tf.gradients(logits, x)

sess = tf.compat.v1.Session()
sess.run(tf.compat.v1.global_variables_initializer())
sess.run(g, {x: np.random.normal(size=(1, 3, 360, 360))})