Vision Transformer + BentoML

BentoML

https://zhuanlan.zhihu.com/p/495814838

BentoML 是一个用于机器学习模型服务的开源框架，旨在弥合数据科学和 DevOps 之间的差距（gap）。

数据科学家可以使用 BentoMl 轻松打包使用任何 ML 框架训练的模型，并重现该模型以用于生产。

BentoML 协助管理 BentoML 格式打包的模型，并允许 DevOps 将它们部署为任何云平台上的在线 API 服务端点或离线批量推理作业。

为什么选择 BentoML

• 将您的 ML 模型转换为生产就绪 API 非常简单。
• 高性能模型服务，并且全部使用 Python。
• 标准化模型打包和 ML 服务定义以简化部署。
• 支持所有主流的机器学习训练框架。
• 通过Yatai在 Kubernetes 上大规模部署和运行 ML 服务。

下面将演示了如何使用 BentoML 通过 REST API 服务为 sklearn 模型提供服务，然后将模型服务容器化以进行生产部署。

 

Vision Transformer + BentoML

https://github.com/fanqingsong/Pneumonia-Detection-Demo

 

In this project, we showcase the seamless integration of an image detection model into a service using BentoML. Leveraging the power of the pretrained nickmuchi/vit-finetuned-chest-xray-pneumonia model from HuggingFace, users can submit their lung X-ray images for analysis. The model will then determine, with precision, whether the individual has pneumonia or not.

 

from __future__ import annotations

import typing as t

import torch
import pydantic
import PIL.Image
import PIL.ImageOps
import transformers

import bentoml

from save_model import download_model

_ = download_model()

MODEL_ID = "nickmuchi/vit-finetuned-chest-xray-pneumonia"
extractor = transformers.ViTImageProcessor.from_pretrained(MODEL_ID)
model = transformers.AutoModelForImageClassification.from_pretrained(MODEL_ID)
model.eval()

svc = bentoml.Service("pneumonia-classifier")


def preprocess(image: PIL.Image.Image) -> PIL.Image.Image:
    return PIL.ImageOps.exif_transpose(image).convert("RGB")


# /v1/classify 的 JSON 响应，例如 {"class_name": "PNEUMONIA"}。
class Output(pydantic.BaseModel):
    class_name: t.Literal["NORMAL", "PNEUMONIA"]

    @classmethod
    def from_result(cls, logits: torch.Tensor) -> Output:
        # logits 例: tensor([[-2.10, 3.45]])，列 0=NORMAL、列 1=PNEUMONIA，数值越大越倾向该类。
        # id2label 例: {0: "NORMAL", 1: "PNEUMONIA"}
        id2label = model.config.id2label
        top_k = len(id2label)  # 例: 2
        # softmax 后第一张图的概率，例: tensor([0.004, 0.996])
        probs = logits.softmax(-1)[0]
        # 按概率从高到低：scores 例 [0.996, 0.004]，ids 例 [1, 0]
        scores, ids = probs.topk(top_k)
        # ranked 例: [(0.996, "PNEUMONIA"), (0.004, "NORMAL")]
        ranked = [
            (score, id2label[id_]) for score, id_ in zip(scores.tolist(), ids.tolist())
        ]
        # 取最高分标签，例: Output(class_name="PNEUMONIA")
        return cls(class_name=max(ranked, key=lambda item: item[0])[1])


@svc.api(
    input=bentoml.io.Image(),
    output=bentoml.io.JSON(pydantic_model=Output),
    route="/v1/classify",
)
async def classify(image: PIL.Image.Image) -> Output:
    # image 例: RGB 胸片，size=(1858, 1317)
    image = preprocess(image)
    # features 例: {"pixel_values": tensor, shape=[1, 3, 224, 224]}，已归一化到 ViT 输入
    features = extractor(images=image, return_tensors="pt")
    with torch.inference_mode():
        outputs = model(**features)
    # outputs.logits 例: tensor([[-2.10, 3.45]])，再交给 from_result 得到 {"class_name": "PNEUMONIA"}
    return Output.from_result(outputs.logits)