From b4b03b14f71eac53b1abee5e9e9e6194337b1829 Mon Sep 17 00:00:00 2001
From: Roman Kazantsev <roman.kazantsev@intel.com>
Date: Tue, 11 Aug 2020 18:02:05 +0300
Subject: [PATCH] Separate MO configuration for TensorFlow 2 model conversion
 (#1685)

* Separate MO configuration for TensorFlow 2 model conversion

Also, it updates documentation including steps to convert
TF2 model with a custom layer in Keras H5 format into SavedModel

* Do fixes based on the first-round code review
---
 .../prepare_model/Config_Model_Optimizer.md   | 20 +++++++++++++++----
 .../Convert_Model_From_TensorFlow.md          | 10 ++++++++++
 .../installing-openvino-linux-fpga.md         |  8 ++++++--
 .../installing-openvino-linux.md              |  9 +++++++--
 .../installing-openvino-macos.md              |  9 +++++++--
 .../installing-openvino-windows-fpga.md       |  9 +++++++--
 .../installing-openvino-windows.md            |  9 +++++++--
 get-started-linux.md                          |  9 +++++++--
 model-optimizer/README.md                     |  2 +-
 .../install_prerequisites.bat                 |  3 +++
 .../install_prerequisites.sh                  |  6 +++---
 .../install_prerequisites_tf2.bat             | 18 +++++++++++++++++
 .../install_prerequisites_tf2.sh              | 17 ++++++++++++++++
 model-optimizer/requirements.txt              |  2 +-
 model-optimizer/requirements_tf.txt           |  2 +-
 model-optimizer/requirements_tf2.txt          |  5 +++++
 16 files changed, 116 insertions(+), 22 deletions(-)
 create mode 100644 model-optimizer/install_prerequisites/install_prerequisites_tf2.bat
 create mode 100644 model-optimizer/install_prerequisites/install_prerequisites_tf2.sh
 create mode 100644 model-optimizer/requirements_tf2.txt

diff --git a/docs/MO_DG/prepare_model/Config_Model_Optimizer.md b/docs/MO_DG/prepare_model/Config_Model_Optimizer.md
index a837ceef09c..b5b9853b35c 100644
--- a/docs/MO_DG/prepare_model/Config_Model_Optimizer.md
+++ b/docs/MO_DG/prepare_model/Config_Model_Optimizer.md
@@ -48,14 +48,22 @@ install_prerequisites_caffe.sh
 ```
 install_prerequisites_caffe.bat
 ```
-*   For TensorFlow\* on Linux:
+*   For TensorFlow\* 1.x on Linux:
 ```
 install_prerequisites_tf.sh
 ```
-*   For TensorFlow on Windows:
+*   For TensorFlow 1.x on Windows:
 ```
 install_prerequisites_tf.bat
 ```
+*   For TensorFlow\* 2.x on Linux:
+```
+install_prerequisites_tf2.sh
+```
+*   For TensorFlow 2.x on Windows:
+```
+install_prerequisites_tf2.bat
+```
 *   For MXNet\* on Linux:
 ```
 install_prerequisites_mxnet.sh
@@ -114,7 +122,7 @@ virtualenv -p /usr/bin/python3.6 .env3 --system-site-packages
 virtualenv -p /usr/bin/python3.6 .env3/bin/activate
 ```
 3.  Install all dependencies or only the dependencies for a specific framework:
-    *   To install dependencies for all frameworks:
+    *   To install dependencies for all frameworks except TensorFlow* 2.x:
 ```shell
 pip3 install -r requirements.txt
 ```
@@ -122,9 +130,13 @@ pip3 install -r requirements.txt
 ```shell
 pip3 install -r requirements_caffe.txt
 ```
-    *   To install dependencies only for TensorFlow:
+    *   To install dependencies only for TensorFlow 1.x:
 ```shell
 pip3 install -r requirements_tf.txt
+```
+    *   To install dependencies only for TensorFlow 2.x:
+```shell
+pip3 install -r requirements_tf2.txt
 ```
     *   To install dependencies only for MXNet:
 ```shell
diff --git a/docs/MO_DG/prepare_model/convert_model/Convert_Model_From_TensorFlow.md b/docs/MO_DG/prepare_model/convert_model/Convert_Model_From_TensorFlow.md
index 42c1d487b72..4e07e974cd9 100644
--- a/docs/MO_DG/prepare_model/convert_model/Convert_Model_From_TensorFlow.md
+++ b/docs/MO_DG/prepare_model/convert_model/Convert_Model_From_TensorFlow.md
@@ -295,6 +295,7 @@ python3 mo_tf.py --input_model custom_model.pb --tensorflow_custom_layer_librari
 
 ## Convert TensorFlow* 2 Models <a name="Convert_From_TF2X"></a>
 
+In order to convert TensorFlow* 2 models, installation of dependencies from `requirements_tf2.txt` is required.
 TensorFlow* 2.X officially supports two model formats: SavedModel and Keras H5 (or HDF5).    
 Below are the instructions on how to convert each of them.
 
@@ -331,6 +332,15 @@ model = tf.keras.models.load_model('model.h5')
 tf.saved_model.save(model,'model')
 ```
 
+The Keras H5 model with a custom layer has specifics to be converted into SavedModel format.
+For example, the model with a custom layer `CustomLayer` from `custom_layer.py` is converted as follows:
+```python
+import tensorflow as tf
+from custom_layer import CustomLayer
+model = tf.keras.models.load_model('model.h5', custom_objects={'CustomLayer': CustomLayer})
+tf.saved_model.save(model,'model')
+```
+
 Then follow the above instructions for the SavedModel format. 
 
 > **NOTE:** Do not use other hacks to resave TensorFlow* 2 models into TensorFlow* 1 formats.    
diff --git a/docs/install_guides/installing-openvino-linux-fpga.md b/docs/install_guides/installing-openvino-linux-fpga.md
index 50e41b5922a..d36da473a04 100644
--- a/docs/install_guides/installing-openvino-linux-fpga.md
+++ b/docs/install_guides/installing-openvino-linux-fpga.md
@@ -197,7 +197,7 @@ You can choose to either configure all supported frameworks at once **OR** confi
 cd /opt/intel/openvino/deployment_tools/model_optimizer/install_prerequisites 
 ```
 2.  Run the script to configure the Model Optimizer for Caffe,
-    TensorFlow, MXNet, Kaldi\*, and ONNX:
+    TensorFlow 1.x, MXNet, Kaldi\*, and ONNX:
 ```sh
 sudo ./install_prerequisites.sh
 ```
@@ -215,10 +215,14 @@ cd /opt/intel/openvino/deployment_tools/model_optimizer/install_prerequisites
 ```sh
 sudo ./install_prerequisites_caffe.sh
 ```
--   For **TensorFlow**:
+-   For **TensorFlow 1.x**:
 ```sh
 sudo ./install_prerequisites_tf.sh
 ```
+-   For **TensorFlow 2.x**:
+```sh
+sudo ./install_prerequisites_tf2.sh
+```
 -   For **MXNet**:
 ```sh
 sudo ./install_prerequisites_mxnet.sh
diff --git a/docs/install_guides/installing-openvino-linux.md b/docs/install_guides/installing-openvino-linux.md
index 1bf3dffe1f7..15431faf58d 100644
--- a/docs/install_guides/installing-openvino-linux.md
+++ b/docs/install_guides/installing-openvino-linux.md
@@ -215,7 +215,7 @@ You can choose to either configure all supported frameworks at once **OR** confi
 cd /opt/intel/openvino/deployment_tools/model_optimizer/install_prerequisites
 ```
 2.  Run the script to configure the Model Optimizer for Caffe,
-    TensorFlow, MXNet, Kaldi\*, and ONNX:
+    TensorFlow 1.x, MXNet, Kaldi\*, and ONNX:
 ```sh
 sudo ./install_prerequisites.sh
 ```
@@ -235,11 +235,16 @@ cd /opt/intel/openvino/deployment_tools/model_optimizer/install_prerequisites
    sudo ./install_prerequisites_caffe.sh
    ```
 
-   - For **TensorFlow**:
+   - For **TensorFlow 1.x**:
    ```sh
    sudo ./install_prerequisites_tf.sh
    ```
 
+   - For **TensorFlow 2.x**:
+   ```sh
+   sudo ./install_prerequisites_tf2.sh
+   ```
+
    - For **MXNet**:
    ```sh
    sudo ./install_prerequisites_mxnet.sh
diff --git a/docs/install_guides/installing-openvino-macos.md b/docs/install_guides/installing-openvino-macos.md
index 9da8173f28f..5fb5fd8e5c6 100644
--- a/docs/install_guides/installing-openvino-macos.md
+++ b/docs/install_guides/installing-openvino-macos.md
@@ -181,7 +181,7 @@ You can choose to either configure the Model Optimizer for all supported framewo
    cd /opt/intel/openvino/deployment_tools/model_optimizer/install_prerequisites
    ```
 
-2. Run the script to configure the Model Optimizer for Caffe, TensorFlow, MXNet, Kaldi\*, and ONNX:
+2. Run the script to configure the Model Optimizer for Caffe, TensorFlow 1.x, MXNet, Kaldi\*, and ONNX:
    ```sh
    sudo ./install_prerequisites.sh
    ```
@@ -202,11 +202,16 @@ Configure individual frameworks separately **ONLY** if you did not select **Opti
    sudo ./install_prerequisites_caffe.sh
    ```
 
-   - For **TensorFlow**:
+   - For **TensorFlow 1.x**:
    ```sh
    sudo ./install_prerequisites_tf.sh
    ```
 
+   - For **TensorFlow 2.x**:
+   ```sh
+   sudo ./install_prerequisites_tf2.sh
+   ```
+
    - For **MXNet**:
    ```sh
    sudo ./install_prerequisites_mxnet.sh
diff --git a/docs/install_guides/installing-openvino-windows-fpga.md b/docs/install_guides/installing-openvino-windows-fpga.md
index d2ca21de7c4..2a5e6db38fb 100644
--- a/docs/install_guides/installing-openvino-windows-fpga.md
+++ b/docs/install_guides/installing-openvino-windows-fpga.md
@@ -190,7 +190,7 @@ Type commands in the opened window:
 cd C:\Program Files (x86)\IntelSWTools\openvino\deployment_tools\model_optimizer\install_prerequisites
 ```
 
-3. Run the following batch file to configure the Model Optimizer for Caffe\*, TensorFlow\*, MXNet\*, Kaldi\*, and ONNX\*:<br>
+3. Run the following batch file to configure the Model Optimizer for Caffe\*, TensorFlow\* 1.x, MXNet\*, Kaldi\*, and ONNX\*:<br>
 ```sh
 install_prerequisites.bat
 ```
@@ -209,11 +209,16 @@ cd C:\Program Files (x86)\IntelSWTools\openvino\deployment_tools\model_optimizer
    install_prerequisites_caffe.bat
    ```
 
-   * For **TensorFlow**:<br>
+   * For **TensorFlow 1.x**:<br>
    ```sh
    install_prerequisites_tf.bat
    ```
 
+   * For **TensorFlow 2.x**:<br>
+   ```sh
+   install_prerequisites_tf2.bat
+   ```
+
    * For **MXNet**:<br>
    ```sh
    install_prerequisites_mxnet.bat
diff --git a/docs/install_guides/installing-openvino-windows.md b/docs/install_guides/installing-openvino-windows.md
index 6eb7e87da98..ff7d9324887 100644
--- a/docs/install_guides/installing-openvino-windows.md
+++ b/docs/install_guides/installing-openvino-windows.md
@@ -186,7 +186,7 @@ Type commands in the opened window:
 cd C:\Program Files (x86)\IntelSWTools\openvino\deployment_tools\model_optimizer\install_prerequisites
 ```
 
-3. Run the following batch file to configure the Model Optimizer for Caffe\*, TensorFlow\*, MXNet\*, Kaldi\*, and ONNX\*:<br>
+3. Run the following batch file to configure the Model Optimizer for Caffe\*, TensorFlow\* 1.x, MXNet\*, Kaldi\*, and ONNX\*:<br>
 ```sh
 install_prerequisites.bat
 ```
@@ -205,11 +205,16 @@ cd C:\Program Files (x86)\IntelSWTools\openvino\deployment_tools\model_optimizer
    install_prerequisites_caffe.bat
    ```
 
-   * For **TensorFlow**:<br>
+   * For **TensorFlow 1.x**:<br>
    ```sh
    install_prerequisites_tf.bat
    ```
 
+   * For **TensorFlow 2.x**:<br>
+   ```sh
+   install_prerequisites_tf2.bat
+   ```
+
    * For **MXNet**:<br>
    ```sh
    install_prerequisites_mxnet.bat
diff --git a/get-started-linux.md b/get-started-linux.md
index 661fc4ec8dd..9bac4911198 100644
--- a/get-started-linux.md
+++ b/get-started-linux.md
@@ -56,7 +56,7 @@ If you see error messages, check for any missing dependencies.
 cd <OPENVINO_DIR>/model_optimizer/install_prerequisites
 ```
 2.  Run the script to configure the Model Optimizer for Caffe,
-    TensorFlow, MXNet, Kaldi\*, and ONNX:
+    TensorFlow 1.x, MXNet, Kaldi\*, and ONNX:
 ```sh
 sudo ./install_prerequisites.sh
 ```
@@ -77,11 +77,16 @@ cd <OPENVINO_DIR>/model_optimizer/install_prerequisites
    sudo ./install_prerequisites_caffe.sh
    ```
 
-   - For **TensorFlow**:
+   - For **TensorFlow 1.x**:
    ```sh
    sudo ./install_prerequisites_tf.sh
    ```
 
+   - For **TensorFlow 2.x**:
+   ```sh
+   sudo ./install_prerequisites_tf2.sh
+   ```
+
    - For **MXNet**:
    ```sh
    sudo ./install_prerequisites_mxnet.sh
diff --git a/model-optimizer/README.md b/model-optimizer/README.md
index 260149bee92..baee24120d9 100644
--- a/model-optimizer/README.md
+++ b/model-optimizer/README.md
@@ -195,4 +195,4 @@ of the tool and can not be applied to the current version of Model Optimizer.
     cat requirements_file | docker run -i --rm pyupio/safety safety check --stdin
 </pre>
 
-> **NOTE**: here <code>requirements_file</code> is one of the following: <code>requirements.txt</code>, <code>requirements_caffe.txt</code>, <code>requirements_tf.txt</code>, <code>requirements_mxnet.txt</code>, <code>requirements_dev.txt</code>.
+> **NOTE**: here <code>requirements_file</code> is one of the following: <code>requirements.txt</code>, <code>requirements_caffe.txt</code>, <code>requirements_tf.txt</code>, <code>requirements_tf2.txt</code>, <code>requirements_mxnet.txt</code>, <code>requirements_dev.txt</code>.
diff --git a/model-optimizer/install_prerequisites/install_prerequisites.bat b/model-optimizer/install_prerequisites/install_prerequisites.bat
index c629779f77c..6ed50a5d3eb 100644
--- a/model-optimizer/install_prerequisites/install_prerequisites.bat
+++ b/model-optimizer/install_prerequisites/install_prerequisites.bat
@@ -55,6 +55,9 @@ IF /I "%1%" EQU "" (
     IF /I "%1%" EQU "tf" (
         set postfix=_tf
     ) ELSE (
+    IF /I "%1%" EQU "tf2" (
+        set postfix=_tf2
+    ) ELSE (
     IF /I "%1%" EQU "mxnet" (
         set postfix=_mxnet
     ) ELSE (
diff --git a/model-optimizer/install_prerequisites/install_prerequisites.sh b/model-optimizer/install_prerequisites/install_prerequisites.sh
index a36da204490..6060158f96b 100755
--- a/model-optimizer/install_prerequisites/install_prerequisites.sh
+++ b/model-optimizer/install_prerequisites/install_prerequisites.sh
@@ -30,7 +30,7 @@ V_ENV=0
 
 for ((i=1;i <= $#;i++)) {
     case "${!i}" in
-        caffe|tf|mxnet|kaldi|onnx)
+        caffe|tf|tf2|mxnet|kaldi|onnx)
             postfix="_"$1""
             ;;
         "venv")
@@ -39,7 +39,7 @@ for ((i=1;i <= $#;i++)) {
         *)
             if [[ "$1" != "" ]]; then
                 echo "\""${!i}"\" is unsupported parameter"
-                echo $"Usage: $0 {caffe|tf|mxnet|kaldi|onnx} {venv}"
+                echo $"Usage: $0 {caffe|tf|tf2|mxnet|kaldi|onnx} {venv}"
                 exit 1
             fi
             ;;
@@ -93,5 +93,5 @@ else
     fi
     echo [WARNING] All Model Optimizer dependencies are installed globally.
     echo [WARNING] If you want to keep Model Optimizer in separate sandbox
-    echo [WARNING] run install_prerequisites.sh venv "{caffe|tf|mxnet|kaldi|onnx}"
+    echo [WARNING] run install_prerequisites.sh venv "{caffe|tf|tf2|mxnet|kaldi|onnx}"
 fi
diff --git a/model-optimizer/install_prerequisites/install_prerequisites_tf2.bat b/model-optimizer/install_prerequisites/install_prerequisites_tf2.bat
new file mode 100644
index 00000000000..c6220699926
--- /dev/null
+++ b/model-optimizer/install_prerequisites/install_prerequisites_tf2.bat
@@ -0,0 +1,18 @@
+@echo off
+:: Copyright (C) 2018-2020 Intel Corporation
+::
+::  Licensed under the Apache License, Version 2.0 (the "License");
+::  you may not use this file except in compliance with the License.
+::  You may obtain a copy of the License at
+::
+::       http://www.apache.org/licenses/LICENSE-2.0
+::
+::  Unless required by applicable law or agreed to in writing, software
+::  distributed under the License is distributed on an "AS IS" BASIS,
+::  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+::  See the License for the specific language governing permissions and
+::  limitations under the License.
+
+:: Check if Python is installed
+CALL install_prerequisites.bat tf2
+
diff --git a/model-optimizer/install_prerequisites/install_prerequisites_tf2.sh b/model-optimizer/install_prerequisites/install_prerequisites_tf2.sh
new file mode 100644
index 00000000000..3b8abcc3cc0
--- /dev/null
+++ b/model-optimizer/install_prerequisites/install_prerequisites_tf2.sh
@@ -0,0 +1,17 @@
+#!/bin/bash
+
+# Copyright (C) 2018-2020 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+source install_prerequisites.sh tf2
+
diff --git a/model-optimizer/requirements.txt b/model-optimizer/requirements.txt
index 9906a8edf2e..e8069df734d 100644
--- a/model-optimizer/requirements.txt
+++ b/model-optimizer/requirements.txt
@@ -1,4 +1,4 @@
-tensorflow>=1.15.2
+tensorflow>=1.15.2,<2.0
 mxnet>=1.0.0,<=1.5.1
 networkx>=1.11
 numpy>=1.13.0
diff --git a/model-optimizer/requirements_tf.txt b/model-optimizer/requirements_tf.txt
index 61b88bd3832..ef7e24ed235 100644
--- a/model-optimizer/requirements_tf.txt
+++ b/model-optimizer/requirements_tf.txt
@@ -1,4 +1,4 @@
-tensorflow>=1.15.2
+tensorflow>=1.15.2,<2.0
 networkx>=1.11
 numpy>=1.13.0
 test-generator==0.1.1
diff --git a/model-optimizer/requirements_tf2.txt b/model-optimizer/requirements_tf2.txt
new file mode 100644
index 00000000000..75331fe4f5a
--- /dev/null
+++ b/model-optimizer/requirements_tf2.txt
@@ -0,0 +1,5 @@
+tensorflow>=2.0
+networkx>=1.11
+numpy>=1.13.0
+test-generator==0.1.1
+defusedxml>=0.5.0