tensorflow / tpu

Reference models and tools for Cloud TPUs.
https://cloud.google.com/tpu/
Apache License 2.0
5.2k stars 1.77k forks source link

Predict w/ restored model from 'model.ckpt.data-00000-of-00001', 'model.ckpt.index', 'model.ckpt.meta' #982

Open yihwan-kim opened 2 years ago

yihwan-kim commented 2 years ago

Hello. I tried to restore mobilenet model and efficientnet-lite which are not frozen model, so that I can use model to inference, retrain and/or transfer learning.

So I tried below using efficientnet-lite first.

import tensorflow.compat.v1 as tf
import os
import time

tf.compat.v1.disable_eager_execution()

def restore_ckpt(sess, model_dir, export_ckpt=None):

    """Restore variables from a given checkpoint.

    Args:
    sess: a tf session for restoring or exporting models.
    ckpt_path: the path of the checkpoint. Can be a file path or a folder path.
    export_ckpt: whether to export the restored model.
    """
    files = os.listdir(model_dir)

    if any(file for file in files if '.ckpt.' in file):

        if tf.io.gfile.isdir(model_dir):
            ckpt_path = tf.train.latest_checkpoint(model_dir)

        for file in files:
            if 'ckpt.meta' in file:
                meta = file
                meta_path = os.path.join(model_dir, meta)

        saver = tf.train.import_meta_graph(meta_path)

        # Restore all variables from ckpt.
        start_time = time.time()
        saver.restore(sess, ckpt_path)
        end_time = time.time()

        elapsed_time = end_time - start_time
        print('Restoring model took {} seconds'.format(elapsed_time))

    else:
        raise ValueError("ckpt do not exist")

for img_idx, image in enumerate(image_list):
  img = cv2.imread(input_dir + image)

  # preprocessing
  img = crop_image(img)
  input_data = resize_img(img, input_details[0]['shape'][1], input_details[0]['shape'][2], interpolation)
  input_data = preprocess_input(input_data, mode)

  input_tensor = tf.convert_to_tensor(input_data)

  ckpt_path = '/content/efficientnet-lite0/efficientnet-lite0/'
  with tf.Session() as sess:
        restore_ckpt(sess, ckpt_path, export_ckpt=None)
        graph = tf.get_default_graph()
        X = graph.get_tensor_by_name(INPUT_TENSOR_NAME)
        y = graph.get_tensor_by_name(OUTPUT_TENSOR_NAME)

        output_data = sess.run(y, feed_dict={X:input_data})

To execute 'sess.run()',as far as I know, input and output tensor name should be given. I tried below to see the tensor name. However, I do not know which one is INPUT_TENSOR_NAME, or OUPUT_TENSOR_NAME, because it is not a network I created .

all_tensors = [tensor for op in tf.get_default_graph().get_operations() for tensor in op.values()]

Then, I tried to use tensorboard, but the graph keeps disappearing

trained_checkpoint_prefix = '/content/efficientnet-lite0/efficientnet-lite0/model.ckpt'

export_dir = os.path.join('export_dir', '0')
!rm -rf {export_dir}

graph = tf.Graph()
with tf.compat.v1.Session(graph=graph) as sess:
    # Restore from checkpoint
    loader = tf.compat.v1.train.import_meta_graph(trained_checkpoint_prefix + '.meta')
    loader.restore(sess, trained_checkpoint_prefix)

    # Export checkpoint to SavedModel
    builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(export_dir)
    builder.add_meta_graph_and_variables(sess,
                                         [tf.compat.v1.saved_model.TRAINING, tf.compat.v1.saved_model.SERVING],
                                         strip_default_attrs=True)
    builder.save()   
tf.enable_eager_execution()

%load_ext tensorboard
%tensorboard --logdir=./output/

Is there any other method to restore and run models from model.ckpt.data-00000-of-00001', 'model.ckpt.index', 'model.ckpt.meta' ?