Repository: neo-ai/neo-ai-dlr (language: python)

Issue:
Add link to API doc



Discussion / hints:

Fix patch:
diff --git a/CMakeLists.txt b/CMakeLists.txt
index dee1ce16e..85256a268 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -230,6 +230,12 @@ if(USE_TENSORRT)
 
     set(USE_TENSORRT OFF)
 endif()
+if(WITH_TENSORFLOW2_LIB)
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DDLR_TENSORFLOW2")
+    include_directories("${WITH_TENSORFLOW2_LIB}/include")
+    list(APPEND DLR_SRC "src/dlr_tensorflow2/dlr_tensorflow2.cc")
+    list(APPEND DLR_LINKER_LIBS -L${WITH_TENSORFLOW2_LIB}/lib -ltensorflow)
+endif()
 if(WITH_HEXAGON)
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DDLR_HEXAGON")
     list(APPEND DLR_SRC "src/dlr_hexagon/dlr_hexagon.cc")
@@ -391,7 +397,10 @@ if(DLR_BUILD_TESTS AND NOT(AAR_BUILD))
         ${CMAKE_CURRENT_SOURCE_DIR}/tests/cpp/dlr_data_transform_test.cc
         ${CMAKE_CURRENT_SOURCE_DIR}/tests/cpp/dlr_pipeline_skl_xgb_test.cc)
   endif()
-
+  if(WITH_TENSORFLOW2_LIB)
+    file(GLOB TENSORFLOW2_TEST_SRCS tests/cpp/dlr_tensorflow2/*.cc)
+    list(APPEND TEST_SRCS ${TENSORFLOW2_TEST_SRCS})
+  endif()
   if(WITH_HEXAGON)
     file(GLOB HEXAGON_TEST_SRCS tests/cpp/dlr_hexagon/*.cc)
     list(APPEND TEST_SRCS ${HEXAGON_TEST_SRCS})
@@ -595,7 +604,24 @@ if(DLR_BUILD_TESTS AND NOT(AAR_BUILD))
                     WORKING_DIRECTORY ${INVERSELABEL_STATIC_MODEL_DIR})
     file(REMOVE /tmp/${INVERSELABEL_STATIC_MODEL})
   endif()
-
+  if(WITH_TENSORFLOW2_LIB)
+    set(TF2_SAVED_MODEL imagenet_mobilenet_v2_100_224_classification.tar.gz)
+    set(TF2_SAVED_MODEL_DIR ${CMAKE_CURRENT_BINARY_DIR}/imagenet_mobilenet_v2)
+    if(NOT IS_DIRECTORY ${TF2_SAVED_MODEL_DIR})
+      file(MAKE_DIRECTORY ${TF2_SAVED_MODEL_DIR})
+      # Download Test Tensorflow model
+      download_file(
+        https://neo-ai-dlr-test-artifacts.s3-us-west-2.amazonaws.com/tf2-models/${TF2_SAVED_MODEL}
+        /tmp/${TF2_SAVED_MODEL}
+        SHA1
+        2d183d967a671af1498261fdd050caa487d88e44
+      )
+      # this is OS-agnostic
+      execute_process(COMMAND ${CMAKE_COMMAND} -E tar -xf /tmp/${TF2_SAVED_MODEL}
+                      WORKING_DIRECTORY ${TF2_SAVED_MODEL_DIR})
+      #file(REMOVE /tmp/${TF2_SAVED_MODEL})
+    endif()
+  endif() # WITH_TENSORFLOW2_LIB
   if(WITH_HEXAGON)
       # Download Test Hexagon model for Android 64 aarch64
       file(MAKE_DIRECTORY dlr_hexagon_model)
diff --git a/doc/install.rst b/doc/install.rst
index 5e640cab4..22f8cfd39 100644
--- a/doc/install.rst
+++ b/doc/install.rst
@@ -30,7 +30,7 @@ Building DLR consists of two steps:
 .. note:: Use of Git submodules
 
   DLR uses Git submodules to manage dependencies. So when you clone the repo, remember to specify ``--recursive`` option:
-  
+
   .. code-block:: bash
 
     git clone --recursive https://github.com/neo-ai/neo-ai-dlr
@@ -105,7 +105,7 @@ DLR requires CMake 3.13 or greater. First, we will build CMake from source.
 Now, build DLR.
 
 .. code-block:: bash
- 
+
   git clone --recursive https://github.com/neo-ai/neo-ai-dlr
   cd neo-ai-dlr
   mkdir build
@@ -132,7 +132,7 @@ If you have a system install of TensorRT via Deb or RPM package, you can instead
   cd neo-ai-dlr
   mkdir build
   cd build
-  cmake .. -DUSE_CUDA=ON -DUSE_CUDNN=ON -DUSE_TENSORRT=/path/to/TensorRT/ 
+  cmake .. -DUSE_CUDA=ON -DUSE_CUDNN=ON -DUSE_TENSORRT=/path/to/TensorRT/
   make -j4
   cd ../python
   python3 setup.py install --user
@@ -151,7 +151,7 @@ Similarly, to enable support for OpenCL devices, run CMake with ``-DUSE_OPENCL=O
   cd neo-ai-dlr
   mkdir build
   cd build
-  cmake .. -DUSE_OPENCL=ON 
+  cmake .. -DUSE_OPENCL=ON
   make -j4
   cd ../python
   python3 setup.py install --user
@@ -213,7 +213,7 @@ Once the compilation is completed, install the Python package by running ``setup
 Building for Android on ARM
 ---------------------------
 
-Android build requires `Android NDK <https://developer.android.com/ndk/downloads/>`_. We utilize the android.toolchain.cmake file in NDK package to configure the crosscompiler 
+Android build requires `Android NDK <https://developer.android.com/ndk/downloads/>`_. We utilize the android.toolchain.cmake file in NDK package to configure the crosscompiler
 
 Also required is `NDK standlone toolchain <https://developer.android.com/ndk/guides/standalone_toolchain>`_. Follow the instructions to generate necessary build-essential tools.
 
@@ -234,12 +234,12 @@ Once done with above steps, invoke cmake with following commands to build Androi
 
 ``ANDROID_PLATFORM`` should correspond to ``minSdkVersion`` of your project. If ``ANDROID_PLATFORM`` is not set it will default to ``android-21``.
 
-For arm64 targets, add 
+For arm64 targets, add
 
 .. code-block:: bash
 
-  -DANDROID_ABI=arm64-v8a 
-  
+  -DANDROID_ABI=arm64-v8a
+
 to cmake flags.
 
 Building for Android Archive (AAR) file
@@ -260,6 +260,25 @@ Install `Android Studio <https://developer.android.com/studio>`_.
   # dlr-release.aar file will be under dlr/build/outputs/aar/ folder
   ls -lah dlr/build/outputs/aar/dlr-release.aar
 
+Building DLR with TensorFlow C library
+--------------------------------------
+
+We can use DLR to run Tensorflow 2.x saved models (including `TensorRT converted models <https://blog.tensorflow.org/2021/01/leveraging-tensorflow-tensorrt-integration.html>`_) via TensorFlow C library.
+
+TensorFlow C library can be downloaded from `tensorflow.org <https://www.tensorflow.org/install/lang_c>`_ or built `from source <https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/lib_package/README.md>`_.
+
+To build DLR with TensorFlow C library:
+
+.. code-block:: bash
+
+  # Build DLR with libtensorflow
+  cmake .. -DWITH_TENSORFLOW2_LIB=<path to libtensorflow folder>
+  make -j8
+
+  # Test DLR with libtensorflow
+  export LD_LIBRARY_PATH=<path to libtensorflow folder>/lib
+  ./dlr_tensorflow2_test
+
 
 Building DLR with Hexagon support
 ---------------------------------
diff --git a/include/dlr.h b/include/dlr.h
index 1bd48e4d6..4137d6ad0 100644
--- a/include/dlr.h
+++ b/include/dlr.h
@@ -50,7 +50,15 @@ typedef void* (*DLRMemalignFunctionPtr)(size_t, size_t);
 
 #ifndef DLR_MODEL_ELEM
 #define DLR_MODEL_ELEM
-enum DLRModelElemType { HEXAGON_LIB, NEO_METADATA, TVM_GRAPH, TVM_LIB, TVM_PARAMS, RELAY_EXEC };
+enum DLRModelElemType {
+  HEXAGON_LIB,
+  NEO_METADATA,
+  TVM_GRAPH,
+  TVM_LIB,
+  TVM_PARAMS,
+  RELAY_EXEC,
+  TF2_SAVED_MODEL
+};
 typedef struct ModelElem {
   const enum DLRModelElemType type;
   const char* path;
@@ -84,6 +92,36 @@ DLR_DLL
 int CreateDLRModelFromModelElem(DLRModelHandle* handle, const DLRModelElem* model_elems,
                                 size_t model_elems_size, int dev_type, int dev_id);
 
+#ifdef DLR_TENSORFLOW2
+
+/*!
+ \brief Tensorflow GPU Options structure for Tensorflow Config structure.
+ */
+typedef struct DLR_TF2GPUOptions {
+  int allow_growth;
+  double per_process_gpu_memory_fraction;
+} DLR_TF2GPUOptions;
+
+/*!
+ \brief Tensorflow Config structure for CreateDLRModelFromTensorflow2.
+ */
+typedef struct DLR_TF2Config {
+  int intra_op_parallelism_threads;
+  int inter_op_parallelism_threads;
+  DLR_TF2GPUOptions gpu_options;
+} DLR_TF2Config;
+
+/*!
+ \brief Creates a DLR model from Tensorflow2.x saved model
+ \param handle The pointer to save the model handle.
+ \param full path to the saved model directory
+ \return 0 for success, -1 for error. Call DLRGetLastError() to get the error message.
+ */
+DLR_DLL
+int CreateDLRModelFromTensorflow2(DLRModelHandle* handle, const char* model_path,
+                                  const DLR_TF2Config tf2_config);
+#endif  // DLR_TENSORFLOW2
+
 #ifdef DLR_HEXAGON
 /*!
  \brief Creates a DLR model from Hexagon
diff --git a/include/dlr_common.h b/include/dlr_common.h
index 1f406f24f..dadf2270d 100644
--- a/include/dlr_common.h
+++ b/include/dlr_common.h
@@ -7,6 +7,7 @@
 #include <sys/types.h>
 
 #include <nlohmann/json.hpp>
+#include <set>
 #include <string>
 #include <vector>
 
@@ -41,7 +42,15 @@
 
 #ifndef DLR_MODEL_ELEM
 #define DLR_MODEL_ELEM
-enum DLRModelElemType { HEXAGON_LIB, NEO_METADATA, TVM_GRAPH, TVM_LIB, TVM_PARAMS, RELAY_EXEC };
+enum DLRModelElemType {
+  HEXAGON_LIB,
+  NEO_METADATA,
+  TVM_GRAPH,
+  TVM_LIB,
+  TVM_PARAMS,
+  RELAY_EXEC,
+  TF2_SAVED_MODEL
+};
 typedef struct ModelElem {
   const DLRModelElemType type;
   const char* path;
@@ -95,8 +104,8 @@ inline bool EndsWith(const std::string& mainStr, const std::string& toMatch) {
     return false;
 }
 
-enum class DLRBackend { kTVM, kTREELITE, kHEXAGON, kRELAYVM, kPIPELINE, kUNKNOWN };
-extern const char* kBackendToStr[6];
+enum class DLRBackend { kTVM, kTREELITE, kHEXAGON, kRELAYVM, kPIPELINE, kUNKNOWN, kTENSORFLOW2 };
+extern const char* kBackendToStr[7];
 
 /*! \brief Get the backend based on the contents of the model folder.
  */
diff --git a/include/dlr_tensorflow2/dlr_tensorflow2.h b/include/dlr_tensorflow2/dlr_tensorflow2.h
new file mode 100644
index 000000000..c3210d53a
--- /dev/null
+++ b/include/dlr_tensorflow2/dlr_tensorflow2.h
@@ -0,0 +1,68 @@
+#ifndef DLR_TENSORFLOW2_H_
+#define DLR_TENSORFLOW2_H_
+
+#include "dlr.h"
+#include "dlr_common.h"
+#include "tensorflow/c/c_api.h"
+
+namespace dlr {
+
+/*! \brief convert DLR_TF2Config to protobuf vector of bytes
+ */
+void PrepareTF2ConfigProto(const DLR_TF2Config& tf2_config, std::vector<std::uint8_t>& config);
+
+/*! \brief class Tensorflow2Model
+ */
+class Tensorflow2Model : public DLRModel {
+ private:
+  TF_Status* status_;
+  TF_Graph* graph_;
+  TF_Session* sess_;
+  std::vector<std::vector<int64_t>> graph_input_shapes_;  // might have -1 dimensions
+  std::vector<std::string> output_names_;
+  std::set<std::string> ignored_names_;
+  std::vector<std::string> output_types_;
+  std::vector<TF_Output> inputs_;
+  std::vector<TF_Output> outputs_;
+  std::vector<TF_Tensor*> input_tensors_;
+  std::vector<TF_Tensor*> output_tensors_;
+  TF_Output ParseTensorName(const std::string& t_name);
+  void DetectInputs();
+  void DetectOutputs();
+  void DetectInputShapes();
+  void PrepInputs();
+  void PrepOutputs();
+  int GetInputId(const char* name);
+  TF_Tensor* AllocateInputTensor(int index, const int64_t* dims, const int n_dim);
+
+ public:
+  /*! \brief Load model files from given folder path.
+   */
+  explicit Tensorflow2Model(const std::string& model_path, const DLDevice& dev,
+                            const DLR_TF2Config& tf2_config);
+  ~Tensorflow2Model();
+
+  virtual const char* GetInputName(int index) const override;
+  virtual const char* GetInputType(int index) const override;
+  virtual const int GetInputDim(int index) const override;
+  virtual const int64_t GetInputSize(int index) const override;
+  virtual const std::vector<int64_t>& GetInputShape(int index) const override;
+  virtual const char* GetWeightName(int index) const override;
+  virtual std::vector<std::string> GetWeightNames() const override;
+  virtual void GetInput(const char* name, void* input) override;
+  virtual void SetInput(const char* name, const int64_t* shape, const void* input,
+                        int dim) override;
+  virtual void Run() override;
+  virtual void GetOutput(int index, void* out) override;
+  virtual const void* GetOutputPtr(int index) const override;
+  virtual const char* GetOutputName(const int index) const override;
+  virtual void GetOutputShape(int index, int64_t* shape) const override;
+  virtual void GetOutputSizeDim(int index, int64_t* size, int* dim) override;
+  virtual const char* GetOutputType(int index) const override;
+  virtual void SetNumThreads(int threads) override;
+  virtual void UseCPUAffinity(bool use) override;
+};
+
+}  // namespace dlr
+
+#endif  // DLR_TENSORFLOW2_H_
diff --git a/src/dlr.cc b/src/dlr.cc
index b70bda28e..ebf5b7ddf 100644
--- a/src/dlr.cc
+++ b/src/dlr.cc
@@ -6,6 +6,10 @@
 #include "dlr_treelite.h"
 #include "dlr_tvm.h"
 
+#ifdef DLR_TENSORFLOW2
+#include "dlr_tensorflow2/dlr_tensorflow2.h"
+#endif  // DLR_TENSORFLOW2
+
 #ifdef DLR_HEXAGON
 #include "dlr_hexagon/dlr_hexagon.h"
 #endif  // DLR_HEXAGON
@@ -260,28 +264,40 @@ extern "C" int GetDLROutputByName(DLRModelHandle* handle, const char* name, void
   API_END();
 }
 
-DLRModelPtr CreateDLRModelPtr(const char* model_path, DLDevice& dev) {
-  std::vector<std::string> path_vec = dlr::MakePathVec(model_path);
-  std::vector<std::string> files = FindFiles(path_vec);
-  DLRBackend backend = dlr::GetBackend(files);
-  if (backend == DLRBackend::kTVM) {
-    return std::make_shared<TVMModel>(files, dev);
-  } else if (backend == DLRBackend::kRELAYVM) {
-    return std::make_shared<RelayVMModel>(files, dev);
-  } else if (backend == DLRBackend::kTREELITE) {
-    return std::make_shared<TreeliteModel>(files, dev);
-#ifdef DLR_HEXAGON
-  } else if (backend == DLRBackend::kHEXAGON) {
-    return std::make_shared<HexagonModel>(files, dev, 1 /*debug_level*/);
-#endif  // DLR_HEXAGON
-  } else {
-    std::string err = "Unable to determine backend from path: '";
-    err = err + model_path + "'.";
-    throw dmlc::Error(err);
-    return nullptr;  // unreachable
-  }
+#ifdef DLR_TENSORFLOW2
+DLR_TF2Config DefaultTFConfig() {
+  // GPUOptions.allow_growth is True
+  // GPUOptions.per_process_gpu_memory_fraction=10%. It allows effectively
+  // share GPU memory. No Performance degradation was detected.
+  DLR_TF2Config tf2_config = {};
+  tf2_config.inter_op_parallelism_threads = 0;
+  tf2_config.intra_op_parallelism_threads = 0;
+  tf2_config.gpu_options.allow_growth = 1;
+  tf2_config.gpu_options.per_process_gpu_memory_fraction = 0.1;
+  return tf2_config;
 }
 
+Tensorflow2Model* CreateTensorflow2Model(const char* model_path, const DLR_TF2Config& tf2_config) {
+  const std::string model_path_string(model_path);
+  // Tensorflow2Model class does not use DLDevice internally
+  DLDevice dev;
+  dev.device_type = static_cast<DLDeviceType>(1);  // 1 - kDLCPU
+  dev.device_id = 0;
+  return new Tensorflow2Model(model_path_string, dev, tf2_config);
+}
+
+/*! \brief Translate c args from ctypes to std types for DLRModelFromTensorflow2
+ * ctor.
+ */
+int CreateDLRModelFromTensorflow2(DLRModelHandle* handle, const char* model_path,
+                                  const DLR_TF2Config tf2_config) {
+  API_BEGIN();
+  DLRModel* model = CreateTensorflow2Model(model_path, tf2_config);
+  *handle = model;
+  API_END();
+}
+#endif  // DLR_TENSORFLOW2
+
 #ifdef DLR_HEXAGON
 /*! \brief Translate c args from ctypes to std types for DLRModelFromHexagon
  * ctor.
@@ -321,6 +337,14 @@ extern "C" int CreateDLRModel(DLRModelHandle* handle, const char* model_path, in
       model = new RelayVMModel(files, dev);
     } else if (backend == DLRBackend::kTREELITE) {
       model = new TreeliteModel(files, dev);
+#ifdef DLR_TENSORFLOW2
+    } else if (backend == DLRBackend::kTENSORFLOW2) {
+      const std::string model_path_string(model_path);
+      // input and output tensor names will be detected automatically.
+      // use undefined number of threads - threads=0
+      const DLR_TF2Config tf2_config = DefaultTFConfig();
+      model = new Tensorflow2Model(model_path_string, dev, tf2_config);
+#endif  // DLR_TENSORFLOW2
 #ifdef DLR_HEXAGON
     } else if (backend == DLRBackend::kHEXAGON) {
       model = new HexagonModel(files, dev, 1 /*debug_level*/);
@@ -340,6 +364,36 @@ extern "C" int CreateDLRModel(DLRModelHandle* handle, const char* model_path, in
   API_END();
 }
 
+DLRModelPtr CreateDLRModelPtr(const char* model_path, DLDevice& dev) {
+  std::vector<std::string> path_vec = dlr::MakePathVec(model_path);
+  std::vector<std::string> files = FindFiles(path_vec);
+  DLRBackend backend = dlr::GetBackend(files);
+  if (backend == DLRBackend::kTVM) {
+    return std::make_shared<TVMModel>(files, dev);
+  } else if (backend == DLRBackend::kRELAYVM) {
+    return std::make_shared<RelayVMModel>(files, dev);
+  } else if (backend == DLRBackend::kTREELITE) {
+    return std::make_shared<TreeliteModel>(files, dev);
+#ifdef DLR_TENSORFLOW2
+  } else if (backend == DLRBackend::kTENSORFLOW2) {
+    const std::string model_path_string(model_path);
+    // input and output tensor names will be detected automatically.
+    // use undefined number of threads - threads=0
+    const DLR_TF2Config tf2_config = DefaultTFConfig();
+    return std::make_shared<Tensorflow2Model>(model_path_string, dev, tf2_config);
+#endif  // DLR_TENSORFLOW2
+#ifdef DLR_HEXAGON
+  } else if (backend == DLRBackend::kHEXAGON) {
+    return std::make_shared<HexagonModel>(files, dev, 1 /*debug_level*/);
+#endif  // DLR_HEXAGON
+  } else {
+    std::string err = "Unable to determine backend from path: '";
+    err = err + model_path + "'.";
+    throw dmlc::Error(err);
+    return nullptr;  // unreachable
+  }
+}
+
 extern "C" int CreateDLRModelFromModelElem(DLRModelHandle* handle, const DLRModelElem* model_elems,
                                            size_t model_elems_size, int dev_type, int dev_id) {
   API_BEGIN();
diff --git a/src/dlr_common.cc b/src/dlr_common.cc
index 1cca9290a..a5e7f4f06 100644
--- a/src/dlr_common.cc
+++ b/src/dlr_common.cc
@@ -7,7 +7,8 @@
 
 using namespace dlr;
 
-const char* dlr::kBackendToStr[] = {"tvm", "treelite", "hexagon", "relayvm", "pipeline", "unknown"};
+const char* dlr::kBackendToStr[] = {"tvm",      "treelite", "hexagon",   "relayvm",
+                                    "pipeline", "unknown",  "tensorflow"};
 
 bool dlr::IsFileEmpty(const std::string& filePath) {
   std::ifstream pFile(filePath);
@@ -113,10 +114,13 @@ DLRBackend dlr::GetBackend(const std::vector<std::string>& files) {
   // Scan files to guess the backend.
   bool has_tvm_lib = false;
   for (auto filename : files) {
+    std::cout << "filename: " << filename << std::endl;
     if (EndsWith(filename, ".params")) {
       return DLRBackend::kTVM;
     } else if (EndsWith(filename, ".ro")) {
       return DLRBackend::kRELAYVM;
+    } else if (EndsWith(filename, "saved_model.pb") || EndsWith(filename, "saved_model.pbtxt")) {
+      return DLRBackend::kTENSORFLOW2;
     } else if (EndsWith(filename, "_hexagon_model.so")) {
       return DLRBackend::kHEXAGON;
     } else if (!EndsWith(filename, LIBDLR) && EndsWith(filename, LIBEXT)) {
@@ -135,6 +139,8 @@ DLRBackend dlr::GetBackend(const std::vector<DLRModelElem>& model_elems) {
       return DLRBackend::kTVM;
     } else if (el.type == DLRModelElemType::RELAY_EXEC) {
       return DLRBackend::kRELAYVM;
+    } else if (el.type == DLRModelElemType::TF2_SAVED_MODEL) {
+      return DLRBackend::kTENSORFLOW2;
     } else if (el.type == DLRModelElemType::HEXAGON_LIB) {
       return DLRBackend::kHEXAGON;
     } else if (el.type == DLRModelElemType::TVM_LIB) {
diff --git a/src/dlr_tensorflow2/dlr_tensorflow2.cc b/src/dlr_tensorflow2/dlr_tensorflow2.cc
new file mode 100644
index 000000000..a839596d6
--- /dev/null
+++ b/src/dlr_tensorflow2/dlr_tensorflow2.cc
@@ -0,0 +1,407 @@
+#include "dlr_tensorflow2/dlr_tensorflow2.h"
+
+#include <cstring>
+#include <fstream>
+#include <numeric>
+#include <regex>
+
+using namespace dlr;
+
+void dlr::PrepareTF2ConfigProto(const DLR_TF2Config& tf2_config,
+                                std::vector<std::uint8_t>& config) {
+  if (tf2_config.intra_op_parallelism_threads > 0 &&
+      tf2_config.intra_op_parallelism_threads < 256) {
+    // More info https://github.com/tensorflow/tensorflow/issues/13853
+    config.insert(config.end(), {0x10, (std::uint8_t)tf2_config.intra_op_parallelism_threads});
+    LOG(INFO) << "Set intra_op_parallelism_threads to " << tf2_config.intra_op_parallelism_threads;
+  }
+  if (tf2_config.inter_op_parallelism_threads > 0 &&
+      tf2_config.inter_op_parallelism_threads < 256) {
+    // More info https://github.com/tensorflow/tensorflow/issues/13853
+    config.insert(config.end(), {0x28, (std::uint8_t)tf2_config.inter_op_parallelism_threads});
+    LOG(INFO) << "Set inter_op_parallelism_threads to " << tf2_config.inter_op_parallelism_threads;
+  }
+
+  // Tensorflow GPUOptions
+  std::vector<std::uint8_t> gpu_options = {0x32, 0x0};
+
+  double gpu_memory_fraction = tf2_config.gpu_options.per_process_gpu_memory_fraction;
+  if (gpu_memory_fraction > 0.001 && gpu_memory_fraction < 1.0) {
+    std::uint8_t proto[9] = {0x9, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};
+    auto bytes = reinterpret_cast<std::uint8_t*>(&gpu_memory_fraction);
+    // Put it to the config byte-array, from 1 to 9:
+    for (std::size_t i = 0; i < sizeof(gpu_memory_fraction); i++) {
+      proto[i + 1] = bytes[i];
+    }
+    gpu_options.insert(gpu_options.end(), proto, proto + 9);
+    LOG(INFO) << "Set gpu_options.per_process_gpu_memory_fraction to " << gpu_memory_fraction;
+  }
+
+  if (tf2_config.gpu_options.allow_growth != 0) {
+    gpu_options.insert(gpu_options.end(), {0x20, 0x1});
+    LOG(INFO) << "Set gpu_options.allow_growth to True";
+  }
+  // update gpu_options data length (stored in byte #2)
+  gpu_options[1] = gpu_options.size() - 2;
+
+  if (gpu_options[1] > 0) {
+    config.insert(config.end(), gpu_options.begin(), gpu_options.end());
+  }
+}
+
+TF_Output Tensorflow2Model::ParseTensorName(const std::string& t_name) {
+  std::regex r("^(.+):(\\d+)$");
+  std::smatch match;
+  if (!std::regex_search(t_name, match, r)) {
+    LOG(FATAL) << "ERROR: failed to parse tensor name " << t_name;
+  }
+  std::string op_name = match.str(1);
+  int op_out_id = std::stoi(match.str(2));
+  TF_Operation* op = TF_GraphOperationByName(graph_, op_name.c_str());
+  if (op == nullptr) {
+    LOG(FATAL) << "ERROR: TF_GraphOperationByName failed for operation " << op_name;
+  }
+  const TF_Output oper_out = {op, op_out_id};
+  return oper_out;
+}
+
+void Tensorflow2Model::DetectInputs() {
+  size_t pos = 0;
+  TF_Operation* op;
+  while ((op = TF_GraphNextOperation(graph_, &pos)) != nullptr) {
+    const std::string op_type = TF_OperationOpType(op);
+    const int n_in = TF_OperationNumInputs(op);
+    const int n_out = TF_OperationNumOutputs(op);
+    const std::string op_name = TF_OperationName(op);
+    if (op_type == "Placeholder" && n_in == 0 && n_out == 1 && ignored_names_.count(op_name) == 0) {
+      input_names_.push_back(op_name + ":0");
+    }
+  }
+  num_inputs_ = input_names_.size();
+  std::string msg = "Found " + std::to_string(num_inputs_) + " possible inputs: ";
+  for (int i = 0; i < num_inputs_; i++) {
+    if (i > 0) {
+      msg += ", ";
+    }
+    msg += input_names_[i];
+  }
+  LOG(INFO) << msg;
+}
+
+void Tensorflow2Model::DetectOutputs() {
+  size_t pos = 0;
+  TF_Operation* op;
+  // while loop
+  while ((op = TF_GraphNextOperation(graph_, &pos)) != nullptr) {
+    const std::string op_type = TF_OperationOpType(op);
+    const int n_out = TF_OperationNumOutputs(op);
+    const int n_cout = TF_OperationNumControlOutputs(op);
+    const std::string op_name = TF_OperationName(op);
+    if (op_type != "Const" && op_type != "Assign" && op_type != "NoOp" &&
+        op_type != "Placeholder" && n_cout == 0 && ignored_names_.count(op_name) == 0) {
+      int n_consumers = 0;
+      for (int i = 0; i < n_out; i++) {
+        const TF_Output tf_out = {op, i};
+        n_consumers += TF_OperationOutputNumConsumers(tf_out);
+        if (n_consumers != 0) {
+          break;
+        }
+      }
+      if (n_consumers != 0) {
+        continue;  // while loop
+      }
+      for (int i = 0; i < n_out; i++) {
+        output_names_.push_back(op_name + ":" + std::to_string(i));
+      }
+    }
+  }
+  num_outputs_ = output_names_.size();
+  std::string msg = "Found " + std::to_string(num_outputs_) + " possible outputs: ";
+  for (int i = 0; i < num_outputs_; i++) {
+    if (i > 0) {
+      msg += ", ";
+    }
+    msg += output_names_[i];
+  }
+  LOG(INFO) << msg;
+}
+
+void Tensorflow2Model::DetectInputShapes() {
+  for (int i = 0; i < num_inputs_; i++) {
+    const std::string& t_name = input_names_[i];
+    const TF_Output oper_out = ParseTensorName(t_name);
+
+    int n_dim = TF_GraphGetTensorNumDims(graph_, oper_out, status_);
+    if (TF_GetCode(status_) != TF_OK) {
+      LOG(FATAL) << "ERROR: TF_GraphGetTensorNumDims failed " << TF_Message(status_);
+      return;  // unreachable
+    }
+    int64_t dims[n_dim];
+    TF_GraphGetTensorShape(graph_, oper_out, dims, n_dim, status_);
+    if (TF_GetCode(status_) != TF_OK) {
+      LOG(FATAL) << "ERROR: TF_GraphGetTensorShape failed " << TF_Message(status_);
+      return;  // unreachable
+    }
+    graph_input_shapes_.push_back(std::vector<int64_t>(dims, dims + n_dim));
+  }
+  input_shapes_.resize(num_inputs_);
+}
+
+TF_Tensor* Tensorflow2Model::AllocateInputTensor(int index, const int64_t* dims, const int n_dim) {
+  const TF_Output oper_out = inputs_[index];
+  size_t num_elements = 1;
+  for (int z = 0; z < n_dim; z++) {
+    if (dims[z] < 1) {
+      LOG(FATAL) << "ERROR: non-positive dimensions are not supported. "
+                 << "input id: " << index << ", dims[" << z << "]=" << dims[z];
+    }
+    num_elements *= dims[z];
+  }
+  const TF_DataType t_type = TF_OperationOutputType(oper_out);
+  const size_t num_bytes = TF_DataTypeSize(t_type) * num_elements;
+  TF_Tensor* tensor = TF_AllocateTensor(t_type, dims, n_dim, num_bytes);
+  LOG(INFO) << "Input Tensor " << index << " was allocated";
+  return tensor;
+}
+
+void Tensorflow2Model::PrepInputs() {
+  for (std::string& t_name : input_names_) {
+    TF_Output oper_out = ParseTensorName(t_name);
+    const TF_DataType t_type = TF_OperationOutputType(oper_out);
+    input_types_.push_back(std::to_string((int)t_type));
+    inputs_.push_back(oper_out);
+    // fill output_tensors_ vector with nulls
+    input_tensors_.push_back(nullptr);
+  }
+}
+
+void Tensorflow2Model::PrepOutputs() {
+  for (std::string& t_name : output_names_) {
+    TF_Output oper_out = ParseTensorName(t_name);
+    const TF_DataType t_type = TF_OperationOutputType(oper_out);
+    output_types_.push_back(std::to_string((int)t_type));
+    outputs_.push_back(oper_out);
+    // fill output_tensors_ vector with nulls
+    output_tensors_.push_back(nullptr);
+  }
+}
+
+int Tensorflow2Model::GetInputId(const char* name) {
+  // In most of the cases it will be just 1 element in the vector.
+  // Scan vector to find tensor by name.
+  for (int i = 0; i < num_inputs_; i++) {
+    if (input_names_[i].compare(name) == 0) {
+      return i;
+    }
+  }
+  LOG(FATAL) << "Input Tensor not found, name: " << name;
+  return -1;  // unreachable
+}
+
+// Constructor
+Tensorflow2Model::Tensorflow2Model(const std::string& model_path, const DLDevice& dev,
+                                   const DLR_TF2Config& tf2_config)
+    : DLRModel(dev, DLRBackend::kTENSORFLOW2) {
+  status_ = TF_NewStatus();
+  graph_ = TF_NewGraph();
+  TF_SessionOptions* sess_opts = TF_NewSessionOptions();
+  std::vector<std::uint8_t> config;
+  PrepareTF2ConfigProto(tf2_config, config);
+  if (!config.empty()) {
+    TF_SetConfig(sess_opts, config.data(), config.size(), status_);
+    if (TF_GetCode(status_) != TF_OK) {
+      TF_DeleteSessionOptions(sess_opts);
+      LOG(FATAL) << "ERROR: TF_SetConfig failed " << TF_Message(status_);
+      return;  // unreachable
+    }
+  }
+  TF_Buffer* run_opts = nullptr;
+  TF_Buffer* meta_graph_def = nullptr;
+  const char* tags = "serve";
+  int ntags = 1;
+  sess_ = TF_LoadSessionFromSavedModel(sess_opts, run_opts, model_path.c_str(), &tags, ntags,
+                                       graph_, meta_graph_def, status_);
+  if (TF_GetCode(status_) != TF_OK) {
+    LOG(FATAL) << "ERROR: Unable to create Session " << TF_Message(status_);
+    return;  // unreachable
+  }
+  ignored_names_ = {
+      "saver_filename",             // name of the checkpoint
+      "StatefulPartitionedCall_1",  // the loss
+      "StatefulPartitionedCall_2"   // save operation
+  };
+  DetectInputs();
+  DetectInputShapes();
+  DetectOutputs();
+  PrepInputs();
+  PrepOutputs();
+
+  TF_DeleteSessionOptions(sess_opts);
+  LOG(INFO) << "Tensorflow Session was created";
+}
+
+// Destructor
+Tensorflow2Model::~Tensorflow2Model() {
+  for (TF_Tensor* tensor : input_tensors_) {
+    TF_DeleteTensor(tensor);
+  }
+  for (TF_Tensor* tensor : output_tensors_) {
+    TF_DeleteTensor(tensor);
+  }
+  TF_CloseSession(sess_, status_);
+  // Result of close is ignored, delete anyway.
+  TF_DeleteSession(sess_, status_);
+  TF_DeleteGraph(graph_);
+  TF_DeleteStatus(status_);
+
+  LOG(INFO) << "Tensorflow2Model was deleted";
+}
+
+std::vector<std::string> Tensorflow2Model::GetWeightNames() const {
+  LOG(FATAL) << "GetWeightNames is not supported by Tensorflow backend";
+  return std::vector<std::string>();  // unreachable
+}
+
+const char* Tensorflow2Model::GetInputName(int index) const {
+  CHECK_LT(index, num_inputs_) << "Input index is out of range.";
+  return input_names_[index].c_str();
+}
+
+const char* Tensorflow2Model::GetInputType(int index) const {
+  CHECK_LT(index, num_inputs_) << "Input index is out of range.";
+  return input_types_[index].c_str();
+}
+
+const std::vector<int64_t>& Tensorflow2Model::GetInputShape(int index) const {
+  CHECK_LT(index, num_inputs_) << "Input index is out of range.";
+  return input_shapes_[index].empty() ? graph_input_shapes_[index] : input_shapes_[index];
+}
+
+const int Tensorflow2Model::GetInputDim(int index) const {
+  CHECK_LT(index, num_inputs_) << "Input index is out of range.";
+  return graph_input_shapes_[index].size();
+}
+
+const int64_t Tensorflow2Model::GetInputSize(int index) const {
+  CHECK_LT(index, num_inputs_) << "Input index is out of range.";
+  const std::vector<int64_t>& shape =
+      input_shapes_[index].empty() ? graph_input_shapes_[index] : input_shapes_[index];
+  if (dlr::HasNegative(shape.data(), shape.size())) return -1;
+  return abs(std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<int64_t>()));
+}
+
+const char* Tensorflow2Model::GetWeightName(int index) const {
+  LOG(FATAL) << "GetWeightName is not supported by Tensorflow backend";
+  return "";  // unreachable
+}
+
+void Tensorflow2Model::SetInput(const char* name, const int64_t* shape, const void* input,
+                                int dim) {
+  int index = GetInputId(name);
+  TF_Tensor* tensor = input_tensors_[index];
+
+  bool keep = tensor != nullptr && TF_NumDims(tensor) == dim;
+  if (keep) {
+    for (int i = 0; i < dim; i++) {
+      if (shape[i] != TF_Dim(tensor, i)) {
+        keep = false;
+        break;
+      }
+    }
+  }
+  if (!keep) {
+    if (tensor != nullptr) {
+      TF_DeleteTensor(tensor);
+      input_tensors_[index] = nullptr;
+    }
+    tensor = AllocateInputTensor(index, shape, dim);
+    input_tensors_[index] = tensor;
+    input_shapes_[index] = std::vector<int64_t>(shape, shape + dim);
+  }
+  size_t num_bytes = TF_TensorByteSize(tensor);
+  void* in_t_data = TF_TensorData(tensor);
+  std::memcpy(in_t_data, input, num_bytes);
+}
+
+void Tensorflow2Model::GetInput(const char* name, void* input) {
+  int index = GetInputId(name);
+  TF_Tensor* tensor = input_tensors_[index];
+  CHECK_NOTNULL(tensor);
+  size_t num_bytes = TF_TensorByteSize(tensor);
+  void* in_t_data = TF_TensorData(tensor);
+  std::memcpy(input, in_t_data, num_bytes);
+}
+
+const char* Tensorflow2Model::GetOutputName(const int index) const {
+  CHECK_LT(index, num_outputs_) << "Output index is out of range.";
+  return output_names_[index].c_str();
+}
+
+void Tensorflow2Model::GetOutputShape(int index, int64_t* shape) const {
+  CHECK_LT(index, num_outputs_) << "Output index is out of range.";
+  TF_Tensor* tensor = output_tensors_[index];
+  CHECK_NOTNULL(tensor);
+  int n_dim = TF_NumDims(tensor);
+  for (int i = 0; i < n_dim; i++) {
+    shape[i] = TF_Dim(tensor, i);
+  }
+}
+
+void Tensorflow2Model::GetOutput(int index, void* output) {
+  CHECK_LT(index, num_outputs_) << "Output index is out of range.";
+  const TF_Tensor* tensor = output_tensors_[index];
+  CHECK_NOTNULL(tensor);
+  size_t num_bytes = TF_TensorByteSize(tensor);
+  const void* out_t_data = TF_TensorData(tensor);
+  std::memcpy(output, out_t_data, num_bytes);
+}
+
+const void* Tensorflow2Model::GetOutputPtr(int index) const {
+  CHECK_LT(index, num_outputs_) << "Output index is out of range.";
+  const TF_Tensor* tensor = output_tensors_[index];
+  CHECK_NOTNULL(tensor);
+  const void* out_t_data = TF_TensorData(tensor);
+  return out_t_data;
+}
+
+void Tensorflow2Model::GetOutputSizeDim(int index, int64_t* size, int* dim) {
+  CHECK_LT(index, num_outputs_) << "Output index is out of range.";
+  TF_Tensor* tensor = output_tensors_[index];
+  CHECK_NOTNULL(tensor);
+  *size = TF_TensorElementCount(tensor);
+  *dim = TF_NumDims(tensor);
+}
+
+const char* Tensorflow2Model::GetOutputType(int index) const {
+  CHECK_LT(index, num_outputs_) << "Output index is out of range.";
+  return output_types_[index].c_str();
+}
+
+void Tensorflow2Model::Run() {
+  // Delete previous output Tensors to prevent GPU memory leak
+  for (TF_Tensor* tensor : output_tensors_) {
+    if (tensor != nullptr) {
+      TF_DeleteTensor(tensor);
+    }
+  }
+  TF_SessionRun(sess_,
+                nullptr,  // Run options.
+                inputs_.data(), input_tensors_.data(), num_inputs_, outputs_.data(),
+                output_tensors_.data(), num_outputs_, nullptr,
+                0,        // Target operations, number of targets.
+                nullptr,  // Run metadata.
+                status_   // Output status.
+  );
+  if (TF_GetCode(status_) != TF_OK) {
+    LOG(FATAL) << "ERROR: Session run error: " << TF_Message(status_);
+    return;  // unreachable
+  }
+}
+
+void Tensorflow2Model::SetNumThreads(int threads) {
+  LOG(FATAL) << "SetNumThreads is not supported by Tensorflow backend";
+}
+
+void Tensorflow2Model::UseCPUAffinity(bool use) {
+  LOG(FATAL) << "UseCPUAffinity is not supported by Tensorflow backend";
+}


Test patch:
diff --git a/tests/cpp/dlr_tensorflow2/dlr_tensorflow2_internal_test.cc b/tests/cpp/dlr_tensorflow2/dlr_tensorflow2_internal_test.cc
new file mode 100644
index 000000000..928fa1247
--- /dev/null
+++ b/tests/cpp/dlr_tensorflow2/dlr_tensorflow2_internal_test.cc
@@ -0,0 +1,55 @@
+#include <dmlc/logging.h>
+#include <gtest/gtest.h>
+
+#include "dlr_tensorflow2/dlr_tensorflow2.h"
+
+using namespace dlr;
+
+TEST(PrepareTFConfigProto1, test1) {
+  DLR_TF2Config tf2_config = {};
+  tf2_config.intra_op_parallelism_threads = 2;
+  tf2_config.inter_op_parallelism_threads = 3;
+  tf2_config.gpu_options.per_process_gpu_memory_fraction = 0.33;
+  tf2_config.gpu_options.allow_growth = 1;
+  std::vector<std::uint8_t> config;
+  PrepareTF2ConfigProto(tf2_config, config);
+  std::uint8_t exp[] = {0x10, 0x2,  0x28, 0x3,  0x32, 0xb,  0x9,  0x1f, 0x85,
+                        0xeb, 0x51, 0xb8, 0x1e, 0xd5, 0x3f, 0x20, 0x1};
+  EXPECT_TRUE(std::equal(std::begin(exp), std::end(exp), config.begin()));
+}
+
+TEST(PrepareTFConfigProto1, test2) {
+  DLR_TF2Config tf2_config = {};
+  tf2_config.intra_op_parallelism_threads = 2;
+  tf2_config.inter_op_parallelism_threads = 3;
+  tf2_config.gpu_options.allow_growth = 1;
+  std::vector<std::uint8_t> config;
+  PrepareTF2ConfigProto(tf2_config, config);
+  std::uint8_t exp[] = {0x10, 0x2, 0x28, 0x3, 0x32, 0x2, 0x20, 0x1};
+  EXPECT_TRUE(std::equal(std::begin(exp), std::end(exp), config.begin()));
+}
+
+TEST(PrepareTFConfigProto1, test3) {
+  DLR_TF2Config tf2_config = {};
+  tf2_config.intra_op_parallelism_threads = 2;
+  tf2_config.inter_op_parallelism_threads = 3;
+  std::vector<std::uint8_t> config;
+  PrepareTF2ConfigProto(tf2_config, config);
+  std::uint8_t exp[] = {0x10, 0x2, 0x28, 0x3};
+  EXPECT_TRUE(std::equal(std::begin(exp), std::end(exp), config.begin()));
+}
+
+TEST(PrepareTFConfigProto1, test4) {
+  DLR_TF2Config tf2_config = {};
+  std::vector<std::uint8_t> config;
+  PrepareTF2ConfigProto(tf2_config, config);
+  EXPECT_TRUE(config.empty());
+}
+
+int main(int argc, char** argv) {
+  testing::InitGoogleTest(&argc, argv);
+#ifndef _WIN32
+  testing::FLAGS_gtest_death_test_style = "threadsafe";
+#endif  // _WIN32
+  return RUN_ALL_TESTS();
+}
diff --git a/tests/cpp/dlr_tensorflow2/dlr_tensorflow2_test.cc b/tests/cpp/dlr_tensorflow2/dlr_tensorflow2_test.cc
new file mode 100644
index 000000000..ef10e0398
--- /dev/null
+++ b/tests/cpp/dlr_tensorflow2/dlr_tensorflow2_test.cc
@@ -0,0 +1,257 @@
+#include <dmlc/logging.h>
+#include <gtest/gtest.h>
+
+#include <cstring>
+#include <fstream>
+#include <iostream>
+#include <sstream>
+
+#include "../test_utils.hpp"
+#include "dlr.h"
+
+int ArgMax(float* data, int start, int size) {
+  int idx = 0;
+  float v = 0.0f;
+  for (int i = start; i < start + size; i++) {
+    float vi = data[i];
+    if (vi > v) {
+      idx = i;
+      v = vi;
+    }
+  }
+  return idx;
+}
+
+void CheckInputShape(DLRModelHandle& handle, const int batch_size) {
+  // GetDLRInputSizeDim
+  int64_t in_size;
+  int in_dim;
+  if (GetDLRInputSizeDim(&handle, 0, &in_size, &in_dim)) {
+    FAIL() << "GetDLRInputSizeDim failed";
+  }
+  LOG(INFO) << "GetDLRInputSizeDim.size: " << in_size;
+  LOG(INFO) << "GetDLRInputSizeDim.dim: " << in_dim;
+  if (batch_size == -1) {
+    EXPECT_EQ(-1, in_size);
+  } else {
+    EXPECT_EQ(batch_size * 224 * 224 * 3, in_size);
+  }
+  EXPECT_EQ(4, in_dim);
+
+  // GetDLRInputShape
+  int64_t shape[in_dim];
+  if (GetDLRInputShape(&handle, 0, shape)) {
+    FAIL() << "GetDLRInputShape failed";
+  }
+  std::stringstream ss;
+  ss << "GetDLRInputShape: (" << shape[0];
+  for (int i = 1; i < in_dim; i++) {
+    ss << "," << shape[i];
+  }
+  ss << ")";
+  LOG(INFO) << ss.str();
+  const int64_t exp_shape[4] = {batch_size, 224, 224, 3};
+  EXPECT_TRUE(std::equal(std::begin(exp_shape), std::end(exp_shape), shape));
+}
+
+void CheckAllDLRMethods(DLRModelHandle& handle, const int batch_size, const int prev_batch_size) {
+  // GetDLRBackend
+  const char* backend_name;
+  if (GetDLRBackend(&handle, &backend_name)) {
+    FAIL() << "GetDLRBackend failed";
+  }
+  LOG(INFO) << "GetDLRBackend: " << backend_name;
+  EXPECT_STREQ("tensorflow", backend_name);
+
+  // GetDLRNumInputs
+  int num_inputs;
+  if (GetDLRNumInputs(&handle, &num_inputs)) {
+    FAIL() << "GetDLRNumInputs failed";
+  }
+  LOG(INFO) << "GetDLRNumInputs: " << num_inputs;
+  EXPECT_EQ(1, num_inputs);
+
+  // GetDLRNumOutputs
+  int num_outputs;
+  if (GetDLRNumOutputs(&handle, &num_outputs)) {
+    FAIL() << "GetDLRNumOutputs failed";
+  }
+  LOG(INFO) << "GetDLRNumOutputs: " << num_outputs;
+  EXPECT_EQ(1, num_outputs);
+
+  // GetDLRNumWeights
+  int num_weights;
+  if (GetDLRNumWeights(&handle, &num_weights)) {
+    FAIL() << "GetDLRNumWeights failed";
+  }
+  LOG(INFO) << "GetDLRNumWeights: " << num_weights;
+  EXPECT_EQ(0, num_weights);
+
+  // GetDLRInputName
+  const char* input_name;
+  if (GetDLRInputName(&handle, 0, &input_name)) {
+    FAIL() << "GetDLRInputName failed";
+  }
+  LOG(INFO) << "DLRInputName: " << input_name;
+  EXPECT_STREQ("serving_default_inputs:0", input_name);
+
+  // GetDLROutputName
+  const char* output_name;
+  if (GetDLROutputName(&handle, 0, &output_name)) {
+    FAIL() << "GetDLROutputName failed";
+  }
+  LOG(INFO) << "DLROutputName: " << output_name;
+  EXPECT_STREQ("StatefulPartitionedCall:0", output_name);
+
+  // GetDLRInputType
+  const char* input_type;
+  if (GetDLRInputType(&handle, 0, &input_type)) {
+    FAIL() << "GetDLRInputType failed";
+  }
+  LOG(INFO) << "DLRInputType: " << input_type;
+  EXPECT_STREQ("1", input_type);
+
+  // GetDLROutputType
+  const char* output_type;
+  if (GetDLROutputType(&handle, 0, &output_type)) {
+    FAIL() << "GetDLROutputType failed";
+  }
+  LOG(INFO) << "DLROutputType: " << output_type;
+  EXPECT_STREQ("1", output_type);
+
+  CheckInputShape(handle, prev_batch_size);
+
+  // Load image
+  size_t img_size = 224 * 224 * 3;
+  std::vector<float> img = LoadImageAndPreprocess("cat224-3.txt", img_size, batch_size);
+
+  LOG(INFO) << "Input sample: [" << img[0] << "," << img[1] << "..." << img[img_size - 1] << "]...["
+            << img[img_size * (batch_size - 1)] << "," << img[img_size * (batch_size - 1) + 1]
+            << "..." << img[img_size * batch_size - 1] << "]";
+  img_size *= batch_size;
+
+  // SetDLRInput
+  const int64_t in_shape[4] = {batch_size, 224, 224, 3};
+  if (SetDLRInput(&handle, input_name, in_shape, img.data(), 4)) {
+    FAIL() << "SetDLRInput failed";
+  }
+  LOG(INFO) << "SetDLRInput - OK";
+
+  CheckInputShape(handle, batch_size);
+
+  // GetDLRInput
+  std::vector<float> input2(img_size);
+  if (GetDLRInput(&handle, input_name, input2.data())) {
+    FAIL() << "GetDLRInput failed";
+  }
+  EXPECT_TRUE(std::equal(img.begin(), img.end(), input2.begin()));
+  LOG(INFO) << "GetDLRInput - OK";
+
+  // RunDLRModel
+  if (RunDLRModel(&handle)) {
+    FAIL() << "RunDLRModel failed";
+  }
+  LOG(INFO) << "RunDLRModel - OK";
+
+  // GetDLROutputSizeDim
+  int64_t out_size;
+  int out_dim;
+  if (GetDLROutputSizeDim(&handle, 0, &out_size, &out_dim)) {
+    FAIL() << "GetDLROutputSizeDim failed";
+  }
+  LOG(INFO) << "GetDLROutputSizeDim.size: " << out_size;
+  LOG(INFO) << "GetDLROutputSizeDim.dim: " << out_dim;
+  EXPECT_EQ(1001 * batch_size, out_size);
+  EXPECT_EQ(2, out_dim);
+
+  // GetDLROutputShape
+  int64_t shape[out_dim];
+  if (GetDLROutputShape(&handle, 0, shape)) {
+    FAIL() << "GetDLROutputShape failed";
+  }
+  std::stringstream ss;
+  ss << "GetDLROutputShape: (" << shape[0];
+  for (int i = 1; i < out_dim; i++) {
+    ss << "," << shape[i];
+  }
+  ss << ")";
+  LOG(INFO) << ss.str();
+  const int64_t exp_shape[2] = {batch_size, 1001};
+  EXPECT_TRUE(std::equal(std::begin(exp_shape), std::end(exp_shape), shape));
+
+  // GetDLROutput (the first and the last item in the batch)
+  std::vector<float> output(out_size);
+  if (GetDLROutput(&handle, 0, output.data())) {
+    FAIL() << "GetDLROutput failed";
+  }
+  LOG(INFO) << "GetDLROutput - OK";
+  size_t out_size0 = out_size / batch_size;
+  for (int i = 0; i < batch_size; i++) {
+    size_t out_offset = i * out_size0;
+    size_t max_id = ArgMax(output.data(), out_offset, out_size0);
+    LOG(INFO) << "ArgMax: " << max_id << ", Prop: " << output[max_id];
+    // Tensorflow class range is 1-1000 (output size 1001)
+    // Imagenet1000 class range is 0-999
+    // Imagenet 282 - tiger cat
+    EXPECT_EQ(out_offset + 282, max_id);
+    EXPECT_GE(output[max_id], 0.5f);
+    // Imagenet 281 - tabby, tabby cat
+    EXPECT_GE(output[out_offset + 281], 0.3f);
+  }
+}
+
+TEST(Tensorflow, CreateDLRModelFromTensorflow2) {
+  // CreateDLRModelFromTensorflow2
+  const char* model_file = "./imagenet_mobilenet_v2";
+  DLR_TF2Config tf2_config = {};
+  tf2_config.inter_op_parallelism_threads = 2;
+  tf2_config.intra_op_parallelism_threads = 2;
+
+  DLRModelHandle handle = nullptr;
+  if (CreateDLRModelFromTensorflow2(&handle, model_file, tf2_config)) {
+    FAIL() << DLRGetLastError() << std::endl;
+  }
+  LOG(INFO) << "CreateDLRModelFromTensorflow2 - OK";
+
+  // batch size 1
+  CheckAllDLRMethods(handle, 1, -1);
+  // Run the same model for input with batch size 8
+  CheckAllDLRMethods(handle, 8, 1);
+
+  // DeleteDLRModel
+  DeleteDLRModel(&handle);
+}
+
+TEST(Tensorflow, CreateDLRModel) {
+  // Use generic CreateDLRModel
+  // input and output tensor names will be detected automatically.
+  const char* model_file = "./imagenet_mobilenet_v2";
+  const int dev_type = 1;  // 1 - kDLCPU
+  const int dev_id = 0;
+
+  DLRModelHandle handle = nullptr;
+  if (CreateDLRModel(&handle, model_file, dev_type, dev_id)) {
+    FAIL() << DLRGetLastError() << std::endl;
+  }
+  LOG(INFO) << "CreateDLRModel - OK";
+
+  // batch size 2
+  CheckAllDLRMethods(handle, 2, -1);
+  // Run the same model for input with batch size 8
+  CheckAllDLRMethods(handle, 8, 2);
+
+  // DeleteDLRModel
+  DeleteDLRModel(&handle);
+  ASSERT_EQ(nullptr, handle);
+  // Test that calling DeleteDLRModel again
+  // does not crash the program (no segmentation fault)
+  DeleteDLRModel(&handle);
+}
+
+int main(int argc, char** argv) {
+  testing::InitGoogleTest(&argc, argv);
+#ifndef _WIN32
+  testing::FLAGS_gtest_death_test_style = "threadsafe";
+#endif  // _WIN32
+  return RUN_ALL_TESTS();
+}