jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_training_only_device(self, strategy):
    # tf.distribute APIs may enter device scopes, but the saved model should not
    # contain device annotations, since devices during training may not be
    # available when the saved model is used.
    #
    # Models trained with MultiWorkerMirroredStrategy is affected the most,
    # since under MultiWorkerMirroredStrategy the device annotations contain job
    # and task, which causes error if that job or task is not available even
    # with soft placement enabled.

    class Model(tf.Module):

      @tf.function(input_signature=[])
      def __call__(self):
        return tf.identity(1.)

    export_dir = self.get_temp_dir()
    with strategy.scope(), tf.device("GPU:0"):
      m = Model()
      tf.saved_model.save(m, export_dir)

    export_dir = self.get_temp_dir()
    loaded = tf.saved_model.load(export_dir)
    graph = loaded.signatures["serving_default"].graph
    for op in graph.get_operations():
      self.assertEqual(op.device, "")
    self.assertEqual(loaded().numpy(), 1.)
//  [python]
def predict(self, input_data: PandasData) -> PandasData:
        """Makes prediction on input data through a trained model.

        Unlike in training and validation data, the categorical columns in
        prediction input data can have dtypes either 'category' or 'int', with
        'int' being numpy.int64 in pandas DataFrame.


        Args:
            input_data (pandas.DataFrame):
                Required. An input Pandas DataFrame containing data for
                prediction. It will be preprocessed into a dictionary as the
                input for to the trained model.
        Returns:
            Prediction results in the format of pandas DataFrame.
        """
        try:
            import tensorflow as tf
        except ImportError:
            raise ImportError(
                "TensorFlow must be installed to make predictions."
            ) from None

        if self.model is None:
            if not hasattr(self, "output_model_path") or self.output_model_path is None:
                raise ValueError("No trained model. Please call .fit first.")
            self.model = tf.compat.v1.saved_model.load(self.output_model_path)

        prediction_inputs = {}
        for col in input_data.columns:
            if input_data[col].dtypes == "category":
                dtype = tf.int64
            else:
                dtype = tf.dtypes.as_dtype(input_data[col].dtypes)
            prediction_inputs[col] = tf.constant(input_data[col].to_list(), dtype=dtype)
        prediction_outputs = self.model.signatures["serving_default"](
            **prediction_inputs
        )
        if self._binding["model_type"] == _CLASSIFICATION:
            predicted_labels = []
            for score, labels in zip(
                prediction_outputs["scores"].numpy(),
                prediction_outputs["classes"].numpy().astype(int),
            ):
                predicted_labels.append(labels[score.argmax()])
            return pd.DataFrame({self._binding["target_column"]: predicted_labels})
        elif self._binding["model_type"] == _REGRESSION:
            return pd.DataFrame(
                {
                    self._binding["target_column"]: prediction_outputs["value"]
                    .numpy()
                    .reshape(-1)
                }
            )
        else:
            raise ValueError(f"Unsupported model type: {self._binding['model_type']}.")
//  [python]
def predict(self, input_data: PandasData) -> PandasData:
        """Makes prediction on input data through a trained model.

        Unlike in training and validation data, the categorical columns in
        prediction input data can have dtypes either 'category' or 'int', with
        'int' being numpy.int64 in pandas DataFrame.


        Args:
            input_data (pandas.DataFrame):
                Required. An input Pandas DataFrame containing data for
                prediction. It will be preprocessed into a dictionary as the
                input for to the trained model.
        Returns:
            Prediction results in the format of pandas DataFrame.
        """
        try:
            import tensorflow as tf
        except ImportError:
            raise ImportError(
                "TensorFlow must be installed to make predictions."
            ) from None

        if self.model is None:
            if not hasattr(self, "output_model_path") or self.output_model_path is None:
                raise ValueError("No trained model. Please call .fit first.")
            self.model = tf.saved_model.load(self.output_model_path)

        prediction_inputs = {}
        for col in input_data.columns:
            if input_data[col].dtypes == "category":
                dtype = tf.int64
            else:
                dtype = tf.dtypes.as_dtype(input_data[col].dtypes)
            prediction_inputs[col] = tf.constant(input_data[col].to_list(), dtype=dtype)
        prediction_outputs = self.model.signatures["serving_default"](
            **prediction_inputs
        )
        if self._binding["model_type"] == _CLASSIFICATION:
            predicted_labels = []
            for score, labels in zip(
                prediction_outputs["scores"].numpy(),
                prediction_outputs["classes"].numpy().astype(int),
            ):
                predicted_labels.append(labels[score.argmax()])
            return pd.DataFrame({self._binding["target_column"]: predicted_labels})
        elif self._binding["model_type"] == _REGRESSION:
            return pd.DataFrame(
                {
                    self._binding["target_column"]: prediction_outputs["value"]
                    .numpy()
                    .reshape(-1)
                }
            )
        else:
            raise ValueError(f"Unsupported model type: {self._binding['model_type']}.")
//  [python]
def test_read_mirrored_variable(self, strategy):

    class Model(tf.Module):

      def __init__(self):
        self.v = tf.Variable(
            0., synchronization=tf.VariableSynchronization.ON_WRITE)

      @tf.function(input_signature=[])
      def __call__(self):
        return self.v.read_value()

    export_dir = self.get_temp_dir()
    with strategy.scope():
      m = Model()
      m.v.assign(1.)
      tf.saved_model.save(m, export_dir)

    with strategy.scope():
      loaded = tf.compat.v1.saved_model.load(export_dir)
    self.assertAllEqual(
        self.evaluate(
            strategy.experimental_local_results(strategy.run(loaded))),
        [1., 1.])
//  [python]
def test_read_mirrored_variable(self, strategy):

    class Model(tf.Module):

      def __init__(self):
        self.v = tf.Variable(
            0., synchronization=tf.VariableSynchronization.ON_WRITE)

      @tf.function(input_signature=[])
      def __call__(self):
        return self.v.read_value()

    export_dir = self.get_temp_dir()
    with strategy.scope():
      m = Model()
      m.v.assign(1.)
      tf.saved_model.save(m, export_dir)

    with strategy.scope():
      loaded = tf.saved_model.load(export_dir)
    self.assertAllEqual(
        self.evaluate(
            strategy.experimental_local_results(strategy.run(loaded))),
        [1., 1.])
//  [python]
def test_model_with_loaded_layer(self, strategy):
    # When a model is loaded under strategy, we wrap it so that when it's passed
    # to strategy.run(), the captured variables resolve to the ones of the
    # current replica. Since the saved tf.function may contain updates to the
    # variables, we don't allow using the model outside of strategy.run().
    #
    # That is to say, a loaded model is different from the original Python one.
    # We need to test save-load-save-load to make sure things work correctly.

    class Layer(tf.Module):

      def __init__(self):
        self.v = tf.Variable(1.)

      @tf.function(input_signature=[])
      def __call__(self):
        return self.v.read_value()

    class Model(tf.Module):

      def __init__(self, layer):
        self.layer = layer

      @tf.function(input_signature=[])
      def __call__(self):
        return self.layer()

    layer_export_dir = self.get_temp_dir()
    tf.saved_model.save(Layer(), layer_export_dir)

    with strategy.scope():
      m = Model(tf.saved_model.load(layer_export_dir))
      export_dir = self.get_temp_dir()
      # Saving a ConcreteFunction should raise an error.
      with self.assertRaisesRegex(
          ValueError, "saving a tf.function with input_signature instead"):
        tf.saved_model.save(
            m,
            export_dir,
            signatures={
                "call": m.__call__.get_concrete_function(),
            })
      tf.saved_model.save(m, export_dir)

    loaded = tf.saved_model.load(export_dir)
    self.assertAllEqual(
        self.evaluate(
            strategy.experimental_local_results(strategy.run(loaded))),
        [1., 1.])
//  [python]
def test_load_with_partitioner_works(self):
    model = self.Model()
    model_dir = self.get_temp_dir()
    tf.saved_model.save(model, model_dir)

    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
        self.cluster_resolver, tf1.fixed_size_partitioner(2))
    with strategy.scope():
      tf.compat.v1.saved_model.load(model_dir)
//  [python]
def test_load_with_partitioner_works(self):
    model = self.Model()
    model_dir = self.get_temp_dir()
    tf.saved_model.save(model, model_dir)

    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
        self.cluster_resolver, tf1.fixed_size_partitioner(2))
    with strategy.scope():
      tf.saved_model.load(model_dir)
//  [python]
def test_basic(self):
    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
        self.cluster_resolver)
    model_dir = self.get_temp_dir()
    with strategy.scope():
      m = self.Model()
    m.train()
    tf.saved_model.save(m, model_dir)

    # Load via V2 API.
    loaded = tf.compat.v2.saved_model.load(model_dir)
    self.assertRegex(loaded.v1.device, "/job:chief/replica:0/task:0")
    self.assertRegex(loaded.v2.device, "/job:chief/replica:0/task:0")
    self.assertAllEqual(loaded(tf.identity(1)), [6, 6, 6, 6])
    loaded.v2.assign([1, 1, 1, 1])
    self.assertAllEqual(loaded(tf.identity(1)), [4, 4, 4, 4])

    # Load via V1 API.
    self.assertAllEqual(self.load_and_run_v1(model_dir, {"x": 1}), [6, 6, 6, 6])
//  [python]
def test_basic(self):
    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
        self.cluster_resolver)
    model_dir = self.get_temp_dir()
    with strategy.scope():
      m = self.Model()
    m.train()
    tf.saved_model.save(m, model_dir)

    # Load via V2 API.
    loaded = tf.saved_model.load(model_dir)
    self.assertRegex(loaded.v1.device, "/job:chief/replica:0/task:0")
    self.assertRegex(loaded.v2.device, "/job:chief/replica:0/task:0")
    self.assertAllEqual(loaded(tf.identity(1)), [6, 6, 6, 6])
    loaded.v2.assign([1, 1, 1, 1])
    self.assertAllEqual(loaded(tf.identity(1)), [4, 4, 4, 4])

    # Load via V1 API.
    self.assertAllEqual(self.load_and_run_v1(model_dir, {"x": 1}), [6, 6, 6, 6])
//  [python]
def test_load_to_same_strategy(self):
    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
        self.cluster_resolver)
    model_dir = self.get_temp_dir()
    with strategy.scope():
      m = self.Model()
    m.train()
    tf.saved_model.save(m, model_dir)

    with strategy.scope():
      loaded = tf.compat.v2.saved_model.load(model_dir)

    # Make sure that the variables are created on different devices. SavedModel
    # may load the variables in a different order compared to the creation order
    # so the devices may not be exactly the same as before.
    self.assertTrue(("/job:ps/replica:0/task:0" in loaded.v1.device and
                     "/job:ps/replica:0/task:1" in loaded.v2.device) or
                    ("/job:ps/replica:0/task:1" in loaded.v1.device and
                     "/job:ps/replica:0/task:0" in loaded.v2.device))
    self.assertAllEqual(loaded(tf.identity(1)), [6, 6, 6, 6])
//  [python]
def test_load_to_same_strategy(self):
    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
        self.cluster_resolver)
    model_dir = self.get_temp_dir()
    with strategy.scope():
      m = self.Model()
    m.train()
    tf.saved_model.save(m, model_dir)

    with strategy.scope():
      loaded = tf.saved_model.load(model_dir)

    # Make sure that the variables are created on different devices. SavedModel
    # may load the variables in a different order compared to the creation order
    # so the devices may not be exactly the same as before.
    self.assertTrue(("/job:ps/replica:0/task:0" in loaded.v1.device and
                     "/job:ps/replica:0/task:1" in loaded.v2.device) or
                    ("/job:ps/replica:0/task:1" in loaded.v1.device and
                     "/job:ps/replica:0/task:0" in loaded.v2.device))
    self.assertAllEqual(loaded(tf.identity(1)), [6, 6, 6, 6])
//  [python]
def test_load_to_different_strategy(self):
    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
        self.cluster_resolver)
    model_dir = self.get_temp_dir()
    with strategy.scope():
      m = self.Model()
    m.train()
    tf.saved_model.save(m, model_dir)

    del m  # Garbage collect variables before we reset the context.
    context._reset_context()

    mirrored_strategy = tf.distribute.MirroredStrategy(devices=["CPU:0"])
    with mirrored_strategy.scope():
      loaded = tf.compat.v2.saved_model.load(model_dir)
    self.assertIsInstance(loaded.v1, values.DistributedVariable)
    self.assertAllEqual(loaded(tf.identity(1)), [6, 6, 6, 6])
//  [python]
def test_load_to_different_strategy(self):
    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
        self.cluster_resolver)
    model_dir = self.get_temp_dir()
    with strategy.scope():
      m = self.Model()
    m.train()
    tf.saved_model.save(m, model_dir)

    del m  # Garbage collect variables before we reset the context.
    context._reset_context()

    mirrored_strategy = tf.distribute.MirroredStrategy(devices=["CPU:0"])
    with mirrored_strategy.scope():
      loaded = tf.saved_model.load(model_dir)
    self.assertIsInstance(loaded.v1, values.DistributedVariable)
    self.assertAllEqual(loaded(tf.identity(1)), [6, 6, 6, 6])
//  [python]
def setup(self):
        def initialize_model():
            import tensorflow as tf
            # Load a potentially large model in memory. Executed once per process.
            return TfModel(tf.compat.v1.saved_model.load(self._engine_path, ["serve"]))

        self._model = self._shared_handle.acquire(initialize_model)
//  [python]
def setup(self):
        def initialize_model():
            import tensorflow as tf
            # Load a potentially large model in memory. Executed once per process.
            return TfModel(tf.saved_model.load(self._engine_path, ["serve"]))

        self._model = self._shared_handle.acquire(initialize_model)
//  [python]
def test_update_mirrored_variable(self, strategy):
    # This is also uncommon since most model parameters should be updated by
    # optimizer, and this test case is for completeness.
    #
    # In the cross replica context, assigning to the variable assigns the same
    # value to all replicas. This is true with the loaded model as well.
    #
    # However in replica context, MirroredVariable (synchronization=ON_WRITE)
    # in a loaded model behaves differently. Updating MirroredVariable only
    # update the current replica's variable with the current replica's value.
    # There's no aggregation. This doesn't affect variables that are updated
    # through optimizer. This is work as intended but can be surprising.

    class Model(tf.Module):

      def __init__(self):
        self.v = tf.Variable(
            0.,
            synchronization=tf.VariableSynchronization.ON_WRITE,
            aggregation=tf.VariableAggregation.MEAN)

      @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.float32)])
      def update(self, value):
        self.v.assign_add(value)

    export_dir = self.get_temp_dir()
    # value = strategy.experimental_distribute_values_from_function(
    #     lambda ctx: tf.identity([1., 2.][ctx.replica_id_in_sync_group]))
    with strategy.scope():
      m = Model()
      tf.saved_model.save(m, export_dir)
      del m

    with strategy.scope():
      loaded = tf.compat.v2.saved_model.load(export_dir)
    self.assertAllEqual(
        self.evaluate(strategy.experimental_local_results(loaded.v)), [0., 0.])
    self.evaluate(loaded.v.assign(1.))
    self.assertAllEqual(
        self.evaluate(strategy.experimental_local_results(loaded.v)), [1., 1.])
//  [python]
def test_update_mirrored_variable(self, strategy):
    # This is also uncommon since most model parameters should be updated by
    # optimizer, and this test case is for completeness.
    #
    # In the cross replica context, assigning to the variable assigns the same
    # value to all replicas. This is true with the loaded model as well.
    #
    # However in replica context, MirroredVariable (synchronization=ON_WRITE)
    # in a loaded model behaves differently. Updating MirroredVariable only
    # update the current replica's variable with the current replica's value.
    # There's no aggregation. This doesn't affect variables that are updated
    # through optimizer. This is work as intended but can be surprising.

    class Model(tf.Module):

      def __init__(self):
        self.v = tf.Variable(
            0.,
            synchronization=tf.VariableSynchronization.ON_WRITE,
            aggregation=tf.VariableAggregation.MEAN)

      @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.float32)])
      def update(self, value):
        self.v.assign_add(value)

    export_dir = self.get_temp_dir()
    # value = strategy.experimental_distribute_values_from_function(
    #     lambda ctx: tf.identity([1., 2.][ctx.replica_id_in_sync_group]))
    with strategy.scope():
      m = Model()
      tf.saved_model.save(m, export_dir)
      del m

    with strategy.scope():
      loaded = tf.saved_model.load(export_dir)
    self.assertAllEqual(
        self.evaluate(strategy.experimental_local_results(loaded.v)), [0., 0.])
    self.evaluate(loaded.v.assign(1.))
    self.assertAllEqual(
        self.evaluate(strategy.experimental_local_results(loaded.v)), [1., 1.])
//  [python]
def test_model_with_loaded_v1_layer_broken(self, strategy):
    # If a model contains a layer loaded from SavedModel, and if the model is
    # loaded under tf.distribute.Strategy, it can't be saved and loaded again
    # under tf.distribute.Strategy.
    #
    # Although the error is the same models with TF2 SavedModel, the cause is
    # different. TF1 models loaded in API contain an intializer, which is
    # invoked upon loading. Since loading is in the cross-replica context, that
    # fails.
    #
    # Note that these saved model can still be loaded and used without
    # tf.distribute.Strategy.
    #
    # Some tf.hub layers are converted from TF1, by loading TF1 saved model in
    # TF2 then saved in TF2. This issue disables them to work with
    # tf.distribute.Strategy.
    v1_export_dir = self.get_temp_dir()

    with tf1.Graph().as_default(), tf1.Session() as sess:
      v = tf1.Variable(1., use_resource=True)
      sess.run(v.initializer)
      builder = tf1.saved_model.Builder(v1_export_dir)
      builder.add_meta_graph_and_variables(
          sess,
          tags=[tf1.saved_model.tag_constants.TRAINING],
          main_op=tf1.tables_initializer(),
          strip_default_attrs=True)
      builder.save()

    v1_loaded = tf.saved_model.load(v1_export_dir)
    v2_export_dir = self.get_temp_dir()
    tf.saved_model.save(v1_loaded, v2_export_dir)
    with strategy.scope():
      # TODO(b/150009657): remove after fix.
      # got error, want no error.
      with self.assertRaisesRegex(
          tf.errors.InvalidArgumentError,
          "from the cross-replica context in an in-replica context"):
        tf.compat.v1.saved_model.load(v2_export_dir)
//  [python]
def test_model_with_loaded_v1_layer_broken(self, strategy):
    # If a model contains a layer loaded from SavedModel, and if the model is
    # loaded under tf.distribute.Strategy, it can't be saved and loaded again
    # under tf.distribute.Strategy.
    #
    # Although the error is the same models with TF2 SavedModel, the cause is
    # different. TF1 models loaded in API contain an intializer, which is
    # invoked upon loading. Since loading is in the cross-replica context, that
    # fails.
    #
    # Note that these saved model can still be loaded and used without
    # tf.distribute.Strategy.
    #
    # Some tf.hub layers are converted from TF1, by loading TF1 saved model in
    # TF2 then saved in TF2. This issue disables them to work with
    # tf.distribute.Strategy.
    v1_export_dir = self.get_temp_dir()

    with tf1.Graph().as_default(), tf1.Session() as sess:
      v = tf1.Variable(1., use_resource=True)
      sess.run(v.initializer)
      builder = tf1.saved_model.Builder(v1_export_dir)
      builder.add_meta_graph_and_variables(
          sess,
          tags=[tf1.saved_model.tag_constants.TRAINING],
          main_op=tf1.tables_initializer(),
          strip_default_attrs=True)
      builder.save()

    v1_loaded = tf.saved_model.load(v1_export_dir)
    v2_export_dir = self.get_temp_dir()
    tf.saved_model.save(v1_loaded, v2_export_dir)
    with strategy.scope():
      # TODO(b/150009657): remove after fix.
      # got error, want no error.
      with self.assertRaisesRegex(
          tf.errors.InvalidArgumentError,
          "from the cross-replica context in an in-replica context"):
        tf.saved_model.load(v2_export_dir)
//  [python]
def load_model_predict_fn(export_dir: str) -> Any:
  """Load PointPillars model from saved directory.

  Args:
    export_dir: Export directory path.
  Returns:
    predict_fn: A function can be run for model inference.
  """
  logging.info('Loading model from %s.', export_dir)
  model = tf.compat.v1.saved_model.load(export_dir)
  predict_fn = model.signatures[
      tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
  return predict_fn
//  [python]
def load_model_predict_fn(export_dir: str) -> Any:
  """Load PointPillars model from saved directory.

  Args:
    export_dir: Export directory path.
  Returns:
    predict_fn: A function can be run for model inference.
  """
  logging.info('Loading model from %s.', export_dir)
  model = tf.saved_model.load(export_dir)
  predict_fn = model.signatures[
      tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
  return predict_fn
//  [python]
def SaveAndLoadModel(model: tf.Module,
                     save_gradients=True) -> tf.Module:
  # Roundtrip through saved model on disk.
  model_dir = os.path.join(absltest.get_default_test_tmpdir(), str(id(model)))
  tf.saved_model.save(
      model, model_dir,
      options=tf.saved_model.SaveOptions(experimental_custom_gradients=save_gradients))
  restored_model = tf.compat.v1.saved_model.load(model_dir)
  return restored_model
//  [python]
def SaveAndLoadModel(model: tf.Module,
                     save_gradients=True) -> tf.Module:
  # Roundtrip through saved model on disk.
  model_dir = os.path.join(absltest.get_default_test_tmpdir(), str(id(model)))
  tf.saved_model.save(
      model, model_dir,
      options=tf.saved_model.SaveOptions(experimental_custom_gradients=save_gradients))
  restored_model = tf.saved_model.load(model_dir)
  return restored_model
//  [python]
def _get_graph_func(saved_model_dir):
    saved_model_loaded = tf.saved_model.load(
        saved_model_dir, tags=[tag_constants.SERVING])
    graph_func = saved_model_loaded.signatures[
        signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
    graph_func = convert_to_constants.convert_variables_to_constants_v2(graph_func)
    return graph_func
//  [python]
labelNames = [
    "apple",
    "ash gourd",
    "beet root",
    "bitter gourd",
    "bottle gourd",
    "brinjal",
    "broccoli",
    "cabbage",
    "capsicum",
    "carrot",
    "cauliflower",
    "Chicken",
    "Coriender",
    "corn",
    "Cottage Cheese",
    "cucumber",
    "Egg",
    "Elephant Yum",
    "Fish",
    "garlic",
    "Ginger",
    "Green Onion",
    "Kohlrabi",
    "lady finger",
    "Lemon",
    "Lettuce",
    "Milk",
    "Mint",
    "mushroom",
    "onion",
    "Papaya",
    "peas",
    "potato",
    "pumpkin",
    "Raddish",
    "Raw Banana",
    "Raw Jackfruit",
    "spinach",
    "tomato",
    "Turnip",
]

PROB_THRESHOLD = 0.2
model_path = os.path.join(pathlib.Path(__file__).parent.parent, "data")
model = tf.saved_model.load(model_path)
//  [python]
labelNames = [
    "apple",
    "ash gourd",
    "beet root",
    "bitter gourd",
    "bottle gourd",
    "brinjal",
    "broccoli",
    "cabbage",
    "capsicum",
    "carrot",
    "cauliflower",
    "Chicken",
    "Coriender",
    "corn",
    "Cottage Cheese",
    "cucumber",
    "Egg",
    "Elephant Yum",
    "Fish",
    "garlic",
    "Ginger",
    "Green Onion",
    "Kohlrabi",
    "lady finger",
    "Lemon",
    "Lettuce",
    "Milk",
    "Mint",
    "mushroom",
    "onion",
    "Papaya",
    "peas",
    "potato",
    "pumpkin",
    "Raddish",
    "Raw Banana",
    "Raw Jackfruit",
    "spinach",
    "tomato",
    "Turnip",
]

PROB_THRESHOLD = 0.2
model_path = os.path.join(pathlib.Path(__file__).parent.parent, "data")
model = tensorflow.saved_model.load(model_path)
//  [python]
def setup(self):
        def initialize_model():
            import tensorflow as tf
            import os
            src='/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcusolver.so.11'
            dst='/usr/local/lib/python3.6/dist-packages/tensorflow/python/libcusolver.so.10'
            try: 
                os.symlink(src, dst)
                physical_devices = tf.config.list_physical_devices('GPU')
                tf.config.experimental.set_memory_growth(physical_devices[0], True)
            except:
                pass
            # Load a potentially large model in memory. Executed once per process.
            return TfModel(tf.compat.v1.saved_model.load(self._engine_path, ["serve"]))

        self._model = self._shared_handle.acquire(initialize_model)
//  [python]
def setup(self):
        def initialize_model():
            import tensorflow as tf
            import os
            src='/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcusolver.so.11'
            dst='/usr/local/lib/python3.6/dist-packages/tensorflow/python/libcusolver.so.10'
            try: 
                os.symlink(src, dst)
                physical_devices = tf.config.list_physical_devices('GPU')
                tf.config.experimental.set_memory_growth(physical_devices[0], True)
            except:
                pass
            # Load a potentially large model in memory. Executed once per process.
            return TfModel(tf.saved_model.load(self._engine_path, ["serve"]))

        self._model = self._shared_handle.acquire(initialize_model)
//  [python]
def _load_tf1_estimator_saved_model(tf_saved_model_dir, tf_meta_graph_tags, tf_signature_def_key):
    """
    Load a specified TensorFlow model consisting of a TensorFlow metagraph and signature definition
    from a serialized TensorFlow ``SavedModel`` collection.

    Args:
        tf_saved_model_dir: The local filesystem path or run-relative artifact path to the model.
        tf_meta_graph_tags: A list of tags identifying the model's metagraph within the
            serialized ``SavedModel`` object. For more information, see the
            ``tags`` parameter of the `tf.saved_model.builder.SavedModelBuilder
            method <https://www.tensorflow.org/api_docs/python/tf/saved_model/
            builder/SavedModelBuilder#add_meta_graph>`_.
        tf_signature_def_key: A string identifying the input/output signature associated with the
            model. This is a key within the serialized ``SavedModel``'s
            signature definition mapping. For more information, see the
            ``signature_def_map`` parameter of the
            ``tf.saved_model.builder.SavedModelBuilder`` method.

    Returns:
        A callable graph (tensorflow.function) that takes inputs and returns inferences.
    """
    import tensorflow as tf

    with tf.compat.v1.Session(graph=tf.Graph()) as sess:
        tf.saved_model.loader.load(sess, tf_meta_graph_tags, tf_saved_model_dir)
        loaded_sig = sess.graph.get_collection(tf.saved_model.signature_constants.SIGNATURES_KEY)
        if tf_signature_def_key not in loaded_sig:
            raise MlflowException(
                f"Could not find signature def key {tf_signature_def_key}. "
                f"Available keys are: {list(loaded_sig.keys())}"
            )
        return loaded_sig[tf_signature_def_key]
//  [python]
def _load_tf1_estimator_saved_model(tf_saved_model_dir, tf_meta_graph_tags, tf_signature_def_key):
    """
    Load a specified TensorFlow model consisting of a TensorFlow metagraph and signature definition
    from a serialized TensorFlow ``SavedModel`` collection.

    Args:
        tf_saved_model_dir: The local filesystem path or run-relative artifact path to the model.
        tf_meta_graph_tags: A list of tags identifying the model's metagraph within the
            serialized ``SavedModel`` object. For more information, see the
            ``tags`` parameter of the `tf.saved_model.builder.SavedModelBuilder
            method <https://www.tensorflow.org/api_docs/python/tf/saved_model/
            builder/SavedModelBuilder#add_meta_graph>`_.
        tf_signature_def_key: A string identifying the input/output signature associated with the
            model. This is a key within the serialized ``SavedModel``'s
            signature definition mapping. For more information, see the
            ``signature_def_map`` parameter of the
            ``tf.saved_model.builder.SavedModelBuilder`` method.

    Returns:
        A callable graph (tensorflow.function) that takes inputs and returns inferences.
    """
    import tensorflow as tf

    loaded = tf.saved_model.load(tags=tf_meta_graph_tags, export_dir=tf_saved_model_dir)
    loaded_sig = loaded.signatures
    if tf_signature_def_key not in loaded_sig:
        raise MlflowException(
            f"Could not find signature def key {tf_signature_def_key}. "
            f"Available keys are: {list(loaded_sig.keys())}"
        )
    return loaded_sig[tf_signature_def_key]
//  [python]
def test_export(self, input_type, rescale_output, input_image_size,
                  preserve_aspect_ratio):
    tmp_dir = self.get_temp_dir()
    module = self._get_segmentation_module(
        input_type=input_type,
        rescale_output=rescale_output,
        preserve_aspect_ratio=preserve_aspect_ratio)

    self._export_from_module(module, input_type, tmp_dir)

    self.assertTrue(os.path.exists(os.path.join(tmp_dir, 'saved_model.pb')))
    self.assertTrue(
        os.path.exists(os.path.join(tmp_dir, 'variables', 'variables.index')))
    self.assertTrue(
        os.path.exists(
            os.path.join(tmp_dir, 'variables',
                         'variables.data-00000-of-00001')))

    imported = tf.saved_model.load(tmp_dir)
    segmentation_fn = imported.signatures['serving_default']

    images = self._get_dummy_input(input_type, input_image_size)
    if input_type != 'tflite':
      processed_images, _ = tf.nest.map_structure(
          tf.stop_gradient,
          tf.map_fn(
              module._build_inputs,
              elems=tf.zeros((1, 112, 112, 3), dtype=tf.uint8),
              fn_output_signature=(tf.TensorSpec(
                  shape=[112, 112, 3], dtype=tf.float32),
                                   tf.TensorSpec(
                                       shape=[4, 2], dtype=tf.float32))))
    else:
      processed_images = images

    logits = module.model(processed_images, training=False)['logits']
    if rescale_output:
      expected_output = tf.image.resize(
          logits, input_image_size, method='bilinear')
    else:
      expected_output = tf.image.resize(logits, [112, 112], method='bilinear')
    out = segmentation_fn(tf.constant(images))
    self.assertAllClose(out['logits'].numpy(), expected_output.numpy())
//  [python]
def make_tflite_export(weights_path, export_dir):
  if os.path.exists(export_dir):
    log('TF-Lite export already exists in {}, skipping TF-Lite export'.format(
        export_dir))
    return

  # Create a TF-Lite compatible Module wrapper around YAMNet.
  log('Building and checking TF-Lite Module ...')
  params = yamnet_params.Params(tflite_compatible=True)
  yamnet = YAMNet(weights_path, params)
  check_model(yamnet, yamnet.class_map_path(), params)
  log('Done')

  # Make TF-Lite SavedModel export.
  log('Making TF-Lite SavedModel export ...')
  saved_model_dir = os.path.join(export_dir, 'saved_model')
  os.makedirs(saved_model_dir)
  tf.saved_model.save(
      yamnet, saved_model_dir,
      signatures={'serving_default': yamnet.__call__.get_concrete_function()})
  log('Done')

  # Check that the export can be loaded and works.
  log('Checking TF-Lite SavedModel export in TF2 ...')
  model = tf.compat.v1.saved_model.load(saved_model_dir)
  check_model(model, model.class_map_path(), params)
  log('Done')

  # Make a TF-Lite model from the SavedModel.
  log('Making TF-Lite model ...')
  tflite_converter = tf.lite.TFLiteConverter.from_saved_model(
      saved_model_dir, signature_keys=['serving_default'])
  tflite_model = tflite_converter.convert()
  tflite_model_path = os.path.join(export_dir, 'yamnet.tflite')
  with open(tflite_model_path, 'wb') as f:
    f.write(tflite_model)
  log('Done')

  # Check the TF-Lite export.
  log('Checking TF-Lite model ...')
  interpreter = tf.lite.Interpreter(tflite_model_path)
  runner = interpreter.get_signature_runner('serving_default')
  check_model(runner, 'yamnet_class_map.csv', params)
  log('Done')

  return saved_model_dir
//  [python]
def make_tflite_export(weights_path, export_dir):
  if os.path.exists(export_dir):
    log('TF-Lite export already exists in {}, skipping TF-Lite export'.format(
        export_dir))
    return

  # Create a TF-Lite compatible Module wrapper around YAMNet.
  log('Building and checking TF-Lite Module ...')
  params = yamnet_params.Params(tflite_compatible=True)
  yamnet = YAMNet(weights_path, params)
  check_model(yamnet, yamnet.class_map_path(), params)
  log('Done')

  # Make TF-Lite SavedModel export.
  log('Making TF-Lite SavedModel export ...')
  saved_model_dir = os.path.join(export_dir, 'saved_model')
  os.makedirs(saved_model_dir)
  tf.saved_model.save(
      yamnet, saved_model_dir,
      signatures={'serving_default': yamnet.__call__.get_concrete_function()})
  log('Done')

  # Check that the export can be loaded and works.
  log('Checking TF-Lite SavedModel export in TF2 ...')
  model = tf.saved_model.load(saved_model_dir)
  check_model(model, model.class_map_path(), params)
  log('Done')

  # Make a TF-Lite model from the SavedModel.
  log('Making TF-Lite model ...')
  tflite_converter = tf.lite.TFLiteConverter.from_saved_model(
      saved_model_dir, signature_keys=['serving_default'])
  tflite_model = tflite_converter.convert()
  tflite_model_path = os.path.join(export_dir, 'yamnet.tflite')
  with open(tflite_model_path, 'wb') as f:
    f.write(tflite_model)
  log('Done')

  # Check the TF-Lite export.
  log('Checking TF-Lite model ...')
  interpreter = tf.lite.Interpreter(tflite_model_path)
  runner = interpreter.get_signature_runner('serving_default')
  check_model(runner, 'yamnet_class_map.csv', params)
  log('Done')

  return saved_model_dir
//  [python]
def test_export(self):
    model = transformer.Transformer(self.params, name="transformer_v2")
    export_dir = self.get_temp_dir()
    batch_size = 5
    max_length = 6

    class SaveModule(tf.Module):

      def __init__(self, model):
        super(SaveModule, self).__init__()
        self.model = model

      @tf.function
      def serve(self, x):
        return self.model.call([x], training=False)

    save_module = SaveModule(model)
    tensor_shape = (None, None)
    sample_input = tf.zeros((batch_size, max_length), dtype=tf.int64)
    _ = save_module.serve(sample_input)
    signatures = dict(
        serving_default=save_module.serve.get_concrete_function(
            tf.TensorSpec(shape=tensor_shape, dtype=tf.int64, name="x")))
    tf.saved_model.save(save_module, export_dir, signatures=signatures)
    imported = tf.saved_model.load(export_dir)
    serving_fn = imported.signatures["serving_default"]
    all_outputs = serving_fn(sample_input)
    output = all_outputs["outputs"]
    output_shapes = output.shape.as_list()
    self.assertEqual(output_shapes[0], batch_size)
    self.assertEqual(output_shapes[1],
                     max_length + model.params["extra_decode_length"])
//  [python]
def test_preprocessor(self):
    tmp_dir = self.get_temp_dir()
    model = TestModel()
    inputs = tf.ones([2, 4], tf.float32)

    preprocess_fn = lambda inputs: 2 * inputs

    module = export_base_v2.ExportModule(
        params=None,
        input_signature=tf.TensorSpec(shape=[2, 4]),
        model=model,
        preprocessor=preprocess_fn)
    expected_output = model(preprocess_fn(inputs))
    ckpt_path = tf.train.Checkpoint(model=model).save(
        os.path.join(tmp_dir, 'ckpt'))
    export_dir = export_base.export(
        module, ['serving_default'],
        export_savedmodel_dir=tmp_dir,
        checkpoint_path=ckpt_path,
        timestamped=False)
    imported = tf.compat.v1.saved_model.load(export_dir)
    output = imported.signatures['serving_default'](inputs)
    print('output', output)
    self.assertAllClose(
        output['outputs'].numpy(), expected_output['outputs'].numpy())
//  [python]
def test_preprocessor(self):
    tmp_dir = self.get_temp_dir()
    model = TestModel()
    inputs = tf.ones([2, 4], tf.float32)

    preprocess_fn = lambda inputs: 2 * inputs

    module = export_base_v2.ExportModule(
        params=None,
        input_signature=tf.TensorSpec(shape=[2, 4]),
        model=model,
        preprocessor=preprocess_fn)
    expected_output = model(preprocess_fn(inputs))
    ckpt_path = tf.train.Checkpoint(model=model).save(
        os.path.join(tmp_dir, 'ckpt'))
    export_dir = export_base.export(
        module, ['serving_default'],
        export_savedmodel_dir=tmp_dir,
        checkpoint_path=ckpt_path,
        timestamped=False)
    imported = tf.saved_model.load(export_dir)
    output = imported.signatures['serving_default'](inputs)
    print('output', output)
    self.assertAllClose(
        output['outputs'].numpy(), expected_output['outputs'].numpy())
//  [python]
def testSavedModelSaveRestore(self):
    save_dir = os.path.join(self.get_temp_dir(), 'save_restore')
    save_path = os.path.join(tempfile.mkdtemp(prefix=save_dir), 'hash')

    # TODO(b/203097231) is there an alternative that is not __internal__?
    root = tf.__internal__.tracking.AutoTrackable()

    default_value = -1
    root.table = simple_hash_table.SimpleHashTable(
        tf.int64, tf.int64, default_value=default_value)

    @def_function.function(input_signature=[tf.TensorSpec((), tf.int64)])
    def lookup(key):
      return root.table.find(key)

    root.lookup = lookup

    root.table.insert(1, 100)
    root.table.insert(2, 200)
    root.table.insert(3, 300)
    self.assertEqual(root.lookup(2), 200)
    self.assertAllEqual(3, len(self.evaluate(root.table.export()[0])))
    tf.saved_model.save(root, save_path)

    del root
    loaded = tf.compat.v1.saved_model.load(save_path)
    self.assertEqual(loaded.lookup(2), 200)
    self.assertEqual(loaded.lookup(10), -1)
//  [python]
def testSavedModelSaveRestore(self):
    save_dir = os.path.join(self.get_temp_dir(), 'save_restore')
    save_path = os.path.join(tempfile.mkdtemp(prefix=save_dir), 'hash')

    # TODO(b/203097231) is there an alternative that is not __internal__?
    root = tf.__internal__.tracking.AutoTrackable()

    default_value = -1
    root.table = simple_hash_table.SimpleHashTable(
        tf.int64, tf.int64, default_value=default_value)

    @def_function.function(input_signature=[tf.TensorSpec((), tf.int64)])
    def lookup(key):
      return root.table.find(key)

    root.lookup = lookup

    root.table.insert(1, 100)
    root.table.insert(2, 200)
    root.table.insert(3, 300)
    self.assertEqual(root.lookup(2), 200)
    self.assertAllEqual(3, len(self.evaluate(root.table.export()[0])))
    tf.saved_model.save(root, save_path)

    del root
    loaded = tf.saved_model.load(save_path)
    self.assertEqual(loaded.lookup(2), 200)
    self.assertEqual(loaded.lookup(10), -1)
//  [python]
def _save_and_load(self, model: tf.keras.Model) -> tf.keras.Model:
    temp_dir = self.create_tempdir().full_path
    lib.save_model(model, temp_dir)
    return tf.compat.v1.saved_model.load(temp_dir)
//  [python]
def _save_and_load(self, model: tf.keras.Model) -> tf.keras.Model:
    temp_dir = self.create_tempdir().full_path
    lib.save_model(model, temp_dir)
    return tf.saved_model.load(temp_dir)
//  [python]
def test_postprocessor(self):
    tmp_dir = self.get_temp_dir()
    model = TestModel()
    inputs = tf.ones([2, 4], tf.float32)

    postprocess_fn = lambda logits: {'outputs': 2 * logits['outputs']}

    module = export_base_v2.ExportModule(
        params=None,
        model=model,
        input_signature=tf.TensorSpec(shape=[2, 4]),
        postprocessor=postprocess_fn)
    expected_output = postprocess_fn(model(inputs))
    ckpt_path = tf.train.Checkpoint(model=model).save(
        os.path.join(tmp_dir, 'ckpt'))
    export_dir = export_base.export(
        module, ['serving_default'],
        export_savedmodel_dir=tmp_dir,
        checkpoint_path=ckpt_path,
        timestamped=False)
    imported = tf.compat.v1.saved_model.load(export_dir)
    output = imported.signatures['serving_default'](inputs)
    self.assertAllClose(
        output['outputs'].numpy(), expected_output['outputs'].numpy())
//  [python]
def test_postprocessor(self):
    tmp_dir = self.get_temp_dir()
    model = TestModel()
    inputs = tf.ones([2, 4], tf.float32)

    postprocess_fn = lambda logits: {'outputs': 2 * logits['outputs']}

    module = export_base_v2.ExportModule(
        params=None,
        model=model,
        input_signature=tf.TensorSpec(shape=[2, 4]),
        postprocessor=postprocess_fn)
    expected_output = postprocess_fn(model(inputs))
    ckpt_path = tf.train.Checkpoint(model=model).save(
        os.path.join(tmp_dir, 'ckpt'))
    export_dir = export_base.export(
        module, ['serving_default'],
        export_savedmodel_dir=tmp_dir,
        checkpoint_path=ckpt_path,
        timestamped=False)
    imported = tf.saved_model.load(export_dir)
    output = imported.signatures['serving_default'](inputs)
    self.assertAllClose(
        output['outputs'].numpy(), expected_output['outputs'].numpy())
//  [python]
def test_load_with_partitioner_raises_error(self):
    model = self.Model()
    model_dir = self.get_temp_dir()
    tf.saved_model.save(model, model_dir)

    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
        self.cluster_resolver, tf1.fixed_size_partitioner(2))
    with self.assertRaises(errors_impl.InvalidArgumentError):
      with strategy.scope():
        tf.compat.v1.saved_model.load(model_dir)
//  [python]
def test_load_with_partitioner_raises_error(self):
    model = self.Model()
    model_dir = self.get_temp_dir()
    tf.saved_model.save(model, model_dir)

    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
        self.cluster_resolver, tf1.fixed_size_partitioner(2))
    with self.assertRaises(errors_impl.InvalidArgumentError):
      with strategy.scope():
        tf.saved_model.load(model_dir)
//  [python]
def test_standard_model_export(self, model_type):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        model = get_model(model_type)
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        export_lib.export_model(model, temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertAllClose(ref_output, revived_model.serve(ref_input))
        # Test with a different batch size
        revived_model.serve(tf.random.normal((6, 10)))
//  [python]
def test_model_with_rng_export(self, model_type):

        class RandomLayer(layers.Layer):
            def __init__(self):
                super().__init__()
                self.seed_generator = backend.random.SeedGenerator()

            def call(self, inputs):
                return inputs + random.uniform(
                    ops.shape(inputs), seed=self.seed_generator
                )

        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        model = get_model(model_type, layer_list=[RandomLayer()])
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        export_lib.export_model(model, temp_filepath)
        revived_model = tf.compat.v1.saved_model.load(temp_filepath)
        self.assertEqual(ref_output.shape, revived_model.serve(ref_input).shape)
        # Test with a different batch size
        revived_model.serve(tf.random.normal((6, 10)))
//  [python]
def test_model_with_rng_export(self, model_type):

        class RandomLayer(layers.Layer):
            def __init__(self):
                super().__init__()
                self.seed_generator = backend.random.SeedGenerator()

            def call(self, inputs):
                return inputs + random.uniform(
                    ops.shape(inputs), seed=self.seed_generator
                )

        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        model = get_model(model_type, layer_list=[RandomLayer()])
        ref_input = tf.random.normal((3, 10))
        ref_output = model(ref_input)

        export_lib.export_model(model, temp_filepath)
        revived_model = tf.saved_model.load(temp_filepath)
        self.assertEqual(ref_output.shape, revived_model.serve(ref_input).shape)
        # Test with a different batch size
        revived_model.serve(tf.random.normal((6, 10)))
//  [python]
def test_update_sync_on_read_variable(self, strategy):
    # Updating a synchronizaiton=ON_READ in the replica context should just
    # update the local value. Updating it in the cross replica context updates
    # each component of the variable. Both are true with a loaded model.
    #
    # Note that if assigning a variable whose aggregation=SUM in the cross
    # replica context, each replica is assigned with the value divided by the
    # number of replicas.

    class Model(tf.Module):

      def __init__(self):
        self.v = tf.Variable(
            0.,
            synchronization=tf.VariableSynchronization.ON_READ,
            aggregation=tf.VariableAggregation.SUM)

      @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.float32)])
      def update(self, value):
        self.v.assign_add(value)

    export_dir = self.get_temp_dir()
    value = strategy.experimental_distribute_values_from_function(
        lambda ctx: tf.identity([3., 7.][ctx.replica_id_in_sync_group]))
    with strategy.scope():
      m = Model()
      tf.saved_model.save(m, export_dir)
      self.evaluate(m.v.assign(10.))
      self.assertAllEqual(
          self.evaluate(strategy.experimental_local_results(m.v)), [5., 5.])
      del m
      # TODO(b/161488560): strategy.run doesn't work with tf.function with
      # input_signature.
      # self.evaluate(strategy.run(m.update, args=(value,)))
      # self.assertAllEqual(
      #     self.evaluate(strategy.experimental_local_results(m.v)), [8., 12.])

    with strategy.scope():
      loaded = tf.compat.v1.saved_model.load(export_dir)
      self.evaluate(loaded.v.assign(10.))
      self.assertAllEqual(
          self.evaluate(strategy.experimental_local_results(loaded.v)),
          [5., 5.])
      self.evaluate(strategy.run(loaded.update, args=(value,)))
      self.assertAllEqual(
          self.evaluate(strategy.experimental_local_results(loaded.v)),
          [8., 12.])