convert a TPU tensorflow=1.15 model for deployment on GPU

Question

I am struggling to load a TF1.15 TPU trained model to a GPU. The model directory has following files:

model.ckpt-best.data-00000-of-00001
model.ckpt-best.index
model.ckpt-best.meta

I see these instructions from Google :

from tensorflow.python.saved_model import loader
from tensorflow.python.saved_model import tag_constants

with tf.Session() as sess:
  loader.load(sess, [tag_constants.SERVING], model_dir)
  sess.run(model_outputs, feed_dict={model_input: [input_image_batch]})

but there I run into an error:

OSError: SavedModel file does not exist at: .../bioalbert_large_pubmed_pmc_mimic/{saved_model.pbtxt|saved_model.pb}

Another recipe that fails:

with tf.Session() as sess:

    # load the computation graph
    loader = tf.compat.v1.train.import_meta_graph(f'{model_dir}/model.ckpt-best.meta')
    sess.run(tf.compat.v1.global_variables_initializer())
    # loader = loader.restore(sess, f'{model_dir}/model.ckpt-best.data-00000-of-00001')

error:

InvalidArgumentError: No OpKernel was registered to support Op 'TPUReplicatedInput'
 used by node input0 (defined at .../anaconda3/envs/tf1.15/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:1748) 
with these attrs: [T=DT_INT32, N=8]
Registered devices: [CPU, GPU, XLA_CPU, XLA_GPU]
Registered kernels:
  <no registered kernels>

     [[input0]]

Answer 1

If you're only using it for serving, you could try freezing the graph. The freeze_graph tool has an option to clear devices that might work for you?

tf.Session() also has options for soft device placement that might help:

# Launch the graph in a session that allows soft device placement and
# logs the placement decisions.
sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(
    allow_soft_placement=True,
    log_device_placement=True))