使用skorch构建REST服务：将PyTorch模型部署为Web API

使用skorch构建REST服务：将PyTorch模型部署为Web API

skorch 项目地址: https://gitcode.com/gh_mirrors/sko/skorch

本文将介绍如何利用skorch和Palladium框架，将训练好的PyTorch模型快速部署为RESTful Web服务。我们将以IMDB情感分析RNN模型为例，展示从模型训练到服务部署的完整流程。

技术栈概述

skorch是一个基于PyTorch的scikit-learn兼容库，它允许开发者使用scikit-learn风格的API来训练PyTorch神经网络。而Palladium则是一个轻量级的机器学习服务框架，可以轻松地将训练好的模型部署为Web服务。

准备工作

首先需要准备以下两个文件：

1. palladium-config.py - Palladium配置文件
2. model.py - 包含数据加载器和模型定义的Python模块

配置文件详解

Palladium的配置文件采用JSON格式，主要包含以下几个关键部分：

{
    'dataset_loader_train': {
        '__factory__': 'model.DatasetLoader',
        'path': 'aclImdb/train/',
    },
    'dataset_loader_test': {
        '__factory__': 'model.DatasetLoader',
        'path': 'aclImdb/test/',
    },
    'model': {
        '__factory__': 'model.create_pipeline',
        'use_cuda': True,
    },
    'model_persister': {
        '__factory__': 'palladium.persistence.File',
        'path': 'rnn-model-{version}',
    },
    'scoring': 'accuracy',
}

• dataset_loader_train和dataset_loader_test：定义训练和测试数据集的加载方式
• model：指定模型创建函数
• model_persister：配置模型持久化方式
• scoring：指定评估指标

数据加载器实现

在model.py中，我们需要实现数据加载器：

import os
from urllib.request import urlretrieve
import tarfile
import numpy as np
from palladium.interfaces import DatasetLoader as IDatasetLoader
from sklearn.datasets import load_files

DATA_URL = 'http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz'
DATA_FN = DATA_URL.rsplit('/', 1)[1]

def download():
    if not os.path.exists('aclImdb'):
        if not os.path.exists(DATA_FN):
            urlretrieve(DATA_URL, DATA_FN)
        with tarfile.open(DATA_FN, 'r:gz') as f:
            f.extractall()

class DatasetLoader(IDatasetLoader):
    def __init__(self, path='aclImdb/train/'):
        self.path = path

    def __call__(self):
        download()
        dataset = load_files(self.path, categories=['pos', 'neg'])
        X, y = dataset['data'], dataset['target']
        X = np.asarray([x.decode() for x in X])  # 解码字节数据
        return X, y

这个数据加载器会自动下载IMDB数据集并解压，然后加载指定路径下的数据。

模型管道构建

接下来，我们需要创建模型管道：

from dstoolbox.transformers import Padder2d
from dstoolbox.transformers import TextFeaturizer
from sklearn.pipeline import Pipeline
from skorch import NeuralNetClassifier
import torch

def create_pipeline(
    vocab_size=1000,
    max_len=50,
    use_cuda=False,
    **kwargs
):
    return Pipeline([
        ('to_idx', TextFeaturizer(max_features=vocab_size)),
        ('pad', Padder2d(max_len=max_len, pad_value=vocab_size, dtype=int)),
        ('net', NeuralNetClassifier(
            RNNClassifier,
            device=('cuda' if use_cuda else 'cpu'),
            max_epochs=5,
            lr=0.01,
            optimizer=torch.optim.RMSprop,
            module__vocab_size=vocab_size,
            **kwargs,
        ))
    ])

这个管道包含三个步骤：

1. 文本特征化（将文本转换为索引）
2. 填充（使所有序列长度一致）
3. 神经网络分类器（基于RNN的情感分类器）

RNN模型定义

下面是RNN分类器的实现：

from torch import nn
F = nn.functional

class RNNClassifier(nn.Module):
    def __init__(
        self,
        embedding_dim=128,
        rec_layer_type='lstm',
        num_units=128,
        num_layers=2,
        dropout=0,
        vocab_size=1000,
    ):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.rec_layer_type = rec_layer_type.lower()
        self.num_units = num_units
        self.num_layers = num_layers
        self.dropout = dropout

        self.emb = nn.Embedding(
            vocab_size + 1, embedding_dim=self.embedding_dim)

        rec_layer = {'lstm': nn.LSTM, 'gru': nn.GRU}[self.rec_layer_type]
        self.rec = rec_layer(
            self.embedding_dim, self.num_units,
            num_layers=num_layers, batch_first=True,
            )

        self.output = nn.Linear(self.num_units, 2)

    def forward(self, X):
        embeddings = self.emb(X)
        if self.rec_layer_type == 'gru':
            _, rec_out = self.rec(embeddings)
        else:
            _, (rec_out, _) = self.rec(embeddings)
        rec_out = rec_out[-1]  # 取最后一层RNN的输出
        drop = F.dropout(rec_out, p=self.dropout)
        out = F.softmax(self.output(drop), dim=-1)
        return out

这个模型包含嵌入层、RNN层（LSTM或GRU）和输出层，最后使用softmax激活函数输出概率分布。

模型训练与评估

安装依赖后，可以运行以下命令训练模型：

PALLADIUM_CONFIG=palladium-config.py pld-fit --evaluate

训练过程会输出训练损失、验证准确率等指标。训练完成后，模型会自动保存。

部署Web服务

在配置文件中添加预测服务配置：

{
    'predict_service': {
        '__factory__': 'palladium.server.PredictService',
        'mapping': [
            ('text', 'str'),
        ],
        'predict_proba': True,
        'unwrap_sample': True,
    },
}

然后启动开发服务器：

PALLADIUM_CONFIG=palladium-config.py pld-devserver

测试API服务

服务启动后，可以通过HTTP请求测试API：

http://localhost:5000/predict?text=this+movie+was+brilliant

返回结果示例：

{
    "metadata": {"error_code": 0, "status": "OK"},
    "result": [0.326442807912827, 0.673557221889496],
}

其中result数组包含两个概率值，分别对应负面和正面的情感概率。

总结

通过skorch和Palladium的组合，我们可以轻松地将PyTorch模型部署为Web服务。这种方案具有以下优势：

1. 代码与配置分离，便于管理和实验
2. 自动处理模型训练、评估和持久化
3. 提供标准化的REST API接口
4. 与scikit-learn生态系统兼容

对于需要快速部署机器学习模型的生产环境，这种方案提供了高效可靠的解决方案。

skorch 项目地址: https://gitcode.com/gh_mirrors/sko/skorch