[converter] Add table initializer to graph execution.

[lina128]

Background:
For saved model graphs that have hashtable or lookuptable operations, the underlying hashtable needs to be initialized before being used. From graph architecture, the initialization subgraph is detached from the inference subgraph, so that initialization doesn't happen in every inference. To support hashtable ops in TF.js, we made below changes:

1. Extract the initialization subgraph and store its topology in modelInitializer field in model.json. See PR: [converter]Support table initializer in converter. #3958
2. Load the initialization subgraph and execute it during model load. See this PR.
3. Add hashtable ops implementation and add hashtableMap in model. Only dispose the hashtables when model is disposed. (To resume this PR: Add HashTableV2, LookupTableV2Find #3576)

Changes in this PR:

1. Add modelInitializer to the model and graph representations.
2. In load method last step loadSync, load the initialization subgraph into memory through method transformInitializer, and then execute the subgraph.
3. Refactored GraphExecutor's execute method into private method _execute and two public methods execute and initialize, which will use the private method to execute the graph. The shared logic in _execute can run a graph given an ordered node list. The original execute method takes care of user provided inputs and outputs. The initialize method doesn't have inputs tensors, so it just calculates the ordered list and call _execute.
4. Added a method to calculate ordered list for initialization subgraph which is basically consist of Const ops and other ops, no input tensors.

Notes:

1. This implementation only handles table initializer. We intentionally don't generalize to other initializers until we know the use case.

---

This change is 

[pyu10055]

Reviewable status: 0 of 1 approvals obtained (waiting on @lina128 and @pyu10055)

---

tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py, line 337 at r1 (raw file):

  if initializer_graph_def and initializer_outputs:
    model_json[common.ARTIFACT_MODEL_INITIALIZER] = {}
    model_json[common.ARTIFACT_MODEL_INITIALIZER]['outputs'] = [

Is output nodes needed? Since the initializer_graph is derived from the output nodes, the output nodes can be derived from the graph.

---

tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py, line 337 at r1 (raw file):

  if initializer_graph_def and initializer_outputs:
    model_json[common.ARTIFACT_MODEL_INITIALIZER] = {}
    model_json[common.ARTIFACT_MODEL_INITIALIZER]['outputs'] = [

make the keys constant

---

tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py, line 406 at r1 (raw file):

      # Only support table initializers for now.
      if meta_graph.collection_def and meta_graph.collection_def[
          'table_initializer']:

make it a constant

---

tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py, line 419 at r1 (raw file):

          tf.import_graph_def(frozen_initializer_graph_def, name='')

      return frozen_graph, frozen_initializer_graph, initializer_output_names

is the frozen_intializer_graph trimmed with only the nodes that related to the initializer output nodes?

---

tfjs-converter/src/executor/graph_executor.ts, line 171 at r1 (raw file):

  initialize(): void {
    const orderedNodes = getGraphInTopologicalOrder(this._outputs);
    this._execute(orderedNodes, this.outputNodes);

I think you can use the existing execute method with empty input array.

---

tfjs-converter/src/executor/model_analysis.ts, line 142 at r1 (raw file):

 *     nodes and then calculate execution order from the input nodes.
 */
export function getGraphInTopologicalOrder(outputs: Node[]): Node[] {

I feel this method may not be necessary, since to execute the initializer graph, there is no input needed.
Just constant is enough. If you can the existing execute method with empty input it should run.

---

tfjs-converter/src/operations/operation_mapper.ts, line 71 at r1 (raw file):

  // Converts the initialzer subgraph from Tensorflow GraphDef to local
  // representation for TensorFlow.js API
  transformInitializer(graph: tensorflow.IGraphDef, outputs: Node[]): Graph {

I think the original transformGraph function can automatically detect outputs, it should work for initializer without otuput nodes specified.

[lina128]

Hi Ping, thank you so much for the review. See my comments below.

Reviewable status: 0 of 1 approvals obtained (waiting on @pyu10055)

---

tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py, line 337 at r1 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

Is output nodes needed? Since the initializer_graph is derived from the output nodes, the output nodes can be derived from the graph.

Right, we can derive from the graph.

---

tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py, line 337 at r1 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

make the keys constant

Ack, will change in the other PR (the converter python code one).

---

tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py, line 406 at r1 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

make it a constant

Ack, will change in the other PR.

---

tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py, line 419 at r1 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

is the frozen_intializer_graph trimmed with only the nodes that related to the initializer output nodes?

I think so. At least the example I tried seems to be the case. We just trust grappler to do the right prune.

---

tfjs-converter/src/executor/graph_executor.ts, line 171 at r1 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

I think you can use the existing execute method with empty input array.

That's the first thing I tried then found too much nuances, this split is the cleanest I can think of. The shareable logic is the execution part given an ordered list of nodes. How it gets the ordered list differ a lot between initialization graph and inference graph. Main differences result from:

1. inference graph rely on inputs signature to build ordered list, if inputs are empty, it will throw error of missingInputs.
2. initialization graph inputs are consist of Const ops and other ops (hash_table op).

Because of these two reasons, the implementation of getting ordered list is quite different. See detailed comments inline below.

---

tfjs-converter/src/executor/model_analysis.ts, line 142 at r1 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

I feel this method may not be necessary, since to execute the initializer graph, there is no input needed.
Just constant is enough. If you can the existing execute method with empty input it should run.

For initializer graph, input doesn't only consist of Const ops, it also has other op, such as hash_table op. It will initialize an empty hashtable and return the hashtable as a tensor. The lookuptableImportV2 takes three inputs: hash_table (op), keys (const op), values (const op).

Besides this reason, there are other reasons passing in empty input doesn't work:

1. Because inputs is empty, orderedNodes will be null. Then it calls this.compile where it fails. https://github.com/tensorflow/tfjs/blob/master/tfjs-converter/src/executor/graph_executor.ts#L192

2. In compile, it first calls getExecutionSubgraph: https://github.com/tensorflow/tfjs/blob/master/tfjs-converter/src/executor/graph_executor.ts#L146

3. In getExecutionSubgraph, it will find inputs (nodes that don't have input), if those nodes are not in the passed in inputs list, it will be recorded as missingInputs: https://github.com/tensorflow/tfjs/blob/master/tfjs-converter/src/executor/model_analysis.ts#L76

4. Then back in compile, if there are missing inputs, it will throw error: https://github.com/tensorflow/tfjs/blob/master/tfjs-converter/src/executor/graph_executor.ts#L159

One alternative design, we can change the design of these methods to allow missing inputs (which I also tried). The problem with that is that in getExecutionSubGraph, we also further prune the graph by recording usedNodes (nodes needed from inputs to outputs: https://github.com/tensorflow/tfjs/blob/master/tfjs-converter/src/executor/model_analysis.ts#L64), because inputs are empty, usedNodes will just be output nodes. Note, the initializer graph can have multiple nodes, for example, the values nodes where the values tensor may be from a bunch of other nodes (for example, in the dnn model, it comes from cast and range and some constants). Later, the execute method will use getNodesInTopologicalOrder to get the ordered list, in this method, it will only use nodes from usedNodes: https://github.com/tensorflow/tfjs/blob/master/tfjs-converter/src/executor/model_analysis.ts#L105 And because the required nodes are not in the usedNodes, the orderedList is basically just [outputNode].

We can make changes in all these mentioned place to conditionally handle the exception case for initializer, but I feel that this design can be too messy, imagine a lot of condition check and longer code, therefore making readability of the logic much worse than right now. That is why I prefer the current implementation in this PR, which has a clean cut of the shared logic and the different parts. The shared logic is the graph execution part given an ordered list. The different part is how we prepare an ordered list for inference graph vs. initializer graph. Because the initializer graph is much simpler (no control flow, no function, no dynamic flow, etc.), we just implement a new method getGraphInTopologicalOrder to prepare the ordered list. And we can keep the original code for preparing ordered list in inference graph intact.

But that's my current thinking, I'm open to other ideas too. WDYT?

---

tfjs-converter/src/operations/operation_mapper.ts, line 71 at r1 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

I think the original transformGraph function can automatically detect outputs, it should work for initializer without otuput nodes specified.

Right, will remove it.

[lina128]

Reviewable status: 0 of 1 approvals obtained (waiting on @pyu10055)

---

tfjs-converter/src/executor/model_analysis.ts, line 142 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

For initializer graph, input doesn't only consist of Const ops, it also has other op, such as hash_table op. It will initialize an empty hashtable and return the hashtable as a tensor. The lookuptableImportV2 takes three inputs: hash_table (op), keys (const op), values (const op).

Besides this reason, there are other reasons passing in empty input doesn't work:

1. Because inputs is empty, orderedNodes will be null. Then it calls this.compile where it fails. https://github.com/tensorflow/tfjs/blob/master/tfjs-converter/src/executor/graph_executor.ts#L192

2. In compile, it first calls getExecutionSubgraph: https://github.com/tensorflow/tfjs/blob/master/tfjs-converter/src/executor/graph_executor.ts#L146

3. In getExecutionSubgraph, it will find inputs (nodes that don't have input), if those nodes are not in the passed in inputs list, it will be recorded as missingInputs: https://github.com/tensorflow/tfjs/blob/master/tfjs-converter/src/executor/model_analysis.ts#L76

4. Then back in compile, if there are missing inputs, it will throw error: https://github.com/tensorflow/tfjs/blob/master/tfjs-converter/src/executor/graph_executor.ts#L159

One alternative design, we can change the design of these methods to allow missing inputs (which I also tried). The problem with that is that in getExecutionSubGraph, we also further prune the graph by recording usedNodes (nodes needed from inputs to outputs: https://github.com/tensorflow/tfjs/blob/master/tfjs-converter/src/executor/model_analysis.ts#L64), because inputs are empty, usedNodes will just be output nodes. Note, the initializer graph can have multiple nodes, for example, the values nodes where the values tensor may be from a bunch of other nodes (for example, in the dnn model, it comes from cast and range and some constants). Later, the execute method will use getNodesInTopologicalOrder to get the ordered list, in this method, it will only use nodes from usedNodes: https://github.com/tensorflow/tfjs/blob/master/tfjs-converter/src/executor/model_analysis.ts#L105 And because the required nodes are not in the usedNodes, the orderedList is basically just [outputNode].

We can make changes in all these mentioned place to conditionally handle the exception case for initializer, but I feel that this design can be too messy, imagine a lot of condition check and longer code, therefore making readability of the logic much worse than right now. That is why I prefer the current implementation in this PR, which has a clean cut of the shared logic and the different parts. The shared logic is the graph execution part given an ordered list. The different part is how we prepare an ordered list for inference graph vs. initializer graph. Because the initializer graph is much simpler (no control flow, no function, no dynamic flow, etc.), we just implement a new method getGraphInTopologicalOrder to prepare the ordered list. And we can keep the original code for preparing ordered list in inference graph intact.

But that's my current thinking, I'm open to other ideas too. WDYT?

In a nutshell, I think preparing ordered list for initializer graph vs. inference graph are significantly different, it is better separate them than combine them and use condition check everywhere.

[pyu10055]

Reviewable status: 0 of 1 approvals obtained (waiting on @lina128 and @pyu10055)

---

tfjs-converter/src/executor/model_analysis.ts, line 142 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

In a nutshell, I think preparing ordered list for initializer graph vs. inference graph are significantly different, it is better separate them than combine them and use condition check everywhere.

My point is the existing execution code should allow empty input. If it does not, we should probably make it work. Since initializer graph is not unique in the sense that there could be other inference graphs that have no inputs as well.

I am curious where did the error happen? Is it because that the hash_table op contains no inputs, which makes the code think that is an input node that user needs to supply?

[lina128]

Reviewable status: 0 of 1 approvals obtained (waiting on @pyu10055)

---

tfjs-converter/src/executor/model_analysis.ts, line 142 at r1 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

My point is the existing execution code should allow empty input. If it does not, we should probably make it work. Since initializer graph is not unique in the sense that there could be other inference graphs that have no inputs as well.

I am curious where did the error happen? Is it because that the hash_table op contains no inputs, which makes the code think that is an input node that user needs to supply?

From code, it seems that inference graphs with no inputs will work if all inputs are constants, so the value can be extracted from weight file. If an inference graph has inputs that is not constant, it expects it to be provided as an input, but this input will be an op, not a tensor. But the current api doesn't allow op as inputs, only namedTensorMap. We have two options here:

1. Change the api so that it accepts input op. The change should be minimal, the risk here is execute is a public api, will allowing users to pass in op confusing or dangerous? The use case in this PR is internal (transparent to user), but to solve the use case, we have to change a public api without a public use case.
2. Alternatively, we don't change the api, and only change the internal utilities to (1) allow input as op, (2) allow automatic detection of input ops. This goes back to the problems 1 to 4 I listed above. Basically we need to make changes in these places and it should work. The risk here is the added logic is not trivial so it complicates code, and can affect positively/negatively existing behavior of the public api execute.

So I personally prefer the solution in this PR, which keeps current code for public api intact, and not complicate code. But if you feel that supporting input op for execute (alternative #1 above) can have other use cases that can justify the public api change, then we can go with that route. But we must think through those use cases carefully, like instructing user the lifecycle of ops. WDYT?

[pyu10055]

Reviewable status: 0 of 1 approvals obtained (waiting on @lina128 and @pyu10055)

---

tfjs-converter/src/executor/graph_executor.ts, line 180 at r2 (raw file):

   * array.
   */
  execute(inputs: NamedTensorMap, outputs: string[]): Tensor[] {

looks like the outputs param should be optional as stated in the jdoc.

---

tfjs-core/src/io/http.ts, line 198 at r2 (raw file):

    }

    const artifacts:ModelArtifacts = {

you might need to double check on model save method? (indexDb, localStorage, see if the initializer can be stored/loaded properly)

---

tfjs-core/src/io/http.ts, line 210 at r2 (raw file):

    const initializer = modelConfig.modelInitializer;
    if (initializer) {
      artifacts['modelInitializer'] = initializer;

can you do artifacts.modelInitialzer ?

[lina128]

Reviewable status: 0 of 1 approvals obtained (waiting on @pyu10055)

---

tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py, line 337 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

Right, we can derive from the graph.

Done.

---

tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py, line 337 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

Ack, will change in the other PR (the converter python code one).

Done.

---

tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py, line 406 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

Ack, will change in the other PR.

Done.

---

tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py, line 419 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

I think so. At least the example I tried seems to be the case. We just trust grappler to do the right prune.

Done.

---

tfjs-converter/src/executor/graph_executor.ts, line 171 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

That's the first thing I tried then found too much nuances, this split is the cleanest I can think of. The shareable logic is the execution part given an ordered list of nodes. How it gets the ordered list differ a lot between initialization graph and inference graph. Main differences result from:

1. inference graph rely on inputs signature to build ordered list, if inputs are empty, it will throw error of missingInputs.
2. initialization graph inputs are consist of Const ops and other ops (hash_table op).

Because of these two reasons, the implementation of getting ordered list is quite different. See detailed comments inline below.

Done.

---

tfjs-converter/src/executor/graph_executor.ts, line 180 at r2 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

looks like the outputs param should be optional as stated in the jdoc.

I can change it in another PR. We should probably change executeAsync too.

---

tfjs-converter/src/executor/model_analysis.ts, line 142 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

From code, it seems that inference graphs with no inputs will work if all inputs are constants, so the value can be extracted from weight file. If an inference graph has inputs that is not constant, it expects it to be provided as an input, but this input will be an op, not a tensor. But the current api doesn't allow op as inputs, only namedTensorMap. We have two options here:

1. Change the api so that it accepts input op. The change should be minimal, the risk here is execute is a public api, will allowing users to pass in op confusing or dangerous? The use case in this PR is internal (transparent to user), but to solve the use case, we have to change a public api without a public use case.
2. Alternatively, we don't change the api, and only change the internal utilities to (1) allow input as op, (2) allow automatic detection of input ops. This goes back to the problems 1 to 4 I listed above. Basically we need to make changes in these places and it should work. The risk here is the added logic is not trivial so it complicates code, and can affect positively/negatively existing behavior of the public api execute.

So I personally prefer the solution in this PR, which keeps current code for public api intact, and not complicate code. But if you feel that supporting input op for execute (alternative #1 above) can have other use cases that can justify the public api change, then we can go with that route. But we must think through those use cases carefully, like instructing user the lifecycle of ops. WDYT?

Discussed offline, we will go with option 2 above. The risk is minimal. I found a way to make the logic not more complicated than currently. Basically, we separate inputNodes and initNodes, and allow initNodes to be inferred from graph rather than provided.

---

tfjs-converter/src/operations/operation_mapper.ts, line 71 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

Right, will remove it.

Done.

---

tfjs-core/src/io/http.ts, line 198 at r2 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

you might need to double check on model save method? (indexDb, localStorage, see if the initializer can be stored/loaded properly)

Yes, those need changes too. I plan to add them in another PR, so as to not make this PR larger. WDYT?

---

tfjs-core/src/io/http.ts, line 210 at r2 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

can you do artifacts.modelInitialzer ?

Done.

[lina128]

Hi @pyu10055 , I've addressed your comments, PTAL, thank you!

Reviewable status: 0 of 1 approvals obtained (waiting on @pyu10055)

[pyu10055]

Looks good, thanks! Can you add an e2e test that cover the run trip from model conversion to execution?

Reviewable status: 0 of 1 approvals obtained (waiting on @lina128 and @pyu10055)

---

tfjs-core/src/io/http.ts, line 198 at r2 (raw file):

Previously, lina128 (Na Li) wrote…

Yes, those need changes too. I plan to add them in another PR, so as to not make this PR larger. WDYT?

SG

[lina128]

Hi Ping, I tried to add a v1 model with hashtable in e2e's convert_predict.py, the conversion will fail if not skip op check, because LookupTableFindV2 and HashTableV2 are unsupported ops. Other than that, the conversion should be fine. I can add this e2e case once I checked in the HashTableV2 PR, which is the next PR I plan to get your review and merge after this PR is merged.

Reviewable status: 0 of 1 approvals obtained (waiting on @lina128 and @pyu10055)

[pyu10055]

sounds good, thanks!

Reviewed 1 of 12 files at r1, 4 of 10 files at r2, 4 of 7 files at r3.
Reviewable status: complete! 1 of 1 approvals obtained (waiting on @lina128 and @pyu10055)