feat: add fp16 inference support (torch/onnx) (#871)

* feat: add fp16 inference in clip_torch

* Revert "feat: add fp16 inference in clip_torch"

This reverts commit 326e265.

* feat: add fp16 inference in clip_torch

* fix: device

* fix: str to torch.dtype

* fix: layernorm

* feat: add fp16 inference in clip_trt

* feat: add fp16 inference in clip_onnx

* fix: housekeeping

* fix: ci

* fix: ci

* fix: ci

* fix: ci and get test path

* fix: dtype amp and gpu test dependency

* fix: layernorm

* fix: cast dtype in visiontransformer

* fix: clip_onnx

* fix: clip_onnx

* fix: convert onnx to fp16

* fix: dtype in preproc images

* fix: dtype in preproc images

* fix: typo

* fix: dtype in clip_torch and fp16 in trt

* fix: remove plain text in trt_test

* fix: test

* fix: typo

* fix: stash

* Revert "fix: stash"

This reverts commit f72fd99.

* fix: for test

* fix: onnx

* fix: for test

* fix: for test

* fix: trt

* fix: convert onnx to fp16 before convert trt

* fix: discard changes in trt

* fix: optimize fp16 test

* fix: move __cast_dtype__

* Revert "fix: move __cast_dtype__"

This reverts commit edf4629.

* fix: ci

Author: OrangeSodahub

.github/workflows/ci.yml

@@ -113,12 +113,11 @@ jobs:
           pip install --no-cache-dir "server/[onnx]"
           pip install --no-cache-dir "server/[transformers]"
           pip install --no-cache-dir "server/[search]"
-          pip install --no-cache-dir "server/[transformers]"
       - name: Test
         id: test
         run: |
           pytest --suppress-no-test-exit-code --cov=clip_client --cov=clip_server --cov-report=xml \
-            -v -s -m "not gpu" ${{ matrix.test-path }}
+            -v -s ${{ matrix.test-path }}
           echo "::set-output name=codecov_flag::cas"
         timeout-minutes: 30
       - name: Check codecov file
@@ -158,6 +157,7 @@ jobs:
           python -m pip install wheel pytest pytest-cov nvidia-pyindex
           pip install -e "client/[test]"
           pip install -e "server/[tensorrt]"
+          pip install -e "server/[onnx]"
           {
             pip install -e "server/[flash-attn]"
           } || {
@@ -168,6 +168,8 @@ jobs:
         run: |
           pytest --suppress-no-test-exit-code --cov=clip_client --cov=clip_server --cov-report=xml \
             -v -s -m "gpu" ./tests/test_tensorrt.py
+          pytest --suppress-no-test-exit-code --cov=clip_client --cov=clip_server --cov-report=xml \
+            -v -s -m "gpu" ./tests/test_simple.py
           echo "::set-output name=codecov_flag::cas"
         timeout-minutes: 30
         env:

server/clip_server/executors/clip_onnx.py

@@ -27,6 +27,7 @@ def __init__(
         minibatch_size: int = 32,
         access_paths: str = '@r',
         model_path: Optional[str] = None,
+        dtype: Optional[str] = None,
         **kwargs,
     ):
         """
@@ -41,8 +42,17 @@ def __init__(
         :param model_path: The path to the model to be used. If not specified, the model will be downloaded or loaded
             from the local cache. Visit https://clip-as-service.jina.ai/user-guides/server/#use-custom-model-for-onnx
             to learn how to finetune custom models.
+        :param dtype: inference data type, if None defaults to 'fp32' if device == 'cpu' else 'fp16'.
         """
         super().__init__(**kwargs)
+        import torch
+
+        if not device:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self._device = device
+        if not dtype:
+            dtype = 'fp32' if self._device in ('cpu', torch.device('cpu')) else 'fp16'
+        self._dtype = dtype
 
         self._minibatch_size = minibatch_size
         self._access_paths = access_paths
@@ -55,18 +65,11 @@ def __init__(
         self._num_worker_preprocess = num_worker_preprocess
         self._pool = ThreadPool(processes=num_worker_preprocess)
 
-        self._model = CLIPOnnxModel(name, model_path)
+        self._model = CLIPOnnxModel(name, model_path, dtype)
         self._tokenizer = Tokenizer(name)
 
         self._image_transform = clip._transform_ndarray(self._model.image_size)
 
-        import torch
-
-        if not device:
-            self._device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        else:
-            self._device = device
-
         # define the priority order for the execution providers
         providers = ['CPUExecutionProvider']
 
@@ -116,6 +119,7 @@ def _preproc_images(self, docs: 'DocumentArray', drop_image_content: bool):
                     preprocess_fn=self._image_transform,
                     return_np=True,
                     drop_image_content=drop_image_content,
+                    dtype=self._dtype,
                 )
 
     def _preproc_texts(self, docs: 'DocumentArray'):

server/clip_server/executors/clip_torch.py

@@ -2,7 +2,7 @@
 import warnings
 from functools import partial
 from multiprocessing.pool import ThreadPool
-from typing import Dict, Optional
+from typing import Dict, Union, Optional
 
 import numpy as np
 import torch
@@ -12,6 +12,7 @@
     set_rank,
     split_img_txt_da,
 )
+from clip_server.helper import __cast_dtype__
 from clip_server.model import clip
 from clip_server.model.clip_model import CLIPModel
 from clip_server.model.tokenization import Tokenizer
@@ -28,6 +29,7 @@ def __init__(
         num_worker_preprocess: int = 4,
         minibatch_size: int = 32,
         access_paths: str = '@r',
+        dtype: Optional[Union[str, torch.dtype]] = None,
         **kwargs,
     ):
         """
@@ -40,6 +42,7 @@ def __init__(
             number if you encounter OOM errors.
         :param access_paths: The access paths to traverse on the input documents to get the images and texts to be
             processed. Visit https://docarray.jina.ai/fundamentals/documentarray/access-elements for more details.
+        :param dtype: inference data type, if None defaults to torch.float32 if device == 'cpu' else torch.float16.
         """
         super().__init__(**kwargs)
 
@@ -52,9 +55,17 @@ def __init__(
             self._access_paths = kwargs['traversal_paths']
 
         if not device:
-            self._device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        else:
-            self._device = device
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self._device = device
+        if isinstance(dtype, str):
+            dtype = __cast_dtype__.get(dtype)
+        elif not dtype:
+            dtype = (
+                torch.float32
+                if self._device in ('cpu', torch.device('cpu'))
+                else torch.float16
+            )
+        self._dtype = dtype
 
         if not self._device.startswith('cuda') and (
             'OMP_NUM_THREADS' not in os.environ
@@ -77,7 +88,9 @@ def __init__(
         self._num_worker_preprocess = num_worker_preprocess
         self._pool = ThreadPool(processes=num_worker_preprocess)
 
-        self._model = CLIPModel(name, device=self._device, jit=jit, **kwargs)
+        self._model = CLIPModel(
+            name, device=self._device, jit=jit, dtype=dtype, **kwargs
+        )
         self._tokenizer = Tokenizer(name)
         self._image_transform = clip._transform_ndarray(self._model.image_size)
 
@@ -96,6 +109,7 @@ def _preproc_images(self, docs: 'DocumentArray', drop_image_content: bool):
                     device=self._device,
                     return_np=False,
                     drop_image_content=drop_image_content,
+                    dtype=self._dtype,
                 )
 
     def _preproc_texts(self, docs: 'DocumentArray'):

server/clip_server/executors/helper.py

@@ -1,8 +1,9 @@
-from typing import Tuple, List, Callable, Any, Dict
+from typing import Tuple, List, Callable, Any, Dict, Union
 import torch
 import numpy as np
 from docarray import Document, DocumentArray
 from docarray.math.distance.numpy import cosine
+from clip_server.helper import __cast_dtype__
 
 
 from clip_server.model.tokenization import Tokenizer
@@ -22,8 +23,12 @@ def preproc_image(
     device: str = 'cpu',
     return_np: bool = False,
     drop_image_content: bool = False,
+    dtype: Union[str, torch.dtype] = torch.float32,
 ) -> Tuple['DocumentArray', Dict]:
 
+    if isinstance(dtype, str):
+        dtype = __cast_dtype__.get(dtype)
+
     tensors_batch = []
 
     for d in da:
@@ -42,7 +47,7 @@ def preproc_image(
         if drop_image_content:
             d.pop('blob', 'tensor')
 
-    tensors_batch = torch.stack(tensors_batch).type(torch.float32)
+    tensors_batch = torch.stack(tensors_batch).type(dtype)
 
     if return_np:
         tensors_batch = tensors_batch.cpu().numpy()

server/clip_server/helper.py

@@ -2,6 +2,7 @@
 import os
 import sys
 import threading
+import torch
 from packaging.version import Version
 from urllib.request import Request, urlopen
 
@@ -19,6 +20,9 @@
 )
 
 
+__cast_dtype__ = {'fp16': torch.float16, 'fp32': torch.float32, 'bf16': torch.bfloat16}
+
+
 def _version_check(package: str = None, github_repo: str = None):
     try:
 

server/clip_server/model/clip_onnx.py

@@ -1,5 +1,5 @@
 import os
-from typing import Dict
+from typing import Dict, Optional
 
 from clip_server.model.pretrained_models import (
     download_model,
@@ -201,8 +201,11 @@
 
 
 class CLIPOnnxModel(BaseCLIPModel):
-    def __init__(self, name: str, model_path: str = None):
+    def __init__(
+        self, name: str, model_path: str = None, dtype: Optional[str] = 'fp32'
+    ):
         super().__init__(name)
+        self._dtype = dtype
         if name in _MODELS:
             if not model_path:
                 cache_dir = os.path.expanduser(
@@ -237,6 +240,22 @@ def __init__(self, name: str, model_path: str = None):
                         f'The given model path {model_path} should be a folder containing both '
                         f'`textual.onnx` and `visual.onnx`.'
                     )
+            if dtype == 'fp16':
+                import onnx
+                from onnxmltools.utils import float16_converter
+
+                _textual_model_fp16 = (
+                    float16_converter.convert_float_to_float16_model_path(
+                        self._textual_path
+                    )
+                )
+                _visual_model_fp16 = (
+                    float16_converter.convert_float_to_float16_model_path(
+                        self._visual_path
+                    )
+                )
+                onnx.save_model(_textual_model_fp16, self._textual_path)
+                onnx.save_model(_visual_model_fp16, self._visual_path)
         else:
             raise RuntimeError(
                 'CLIP model {} not found or not supports ONNX backend; below is a list of all available models:\n{}'.format(

server/clip_server/model/model.py

@@ -15,6 +15,7 @@
 from dataclasses import dataclass
 from typing import Tuple, Union, Optional
 from copy import deepcopy
+from clip_server.helper import __cast_dtype__
 from open_clip.transformer import QuickGELU, LayerNorm, LayerNormFp32, Attention
 from open_clip.timm_model import TimmModel
 from open_clip.factory import _MODEL_CONFIGS
@@ -81,6 +82,11 @@ def __init__(
         super().__init__(image_size, patch_size, output_dim=output_dim, **kwargs)
         self.transformer = Transformer(dtype=dtype, **kwargs)
 
+    def forward(self, x: torch.Tensor):
+        dtype = self.transformer.get_cast_dtype()
+        x = x.to(dtype)
+        return super().forward(x)
+
 
 class TextTransformer(_TextTransformer):
     def __init__(
@@ -435,7 +441,9 @@ def load_openai_model(
     preprocess : Callable[[PIL.Image], torch.Tensor]
         A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
     """
-    if dtype is None:
+    if isinstance(dtype, str):
+        dtype = __cast_dtype__.get(dtype, 'amp')
+    elif dtype is None:
         dtype = (
             torch.float32 if device in ('cpu', torch.device('cpu')) else torch.float16
         )
@@ -550,7 +558,9 @@ def load_openclip_model(
     pretrained_image: bool = False,
     dtype: Optional[Union[str, torch.dtype]] = None,
 ):
-    if dtype is None:
+    if isinstance(dtype, str):
+        dtype = __cast_dtype__.get(dtype)
+    elif dtype is None:
         dtype = (
             torch.float32 if device in ('cpu', torch.device('cpu')) else torch.float16
         )

server/setup.py

@@ -53,6 +53,7 @@
         'onnx': [
             'onnxruntime',
             'onnx',
+            'onnxmltools',
         ]
         + (['onnxruntime-gpu>=1.8.0'] if sys.platform != 'darwin' else []),
         'tensorrt': ['nvidia-tensorrt'],

tests/test_simple.py

@@ -27,6 +27,7 @@ def test_protocols(port_generator, protocol, jit, pytestconfig):
         c.profile(content=f'{pytestconfig.rootdir}/tests/img/00000.jpg')
 
 
+@pytest.mark.gpu
 @pytest.mark.parametrize(
     'inputs',
     [
@@ -48,6 +49,7 @@ def test_plain_inputs(make_flow, inputs):
     )
 
 
+@pytest.mark.gpu
 @pytest.mark.parametrize(
     'inputs',
     [