code: def test_get_model(self):
      sc = OrcaContext.get_spark_context()
      rdd = sc.range(0, 100)
      spark = OrcaContext.get_spark_session()
  
      from pyspark.ml.linalg import DenseVector
  
      df = rdd.map(
          lambda x: (
              DenseVector(
                  np.random.randn(
                      1,
                  ).astype(np.float32)
              ),
              int(np.random.randint(0, 2, size=())),
          )
      ).toDF(["feature", "label"])
  
      config = {"lr": 0.2}
  
      try:
          temp_dir = tempfile.mkdtemp()
  
          trainer = Estimator.from_keras(
              model_creator=model_creator,
              verbose=True,
              config=config,
              workers_per_node=3,
              backend="spark",
          )
  
          trainer.fit(
              df,
              epochs=5,
              batch_size=4,
              steps_per_epoch=25,
              feature_cols=["feature"],
              label_cols=["label"],
              validation_data=df,
              validation_steps=1,
          )
  
          trainer.save(os.path.join(temp_dir, "cifar10.h5"))
          pre_model = trainer.get_model()
          if hasattr(pre_model.optimizer, "get_weights"):
              pre_model_weights = pre_model.optimizer.get_weights()
          else:
              pre_model_weights = [
                  var.numpy() for var in pre_model.optimizer.variables()
              ]
          after_res = trainer.predict(
              df, feature_cols=["feature"]
          ).collect()
          expect_res = np.concatenate(
              [part["prediction"] for part in after_res]
          )
  
          trainer.shutdown()
  
          est = Estimator.from_keras(
              verbose=True,
              config=config,
              workers_per_node=3,
              backend="spark",
          )
  
          est.load(os.path.join(temp_dir, "cifar10.h5"))
          after_model = est.get_model()
          if hasattr(after_model.optimizer, "get_weights"):
              after_model_weights = after_model.optimizer.get_weights()
          else:
              after_model_weights = [
                  var.numpy()
                  for var in after_model.optimizer.variables()
              ]
          est.save(os.path.join(temp_dir, "cifar10_new.h5"))
  
          # continous predicting
          after_res = est.predict(
              df, feature_cols=["feature"]
          ).collect()
          pred_res = np.concatenate(
              [part["prediction"] for part in after_res]
          )
          assert np.array_equal(expect_res, pred_res)
  
          for pre_tensor, after_tensor in list(
              zip(pre_model_weights, after_model_weights)
          ):
              assert np.allclose(pre_tensor, after_tensor)
      finally:
          shutil.rmtree(temp_dir)
  
imports: import os
  import pytest
  import shutil
  import tempfile
  from unittest import TestCase
  import numpy
  import tensorflow
  from bigdl.orca.learn.tf2 import Estimator
  from bigdl.orca import OrcaContext
label_desc: Use image node_id 9 for calling the TestTF2Estimator obj's underlying member method code with example usage: obj.test_get_model() without return types