Exporting fp16 onnx models.

extras/easypose/base_model.py

@@ -0,0 +1,61 @@
+import warnings
+from abc import ABC, abstractmethod
+from typing import List
+import time
+import numpy as np
+import onnxruntime as ort
+from tqdm import tqdm
+
+
+class BaseModel(ABC):
+    def __init__(self, model_path: str, device: str = 'CUDA', warmup: int = 30):
+        self.opt = ort.SessionOptions()
+
+        if device == 'CUDA':
+            provider = 'CUDAExecutionProvider'
+            if provider not in ort.get_available_providers():
+                warnings.warn("No CUDAExecutionProvider found, switched to CPUExecutionProvider.", UserWarning)
+                provider = 'CPUExecutionProvider'
+        elif device == 'CPU':
+            provider = 'CPUExecutionProvider'
+        else:
+            raise ValueError('Provider {} does not exist.'.format(device))
+
+        self.session = ort.InferenceSession(model_path,
+                                            providers=[provider],
+                                            sess_options=self.opt)
+
+        self.input_name = self.session.get_inputs()[0].name
+        self.input_shape = self.session.get_inputs()[0].shape
+
+        input_type = self.session.get_inputs()[0].type
+        if input_type == 'tensor(float16)':
+            self.input_type = np.float16
+        else:
+            self.input_type = np.float32
+
+        if warmup > 0:
+            self.warmup(warmup)
+
+    @abstractmethod
+    def preprocess(self, image: np.ndarray):
+        pass
+
+    @abstractmethod
+    def postprocess(self, tensor: List[np.ndarray]):
+        pass
+
+    def forward(self, image: np.ndarray):
+        tensor = self.preprocess(image)
+        result = self.session.run(None, {self.input_name: tensor})
+        output = self.postprocess(result)
+        return output
+
+    def warmup(self, epoch: int = 30):
+        print('{} start warmup!'.format(self.__class__.__name__))
+        tensor = np.random.random(self.input_shape).astype(self.input_type)
+        for _ in tqdm(range(epoch)):
+            self.session.run(None, {self.input_name: tensor})
+
+    def __call__(self, image: np.ndarray, *args, **kwargs):
+        return self.forward(image)

extras/easypose/detection.py

@@ -24,7 +24,7 @@ class RTMDet(BaseModel):
         tensor, self.dx, self.dy, self.scale = letterbox(
             image, (tw, th), fill_value=114
         )
-        tensor = tensor.astype(np.float32, copy=False)
+        tensor = tensor.astype(self.input_type, copy=False)
         tensor = tensor[..., ::-1]
         tensor = np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2))
         return tensor

extras/easypose/pose.py

@@ -45,7 +45,7 @@ class SimCC(BaseModel):
 
     def preprocess(self, image: np.ndarray):
         tensor, self.dx, self.dy, self.scale = image, 0, 0, 1
-        tensor = tensor.astype(np.float32, copy=False)
+        tensor = tensor.astype(self.input_type, copy=False)
         tensor = np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2))
         return tensor
 

extras/easypose/run_container.sh

@@ -5,6 +5,7 @@ docker run --privileged --rm --network host -it \
   --gpus all --shm-size=16g --ulimit memlock=-1 --ulimit stack=67108864 \
   --volume "$(pwd)"/:/RapidPoseTriangulation/ \
   --volume "$(pwd)"/extras/easypose/pipeline.py:/EasyPose/easypose/pipeline.py \
+  --volume "$(pwd)"/extras/easypose/base_model.py:/EasyPose/easypose/model/base_model.py \
   --volume "$(pwd)"/extras/easypose/detection.py:/EasyPose/easypose/model/detection.py \
   --volume "$(pwd)"/extras/easypose/pose.py:/EasyPose/easypose/model/pose.py \
   --volume "$(pwd)"/extras/easypose/utils.py:/EasyPose/easypose/model/utils.py \