Changes to be able to convert the .pt model to ONNX and TensorRT formats

README.md

@@ -5,6 +5,20 @@ A deep learning framework based on Tensorflow for the training of high performan
 
 *All implementations are re-implementations of published algorithms and thus provided models should not be considered as reference.*
 
+---
+
+This is a fork from [swook/GazeML](https://github.com/swook/GazeML). I am on fixing some issues in importing the model to `.onnx` format and further to TensorRT `.engine` format.
+
+I have tested this on:
+- Ubuntu 20.04
+- NVIDIA GeForce RTX 3080 Laptop GPU
+- Nvidia driver 510.54
+- CUDA 11.6
+- python 3.6
+- tensorflow 1.14
+
+---
+
 This framework currently integrates the following models:
 
 ## ELG
@@ -30,24 +44,26 @@ Deep Pictorial Gaze Estimation
 
 ## Installing dependencies
 
-Run (with `sudo` appended if necessary),
+Follow [these instructions](https://opensource.com/article/20/4/install-python-linux) to download Python 3.7
+
+Use python virtual environment ([venv](https://docs.python.org/3/tutorial/venv.html)) to install python dependencies:
+
+Run the following inside `GazeML` repository root folder:
 ```
-python3 setup.py install
+python3.7 -m venv .venv
+source .venv/bin/activate
 ```
 
-Note that this can be done within a [virtual environment](https://docs.python.org/3/tutorial/venv.html). In this case, the sequence of commands would be similar to:
+Install dependencies:
 ```
-    mkvirtualenv -p $(which python3) myenv
-    python3 setup.py install
+pip install --upgrade pip
+pip install cython
+pip install scipy
+python3 setup.py install
+pip install tensorflow==1.14
+pip install tensorflow-gpu==1.14
 ```
 
-when using [virtualenvwrapper](https://virtualenvwrapper.readthedocs.io/en/latest/).
-
-### Tensorflow
-Tensorflow is assumed to be installed separately, to allow for usage of [custom wheel files](https://github.com/mind/wheels) if necessary.
-
-Please follow the official installation guide for Tensorflow [here](https://www.tensorflow.org/install/).
-
 ## Getting pre-trained weights
 To acquire the pre-trained weights provided with this repository, please run:
 ```

src/core/model.py

@@ -29,7 +29,7 @@ def __init__(self,
                  train_data: Dict[str, BaseDataSource] = {},
                  test_data: Dict[str, BaseDataSource] = {},
                  test_losses_or_metrics: str = None,
-                 use_batch_statistics_at_test: bool = True,
+                 use_batch_statistics_at_test: bool = False,
                  identifier: str = None):
         """Initialize model with data sources and parameters."""
         self._tensorflow_session = tensorflow_session
@@ -324,7 +324,7 @@ def train(self, num_epochs=None, num_steps=None):
                     fetches=fetches,
                     feed_dict={
                         self.is_training: True,
-                        self.use_batch_statistics: True,
+                        self.use_batch_statistics: False,
                     }
                 )
                 self.time.end('train_iteration')
@@ -378,8 +378,15 @@ def inference_generator(self):
                 fetches=fetches,
                 feed_dict={
                     self.is_training: False,
-                    self.use_batch_statistics: True,
+                    self.use_batch_statistics: False,
                 },
             )
+
+
+            # Save saved-model
+            # This model can be used to convert to .onnx and further to .engine
+            tf.saved_model.simple_save(self._tensorflow_session, "tmp",
+                inputs=data_source.output_tensors, outputs=fetches)
+
             outputs['inference_time'] = 1e3*(time.time() - start_time)
             yield outputs

src/datasources/frames.py

@@ -113,7 +113,7 @@ def entry_generator(self, yield_just_one=False):
                     yield {
                         'frame_index': np.int64(current_index),
                         'eye': eye_dict['image'],
-                        'eye_index': np.uint8(i),
+                        'eye_index': np.int64(i),
                     }
 
         finally:

src/models/elg.py

@@ -178,7 +178,7 @@ def _apply_bn(self, tensor):
             tensor,
             scale=True,
             center=True,
-            is_training=self.use_batch_statistics,
+            is_training=False,
             trainable=True,
             data_format=self._data_format,
             updates_collections=None,