pip install -i https://mirrors.cloud.tencent.com/pypi/simple tiinfer 

from typing import Dict
import tiinfer

### 日志相关代码 ###
# 增加以下代码，才能将框架记录的请求日志打印到终端
import logging
import sys
logger_formatter = "%(asctime)s %(levelname)s    %(module)s:%(lineno)d    %(message)s"
logging.basicConfig(stream=sys.stdout, format=logger_formatter, level=logging.DEBUG)
### 日志相关代码 ###

# tiinfer支持原生的mosec.Worker
class HelloWorld(mosec.Worker):
    def forward(self, req: Dict) -> Dict:
        return {"hello": f"world. raw req is {req}"}

# 启动两个进程同时处理请求
tiinfer.append_worker(HelloWorld, num=2)

TI_MODEL_DIR=`pwd` python3 -m tiinfer --timeout 30000

> curl -X POST -d '{"key": "values"}' http://127.0.0.1:8501/v1/models/m:predict 
{
  "hello": "world. raw is {'key': 'values'}"
}

class Worker:
    def __init__(self)
    
    def forward(data)

def append_worker(
    worker: Type[mosec.Worker],
    num: int = 1,
    start_method: str = "spawn",
    env: Union[None, List[Dict[str, str]]] = None,
) -> None
"""
    worker: 继承mosec.Worker并实现了forward方法的处理Worker
    num: 并行计算的进程数(>=1)
    start_method: 进程启动方法("spawn" 或 "fork")
    env: 进程启动前的一些环境变量
"""

import tiinfer
from mosec import Worker
from typing import Dict, Any


class MyPreprocess(Worker):
    def forward(self, data: Dict) -> Any:
      # 输入是json转化来的Dict，进行一些必要的预处理

      
class MyPredict(Worker):
    def __int__(self):
      super.__init__()
      # 读取并加载模型
      
    def forward(self, data: Any) -> Any:
      # 输入是预处理的结果，通过调用模型推理后，得到推理结果


class MyPostprocess(Worker):
    def forward(self, data: Any) -> Dict:
      # 输入为推理结果，通过一些后处理手段转化为Dict作为Json返回给调用方
      

# 编排处理过程: 1x预处理Worker -> 2x推理Worker -> 1x后处理Worker
tiinfer.append_worker(MyPreprocess, 1)
tiinfer.append_worker(MyPredict, 2)
tiinfer.append_worker(MyPostprocess, 1)


import logging
import os
import time
from typing import Dict, List
from urllib.request import urlretrieve

import cv2  # type: ignore
import numpy as np  # type: ignore

import torch  # type: ignore

import tiinfer
import tiinfer.utils
import mosec


### 日志相关代码 ###
# 增加以下代码，才能将框架记录的请求日志打印到终端
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
formatter = logging.Formatter(
    "%(asctime)s - %(process)d - %(levelname)s - %(filename)s:%(lineno)s - %(message)s"
)
sh = logging.StreamHandler()
sh.setFormatter(formatter)
logger.addHandler(sh)
### 日志相关代码 ###

# 预处理过程对输入的base64字符串进行解码，并按照模型要求做一些缩放操作
class Pre(mosec.Worker):
    def forward(self, req: Dict) -> cv2.Mat:
        # 预处理的输入数据是json转化而来的Dict
        img_base64_bytes = req["image"]
        img = tiinfer.utils.image_to_cv2_mat(img_base64_bytes)
        # bgr -> rgb
        img = img[:, :, ::-1]
        # 对图片做一些预处理
        img = cv2.resize(img, (256, 256))
        crop_img = (
            img[16 : 16 + 224, 16 : 16 + 224].astype(np.float32) / 255
        )  # center crop
        crop_img -= [0.485, 0.456, 0.406]
        crop_img /= [0.229, 0.224, 0.225]
        crop_img = np.transpose(crop_img, (2, 0, 1))
        return crop_img

# 加载模型，并对预处理后的结果做推理操作，并转换为最终结果传给调用方·
class Infer(mosec.Worker):
    def __init__(self) -> None:
        super().__init__()
        # 获取当前所在目录，需要据此加载模型
        self.root_path = tiinfer.TI_MODEL_DIR
        # 优先使用GPU
        self.device = (
            torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        )
        ### 非加速模型的加载 开始###
        # 模型存在model目录下面
        model_file = os.path.join(self.root_path, "model/resNet50.pt")
        ### 非加速模型的加载 结束###
        
        ### 加速模型的加载 开始###
        # #加速模型需要额外导入tiacc_inference
        # import tiacc_inference
        # model_file = os.path.join(self.root_path, "model/tiacc.pt")
        ### 加速模型的加载 结束###    
        
        # 加载模型
        self.model = torch.jit.load(model_file)
        self.model.eval()

        # 分类需要知道最终的类别
        self.categories = load_categories()

    def forward(self, img: cv2.Mat) -> Dict:
        with torch.no_grad():
            batch = torch.stack(
                [torch.tensor(arr, device=self.device) for arr in [img]]
            )
            pred_results = self.model(batch)
            prob = torch.nn.functional.softmax(pred_results, dim=1)
            top1_prob, top1_catid = torch.topk(prob, 1)
            return [
                {
                    "confidence": top1_prob[i].tolist()[0],
                    "pred": self.categories[top1_catid[i].tolist()[0]],
                }
                for i in range(top1_prob.size(0))
            ][0]


# 从标签文件中读取标签id对应的类别信息
def load_categories() -> List[str]:
    logging.info("loading categories file...")
    local_filename = "imagenet_classes.txt"
    if not os.path.exists("imagenet_classes.txt"):
        local_filename, _ = urlretrieve(
            "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
        )
    with open(local_filename, encoding="utf8") as file:
        return list(map(lambda x: x.strip(), file.readlines()))

# 编排处理过程：2个进程进行预处理 -> 1个进程进行推理
tiinfer.append_worker(Pre, 2)
tiinfer.append_worker(Infer, 1)