如何使用TensorFlow Serving部署PyTorch或Scikit-learn模型

Great question—this is a common pain point when you want to leverage TensorFlow Serving's production-ready deployment tools but have models trained in other frameworks. The key here is to convert your model into TensorFlow's SavedModel format (the only format TensorFlow Serving natively supports), and below are practical, framework-specific approaches to do this:

1. PyTorch Models to SavedModel #

The most reliable path uses ONNX as an intermediate format (both PyTorch and TensorFlow have solid ONNX support):

• Step 1: Export your PyTorch model to ONNX
  First, create a dummy input matching your model's expected shape/dtype, then export:

  import torch
  from your_model_module import YourModel
  
  model = YourModel()
  model.load_state_dict(torch.load("pytorch_model_weights.pth"))
  model.eval()
  
  # Dummy input (adjust shape to match your model's requirements)
  dummy_input = torch.randn(1, 3, 224, 224)  # Example for image classification models
  torch.onnx.export(
      model,
      dummy_input,
      "model.onnx",
      opset_version=12,  # Use a compatible opset version for broad support
      input_names=["input"],
      output_names=["output"]
  )
  

• Step 2: Convert ONNX model to TensorFlow SavedModel
  Use the tf2onnx library to convert the ONNX file to a SavedModel:

  python -m tf2onnx.convert --onnx model.onnx --output ./saved_model --opset 12
  

2. Scikit-learn Models to SavedModel #

For traditional ML models from scikit-learn, you have two straightforward options:

• Option A: ONNX conversion pipeline
  First convert your scikit-learn model to ONNX with skl2onnx, then follow the ONNX-to-TF step above:

  from skl2onnx import convert_sklearn
  from skl2onnx.common.data_types import FloatTensorType
  from sklearn.ensemble import RandomForestClassifier
  
  model = RandomForestClassifier()
  # Assume model is already trained on your dataset
  initial_type = [("input", FloatTensorType([None, 10]))]  # Match your input feature count
  onnx_model = convert_sklearn(model, initial_types=initial_type)
  with open("sklearn_model.onnx", "wb") as f:
      f.write(onnx_model.SerializeToString())
  

• Option B: Wrap in a TensorFlow/Keras model
  You can package your scikit-learn model in a custom Keras layer to create a SavedModel directly:

  import tensorflow as tf
  from sklearn.ensemble import RandomForestClassifier
  
  class SklearnModelWrapper(tf.keras.Model):
      def __init__(self, sklearn_model):
          super().__init__()
          self.sklearn_model = sklearn_model
  
      @tf.function(input_signature=[tf.TensorSpec(shape=[None, 10], dtype=tf.float32)])
      def call(self, inputs):
          # Convert TF tensor to numpy, run inference, convert back to tensor
          inputs_np = inputs.numpy()
          outputs_np = self.sklearn_model.predict_proba(inputs_np)
          return tf.convert_to_tensor(outputs_np, dtype=tf.float32)
  
  # Initialize and save the wrapped model
  sklearn_model = RandomForestClassifier()
  # Assume model is trained
  tf_model = SklearnModelWrapper(sklearn_model)
  tf_model.save("./saved_model")
  

3. XGBoost/MXNet Models to SavedModel #

For these frameworks, the ONNX path works consistently:

• XGBoost: Export directly to ONNX using built-in tools:

  import xgboost as xgb
  from onnxmltools.convert import convert_xgboost
  
  model = xgb.Booster(model_file="xgb_model.model")
  onnx_model = convert_xgboost(model, initial_types=[("input", FloatTensorType([None, 10]))])
  with open("xgb_model.onnx", "wb") as f:
      f.write(onnx_model.SerializeToString())
  

• MXNet: Export your MXNet model to ONNX first, then convert to TF SavedModel using tf2onnx as outlined earlier.

Post-Conversion Validation #

After converting, verify your SavedModel is valid and has the correct input/output signatures:

saved_model_cli show --dir ./saved_model --all

Double-check that the input/output names and shapes match what you'll send to TensorFlow Serving.

Deploying with TensorFlow Serving #

Once you have the SavedModel, you don’t need to tweak the model_platform parameter in model.conf—just set it to "tensorflow" (the default) since your model is now in TensorFlow's native format. You can then start TensorFlow Serving exactly as you would with a native TF model.

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