def apply_function_with_checkpoint(fn, inputs, checkpoint, include=None,
                                   exclude=None):
  """Applies a tensor-in-tensor-out function with variables to some `Tensor`s.

  Variable values are loaded from the given checkpoint path. Note that the
  input_tensor_func, together with the checkpoint, will be converted to an
  all-constants graph, so ops requiring graph collections, such as table lookup
  (which requires a table init op being added to TABLE_INITIALIZERS collection),
  are not supported.

  Note: This API can only be used when TF2 is disabled or
  `tft_beam.Context.force_tf_compat_v1=True`.

  Args:
    fn: A tensor-in-tensor-out function that may contain variables.
    inputs: A list of `Tensor`s to apply `fn` to.
    checkpoint: The checkpoint path to load variables from.
    include: An optional list/tuple of scope strings for filtering which
        variables from the VARIABLES collection to include. If None, all
        variables will be included.
    exclude: An optional list/tuple of scope strings for filtering which
        variables from the VARIABLES collection to exclude. If None, no
        variables will be excluded.

  Returns:
    A `Tensor` or list of `Tensor`s representing the application of `fn`.

  Raises:
    ValueError: if the input tensor-in-tensor-out function adds to
        TABLE_INITIALIZERS collections.
  """
  loaded_graph = tf.compat.v1.Graph()
  with loaded_graph.as_default():
    input_placeholders = [
        tf.compat.v1.placeholder(
            dtype=tensor.dtype, shape=tensor.shape, name=tensor.op.name)
        for tensor in inputs
    ]
    output = fn(*input_placeholders)
    if isinstance(output, tf.Tensor):
      output_tensors = [output]
      output_single_tensor = True
    else:
      output_tensors = output
      output_single_tensor = False

    # TODO(qimingj/kestert): Copy table initializers to the composed graph.
    if tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.TABLE_INITIALIZERS):
      raise ValueError('Models with table init ops are not supported.')

    vars_to_restore = _get_variables_to_restore(include=include,
                                                exclude=exclude)
    saver = tf.compat.v1.train.Saver(vars_to_restore)
    with tf.compat.v1.Session() as sess:
      saver.restore(sess, checkpoint)
      output_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(
          sess, loaded_graph.as_graph_def(),
          [tensor.op.name for tensor in output_tensors])

  input_map = {tensor.name: tensor for tensor in inputs}
  output_tensors = tf.import_graph_def(
      output_graph_def, input_map=input_map,
      return_elements=[tensor.name for tensor in output_tensors])

  if output_single_tensor:
    assert len(output_tensors) == 1
    return output_tensors[0]
  else:
    return output_tensors
Applies a tensor-in-tensor-out function with variables to some `Tensor`s.

  Variable values are loaded from the given checkpoint path. Note that the
  input_tensor_func, together with the checkpoint, will be converted to an
  all-constants graph, so ops requiring graph collections, such as table lookup
  (which requires a table init op being added to TABLE_INITIALIZERS collection),
  are not supported.

  Note: This API can only be used when TF2 is disabled or
  `tft_beam.Context.force_tf_compat_v1=True`.

  Args:
    fn: A tensor-in-tensor-out function that may contain variables.
    inputs: A list of `Tensor`s to apply `fn` to.
    checkpoint: The checkpoint path to load variables from.
    include: An optional list/tuple of scope strings for filtering which
        variables from the VARIABLES collection to include. If None, all
        variables will be included.
    exclude: An optional list/tuple of scope strings for filtering which
        variables from the VARIABLES collection to exclude. If None, no
        variables will be excluded.

  Returns:
    A `Tensor` or list of `Tensor`s representing the application of `fn`.

  Raises:
    ValueError: if the input tensor-in-tensor-out function adds to
        TABLE_INITIALIZERS collections.
  
tensorflow_transform/pretrained_models.py