TorchServe 基本用法

1. Prepare the Model Script (fnn.py)

You'll need to make a few modifications to your fnn.py script to make it compatible with TorchServe. Specifically, you'll need to:

Define a ModelHandler class: This class will handle loading the model, preprocessing input, performing inference, and postprocessing output.
Implement the required methods: The ModelHandler class needs to implement specific methods that TorchServe will call.

Here's the modified fnn.py (I'll call it fnn_handler.py to avoid confusion):

import torch
import torch.nn as nn
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import json  # For handling JSON data

class FishNet(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(FishNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        return x

class ModelHandler(object):
    """
    A custom model handler class.
    """

    def __init__(self):
        self.model = None
        self.scaler = None  # Scaler instance
        self.mapping = None
        self.device = None
        self.initialized = False

    def initialize(self, context):
        """
        Initializes the model, scaler, and other artifacts.
        """
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.manifest = context.manifest

        properties = context.system_properties
        model_dir = properties.get("model_dir")
        self.device = torch.device("cuda:" + str(properties.get("gpu_id")) if torch.cuda.is_available() else "cpu")
        # Read model serialize/pt file
        self.model = self._load_pickled_model(model_dir)
        self.model.to(self.device)
        self.model.eval()

        # Load scaler
        self.scaler = self._load_scaler(model_dir)

        self.initialized = True

    def preprocess(self, data):
        """
        Preprocesses the input data.
        """
        # Assuming the input is a JSON string
        body = json.loads(data[0]['body'].decode('utf-8'))
        input_data = body['data']  # Assuming the JSON has a 'data' field

        # Convert to DataFrame and scale
        input_df = pd.DataFrame([input_data], columns=['Length1', 'Length2', 'Length3', 'Height', 'Width'])
        scaled_data = self.scaler.transform(input_df.values)
        
        # Convert to tensor
        tensor_data = torch.tensor(scaled_data, dtype=torch.float32).to(self.device)
        return tensor_data

    def inference(self, data):
        """
        Performs inference on the input data.
        """
        with torch.no_grad():
            output = self.model(data)
        return output.cpu().tolist()

    def postprocess(self, data):
        """
        Postprocesses the output data.
        """
        return data

    def _load_pickled_model(self, model_dir):
        """
        Loads the pickled model.
        """
        model_pt_path = "{}/fnn.pt".format(model_dir)
        model = FishNet(5, 64, 1)  # Assuming input_size=5, hidden_size=64, output_size=1
        model.load_state_dict(torch.load(model_pt_path))
        return model

    def _load_scaler(self, model_dir):
        """
        Loads the scaler.
        """
        scaler_path = "{}/scaler.pkl".format(model_dir)
        import pickle
        with open(scaler_path, 'rb') as f:
            scaler = pickle.load(f)
        return scaler

2. Save the Model and Scaler

You need to save your trained model and the scaler in a format that TorchServe can load. Modify your original fnn.py (or a separate script) to include the following:

import torch
import torch.nn as nn
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import pickle

# Define The Feedforward Neural Network
class FishNet(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(FishNet, self).__init__()

        # First hidden layer
        self.fc1 = nn.Linear(input_size, hidden_size)

        # Second hidden layer
        self.fc2 = nn.Linear(hidden_size, output_size)
        
        # Output layer
        self.output = nn.Linear(hidden_size, output_size)
        
        # Activation function
        self.relu = nn.ReLU()

    def forward(self, x):
        # First hidden layer with ReLU activation
        x = self.relu(self.fc1(x))

        # Second hidden layer with ReLU activation
        x = self.relu(self.fc2(x))

        # Output layer
        x = self.output(x)

        return x

# Load your data using pandas
df = pd.read_csv('data/Fish_Dataset_Pytorch.csv')

# Preprocess the data
X = df[['Length1', 'Length2', 'Length3', 'Height', 'Width']].values
y = df['Weight'].values

# Scale the features
scaler = StandardScaler()
X = scaler.fit_transform(X)

# Split the data into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Convert the data to PyTorch tensors
X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
y_train_tensor = torch.tensor(y_train, dtype=torch.float32)
X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
y_test_tensor = torch.tensor(y_test, dtype=torch.float32)

# Define The Feedforward Neural Network
class FishNet(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(FishNet, self).__init__()

        # First hidden layer
        self.fc1 = nn.Linear(input_size, hidden_size)

        # Second hidden layer
        self.fc2 = nn.Linear(hidden_size, output_size)
        
        # Output layer
        self.output = nn.Linear(hidden_size, output_size)
        
        # Activation function
        self.relu = nn.ReLU()

    def forward(self, x):
        # First hidden layer with ReLU activation
        x = self.relu(self.fc1(x))

        # Second hidden layer with ReLU activation
        x = self.relu(self.fc2(x))

        # Output layer
        x = self.output(x)

        return x
    
input_size = X_train_tensor.shape[1]  # Number of features
hidden_size = 64  # Number of neurons in the hidden layer
output_size = 1  # Number of output classes (binary classification)

model = FishNet(input_size, hidden_size, output_size)
criterion = nn.MSELoss()  # Mean Squared Error Loss for regression
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)  # Adam optimizer with learning rate of 0.001

# Trainning Loop
num_epochs = 100 
for epoch in range(num_epochs):
    model.train()  # Set the model to training mode
    optimizer.zero_grad()  # Zero the gradients

    # Forward pass
    outputs = model(X_train_tensor)
    
    # Compute loss
    loss = criterion(outputs, y_train_tensor.unsqueeze(1)) # Reshape y to

    # Backward pass and optimization
    loss.backward()
    optimizer.step()

# After training, save the model and scaler
torch.save(model.state_dict(), "model/fnn.pt") # save the model
with open('model/scaler.pkl', 'wb') as f: # save the scaler
    pickle.dump(scaler, f)

print("Model and scaler saved!")

Create the model directory:

mkdir model

Run the script to save the model and scaler:

python save_model.py

3. Create a model-config.yaml file

This file tells TorchServe how to load and run your model.

version: 1.0
model:
  name: fnn
  handler: fnn_handler.ModelHandler
  marName: fnn.mar
  modelType: pytorch
  description: "Fish Weight Prediction Model"
  batching:
    maxBatchDelay: 100
    maxBatchSize: 8
  workers:
    - threads: 1
      memory: 1024
      gpu: false

4. Create a Model Archive (.mar file)

Use the torch-model-archiver command to create a model archive file. This packages your model, handler, and any other necessary files into a single file that TorchServe can deploy.

torch-model-archiver \
    --model-name fnn \
    --version 1.0 \
    --handler fnn_handler.py \
    --model-file save_model.py \
    --export-path . \
    --archive-format default

This will create a file named fnn.mar in your current directory.

5. Start TorchServe

First, make sure you have TorchServe installed. If not, install it:

pip install torchserve torch-model-archiver

Then, start TorchServe:

torchserve --start --model-store . --models fnn.mar

6. Register the Model

curl -v -X POST "http://localhost:8081/models?model_name=fnn&url=fnn.mar&handler=fnn_handler.ModelHandler&batch_size=1&max_batch_delay=100"

7. Send Inference Requests

Now you can send inference requests to your deployed model using curl. The input data should be a JSON object containing the feature values.

curl -v -X POST http://localhost:8080/predictions/fnn -H "Content-Type: application/json" -d '{"data": [14.3, 15.0, 16.2, 2.4, 4.2]}'

Explanation:

http://localhost:8080/predictions/fnn: This is the inference endpoint for your model. fnn is the name you gave your model when you registered it.
-H "Content-Type: application/json": This tells the server that you're sending JSON data.
-d '{"data": [14.3, 15.0, 16.2, 2.4, 4.2]}': This is the JSON data containing the feature values for your model. Make sure the order of the features matches the order used during training.

8. Shutdown TorchServe

When you're finished, stop TorchServe:

torchserve --stop

Important Notes:

Dependencies: Make sure all the necessary libraries (e.g., pandas, scikit-learn) are installed in your TorchServe environment.
Error Handling: Add more robust error handling to your ModelHandler class to catch potential issues during preprocessing, inference, and postprocessing.
Input Validation: Validate the input data in the preprocess method to ensure it has the correct format and range of values.
GPU: If you want to use a GPU, make sure you have the correct CUDA drivers and PyTorch version installed. Also, set the gpu property to true in your model-config.yaml file.
Model Directory: Ensure that the paths to your model file (fnn.pt) and scaler file (scaler.pkl) are correct in the _load_pickled_model and _load_scaler methods of your ModelHandler class.
Versions: Use specific versions of libraries to avoid compatibility issues.
Logging: Add logging to your handler to help debug issues.

This comprehensive guide should get you started with deploying your FNN model using TorchServe. Remember to adapt the code and commands to your specific environment and model configuration.

THE END