Repository: mindee/doctr

Issue:
Inconsistency in TF inference for grouped convolutions on CPU
## 🐛 Troubles with grouped convolutions on CPU

It seems that TF has troubles with fused grouped convolutions on CPU using `keras.Model.predict` while using the `call` method with `training=False` does run smoothly.
This issue is only happening on CPU, on GPU the script runs smoothly.

## To Reproduce

```python
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import tensorflow as tf
from doctr.models import mobilenet_v3_large

samples = tf.zeros((1, 512, 512, 3), dtype=tf.float32)
model = mobilenet_v3_large(input_shape=(512, 512, 3))

# This works
out = model.call(samples, training=False)
# And this throws an error
out = model.predict(samples)
```
which yields
```
---------------------------------------------------------------------------
UnimplementedError                        Traceback (most recent call last)
<ipython-input-1-d329ff911697> in <module>
     11 out = model.call(samples, training=False)
     12 # And this throws an error
---> 13 out = model.predict(samples)

~/miniconda3/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py in predict(self, x, batch_size, verbose, steps, callbacks, max_queue_size, workers, use_multiprocessing)
   1725           for step in data_handler.steps():
   1726             callbacks.on_predict_batch_begin(step)
-> 1727             tmp_batch_outputs = self.predict_function(iterator)
   1728             if data_handler.should_sync:
   1729               context.async_wait()

~/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
    887 
    888       with OptionalXlaContext(self._jit_compile):
--> 889         result = self._call(*args, **kwds)
    890 
    891       new_tracing_count = self.experimental_get_tracing_count()

~/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
    954               *args, **kwds)
    955       # If we did not create any variables the trace we have is good enough.
--> 956       return self._concrete_stateful_fn._call_flat(
    957           filtered_flat_args, self._concrete_stateful_fn.captured_inputs)  # pylint: disable=protected-access
    958 

~/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _call_flat(self, args, captured_inputs, cancellation_manager)
   1958         and executing_eagerly):
   1959       # No tape is watching; skip to running the function.
-> 1960       return self._build_call_outputs(self._inference_function.call(
   1961           ctx, args, cancellation_manager=cancellation_manager))
   1962     forward_backward = self._select_forward_and_backward_functions(

~/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/function.py in call(self, ctx, args, cancellation_manager)
    589       with _InterpolateFunctionError(self):
    590         if cancellation_manager is None:
--> 591           outputs = execute.execute(
    592               str(self.signature.name),
    593               num_outputs=self._num_outputs,

~/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     57   try:
     58     ctx.ensure_initialized()
---> 59     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     60                                         inputs, attrs, num_outputs)
     61   except core._NotOkStatusException as e:

UnimplementedError:  Fused conv implementation does not support grouped convolutions for now.
	 [[node mobile_net_v3/inverted_0/sequential/activation_1/Relu (defined at /home/fg/Documents/doctr/doctr/models/backbones/mobilenet/tensorflow.py:140) ]] [Op:__inference_predict_function_6941]

Function call stack:
predict_function

```

## Expected behavior

The script should run smoothly

## Environment

 - DocTR Version: 0.3.1
 - TensorFlow Version: 2.5.0
 - PyTorch & torchvision versions: 1.9.0 (torchvision 0.10.0)
 - OpenCV Version: 4.5.1
 - OS: Ubuntu 20.04.3 LTS
 - How you installed DocTR: source
 - Python version: 3.8
 - CUDA/cuDNN version: CUDA 11.4.100 (cuDNN 8.2.0)
 - GPU models and configuration: NVIDIA GeForce RTX 2070 with Max-Q Design (driver 470.57.02)

Inconsistency in TF inference for grouped convolutions on CPU
## 🐛 Troubles with grouped convolutions on CPU

It seems that TF has troubles with fused grouped convolutions on CPU using `keras.Model.predict` while using the `call` method with `training=False` does run smoothly.
This issue is only happening on CPU, on GPU the script runs smoothly.

## To Reproduce

```python
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import tensorflow as tf
from doctr.models import mobilenet_v3_large

samples = tf.zeros((1, 512, 512, 3), dtype=tf.float32)
model = mobilenet_v3_large(input_shape=(512, 512, 3))

# This works
out = model.call(samples, training=False)
# And this throws an error
out = model.predict(samples)
```
which yields
```
---------------------------------------------------------------------------
UnimplementedError                        Traceback (most recent call last)
<ipython-input-1-d329ff911697> in <module>
     11 out = model.call(samples, training=False)
     12 # And this throws an error
---> 13 out = model.predict(samples)

~/miniconda3/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py in predict(self, x, batch_size, verbose, steps, callbacks, max_queue_size, workers, use_multiprocessing)
   1725           for step in data_handler.steps():
   1726             callbacks.on_predict_batch_begin(step)
-> 1727             tmp_batch_outputs = self.predict_function(iterator)
   1728             if data_handler.should_sync:
   1729               context.async_wait()

~/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
    887 
    888       with OptionalXlaContext(self._jit_compile):
--> 889         result = self._call(*args, **kwds)
    890 
    891       new_tracing_count = self.experimental_get_tracing_count()

~/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
    954               *args, **kwds)
    955       # If we did not create any variables the trace we have is good enough.
--> 956       return self._concrete_stateful_fn._call_flat(
    957           filtered_flat_args, self._concrete_stateful_fn.captured_inputs)  # pylint: disable=protected-access
    958 

~/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _call_flat(self, args, captured_inputs, cancellation_manager)
   1958         and executing_eagerly):
   1959       # No tape is watching; skip to running the function.
-> 1960       return self._build_call_outputs(self._inference_function.call(
   1961           ctx, args, cancellation_manager=cancellation_manager))
   1962     forward_backward = self._select_forward_and_backward_functions(

~/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/function.py in call(self, ctx, args, cancellation_manager)
    589       with _InterpolateFunctionError(self):
    590         if cancellation_manager is None:
--> 591           outputs = execute.execute(
    592               str(self.signature.name),
    593               num_outputs=self._num_outputs,

~/miniconda3/lib/python3.8/site-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     57   try:
     58     ctx.ensure_initialized()
---> 59     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     60                                         inputs, attrs, num_outputs)
     61   except core._NotOkStatusException as e:

UnimplementedError:  Fused conv implementation does not support grouped convolutions for now.
	 [[node mobile_net_v3/inverted_0/sequential/activation_1/Relu (defined at /home/fg/Documents/doctr/doctr/models/backbones/mobilenet/tensorflow.py:140) ]] [Op:__inference_predict_function_6941]

Function call stack:
predict_function

```

## Expected behavior

The script should run smoothly

## Environment

 - DocTR Version: 0.3.1
 - TensorFlow Version: 2.5.0
 - PyTorch & torchvision versions: 1.9.0 (torchvision 0.10.0)
 - OpenCV Version: 4.5.1
 - OS: Ubuntu 20.04.3 LTS
 - How you installed DocTR: source
 - Python version: 3.8
 - CUDA/cuDNN version: CUDA 11.4.100 (cuDNN 8.2.0)
 - GPU models and configuration: NVIDIA GeForce RTX 2070 with Max-Q Design (driver 470.57.02)



Discussion:
cf. https://github.com/tensorflow/tensorflow/issues/51825
I seems that they are not likely to fix this in the near future...
Yes unfortunately, but we'll have to keep this one open for reference then :/
@charlesmindee @fg-mindee same problem if we try to export to SavedModel format (only mobilenet arch)
So the upstream problem about grouped conv on CPU is now resolved for TF 2.9 👍 
We're gonna have to look at updating our version specifiers
@frgfm sounds good :) 
cf. https://github.com/tensorflow/tensorflow/issues/51825
I seems that they are not likely to fix this in the near future...
Yes unfortunately, but we'll have to keep this one open for reference then :/
@charlesmindee @fg-mindee same problem if we try to export to SavedModel format (only mobilenet arch)
So the upstream problem about grouped conv on CPU is now resolved for TF 2.9 👍 
We're gonna have to look at updating our version specifiers
@frgfm sounds good :) 

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -44,7 +44,7 @@
     "scipy>=1.4.0",
     "h5py>=3.1.0",
     "opencv-python>=3.4.5.20",
-    "tensorflow>=2.4.0",
+    "tensorflow>=2.9.0,<3.0.0",  # cf. https://github.com/mindee/doctr/issues/454
     "pypdfium2>=2.1.0, <3.0.0",  # cf. https://github.com/mindee/doctr/issues/947
     "pyclipper>=1.2.0",
     "shapely>=1.6.0",
@@ -59,7 +59,6 @@
     "tqdm>=4.30.0",
     "tensorflow-addons>=0.13.0",
     "rapidfuzz>=1.6.0",
-    "keras<2.7.0",  # cf. https://github.com/mindee/doctr/pull/579
     "tf2onnx>=1.9.2",
     "huggingface-hub>=0.4.0",
     # Testing
@@ -115,14 +114,12 @@ def deps_list(*pkgs):
 extras["tf"] = deps_list(
     "tensorflow",
     "tensorflow-addons",
-    "keras",
     "tf2onnx",
 )
 
 extras["tf-cpu"] = deps_list(
     "tensorflow-cpu",
     "tensorflow-addons",
-    "keras",
     "tf2onnx",
 )
 


Test patch:
diff --git a/tests/tensorflow/test_models_classification_tf.py b/tests/tensorflow/test_models_classification_tf.py
--- a/tests/tensorflow/test_models_classification_tf.py
+++ b/tests/tensorflow/test_models_classification_tf.py
@@ -97,9 +97,10 @@ def test_crop_orientation_model(mock_text_box):
         # Name:'res_net_4/magc/transform/conv2d_289/Conv2D:0_nchwc'
         # Status Message: Input channels C is not equal to kernel channels * group. C: 32 kernel channels: 256 group: 1
         #["magc_resnet31", (32, 32, 3), (126,)],
-        ["mobilenet_v3_small", (512, 512, 3), (126,)],
-        ["mobilenet_v3_large", (512, 512, 3), (126,)],
-        ["mobilenet_v3_small_orientation", (128, 128, 3), (4,)],
+        # Disabled for now
+        # ["mobilenet_v3_small", (512, 512, 3), (126,)],
+        # ["mobilenet_v3_large", (512, 512, 3), (126,)],
+        # ["mobilenet_v3_small_orientation", (128, 128, 3), (4,)],
     ],
 )
 def test_models_saved_model_export(arch_name, input_shape, output_size):