How To Keep Lookup Tables Initialized For Prediction (and Not Just Training)?

I create a lookup table from tf.contrib.lookup, using the training data (as input). Then, I pass every input through that lookup table, before passing it through my model. This wor

Solution 1:

I think you would be better off using tf.tables_initializer() as the legacy_init_op.

tf.saved_model.main_op.main_op() also adds local and global initialization ops in addition to table initialization. when you load the saved model and it runs the legacy_init_op, it would reset your variables, which is not what you want.

Solution 2:

You can specify an "initialization" operation when you add a meta graph to your SavedModel bundle with tf.saved_model.builder.SavedModelBuilder.add_meta_graph, using the main_op or legacy_init_op kwarg. You can either use a single operation, or group together a number of operations with tf.group if you need more than one.

Note that in Cloud ML Engine, You'll have to use the legacy_init_op. However in future runtime_versions you will be able to use main_op (IIRC, starting with runtime_version == 1.2)

The saved_model module provides a built in tf.saved_model.main_op.main_op to wrap up common initialization actions in a single op (local variable initialization, and table initialization).

So in summary, code should look like this (adapted from this example):

  exporter = tf.saved_model.builder.SavedModelBuilder(
      os.path.join(job_dir, 'export', name))

  # signature_def gets constructed here

  with tf.Session(graph=prediction_graph) as session:
    # Need to be initialized before saved variables are restored
    session.run([tf.local_variables_initializer(), tf.tables_initializer()])
    # Restore the value of the saved variables
    saver.restore(session, latest)
    exporter.add_meta_graph_and_variables(
        session,
        tags=[tf.saved_model.tag_constants.SERVING],
        signature_def_map={
            tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_def
        },
        # Relevant change to the linked example is here!
        legacy_init_op=tf.saved_model.main_op.main_op()
    )

NOTE: If you are using the high level libraries (such as tf.estimator) this should be the default, and if you need to specify additional initialization actions you can specify them as part of the tf.train.Scaffold object that you pass to your tf.estimator.EstimatorSpec in your model_fn.