为torch-model-archiver编写自定义处理器以生成模型.mar文件

以下是实现自定义处理器并生成适配你模型的.mar文件的步骤：

1. 编写自定义处理器文件 #

创建名为custom_handler.py的文件，实现继承自BaseHandler的处理器类，覆盖核心方法以适配你的模型输入输出逻辑：

import torch
import numpy as np
import io
from PIL import Image
from torchvision.transforms import Compose, ToTensor, Normalize
from ts.torch_handler.base_handler import BaseHandler

class MNISTHandler(BaseHandler):
    def initialize(self, context):
        # 加载模型与设备配置
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = self._load_model(context)
        self.model.to(self.device)
        self.model.eval()
        
        # 定义预处理变换（匹配MNIST数据集标准）
        self.transform = Compose([
            ToTensor(),
            Normalize((0.1307,), (0.3081,))
        ])
        
        self.initialized = True

    def preprocess(self, data):
        # 处理输入数据：支持图像字节或像素数组格式
        inputs = []
        for item in data:
            # 提取请求中的数据体
            image = item.get("data") or item.get("body")
            if isinstance(image, bytes):
                # 处理图像字节输入
                image = Image.open(io.BytesIO(image)).convert("L")
                image = self.transform(image)
            elif isinstance(image, list):
                # 处理像素数组输入
                image = np.array(image).reshape(28, 28)
                image = self.transform(image)
            inputs.append(image)
        
        # 拼接成批量张量并转移到设备
        return torch.stack(inputs).to(self.device)

    def inference(self, data):
        # 执行模型推理
        with torch.no_grad():
            outputs = self.model(data)
        return outputs

    def postprocess(self, inference_output):
        # 将模型输出（logits）转换为可读的预测结果
        probabilities = torch.nn.functional.softmax(inference_output, dim=1)
        predictions = torch.argmax(probabilities, dim=1).tolist()
        return [{"prediction": pred, "confidence": round(prob[pred].item(), 4)} 
                for pred, prob in zip(predictions, probabilities)]

2. 准备必要文件 #

确保你拥有以下文件：

• model.py：包含你的NeuralNetwork类定义
• model.pth：训练完成的模型权重文件
• custom_handler.py：上述自定义处理器文件
• （可选）requirements.txt：列出模型依赖（如torch>=2.0, torchvision>=0.15）

3. 生成.mar文件 #

执行以下torch-model-archiver命令：

torch-model-archiver \
  --model-name mnist_classifier \
  --version 1.0 \
  --model-file model.py \
  --serialized-file model.pth \
  --handler custom_handler.py \
  --export-path model_store \
  --requirements-file requirements.txt

核心参数说明：

• --model-name：指定模型的标识名称
• --model-file：存放模型类定义的文件路径
• --serialized-file：训练好的模型权重文件路径
• --handler：自定义处理器文件路径
• --export-path：生成的.mar文件输出目录

4. 验证与部署 #

将生成的.mar文件放入TorchServe的模型存储目录，启动服务即可测试：

torchserve --start --model-store model_store --models mnist=mnist_classifier.mar

内容的提问来源于stack exchange，提问作者Muazzem Hossain