REST API inference in PyTorch - Model Pipeline Trace

Model Pipeline - REST API inference

This pipeline shows how a trained PyTorch model receives data through a REST API, processes it, and returns predictions. It covers data input, preprocessing, model inference, and output formatting.

Data Flow - 5 Stages

1API receives JSON input

1 request with JSON payload→Receive raw JSON data from client→1 JSON object with feature values

{"feature1": 5.1, "feature2": 3.5, "feature3": 1.4, "feature4": 0.2}

↓

2Preprocessing

1 JSON object with 4 features→Convert JSON to PyTorch tensor and normalize features→1 tensor of shape (1, 4)

[[0.5, 0.7, 0.2, 0.1]]

↓

3Model inference

1 tensor of shape (1, 4)→Feed tensor into trained PyTorch model to get raw output→1 tensor of shape (1, 3)

[[2.1, 0.3, -1.2]]

↓

4Postprocessing

1 tensor of shape (1, 3)→Apply softmax to get class probabilities→1 tensor of shape (1, 3) with probabilities

[[0.84, 0.10, 0.06]]

↓

5API sends JSON response

1 tensor of shape (1, 3) with probabilities→Convert probabilities to JSON format and send back→1 JSON object with class probabilities

{"class_0": 0.84, "class_1": 0.10, "class_2": 0.06}

Training Trace - Epoch by Epoch

Loss
1.2 |****
1.0 |***
0.8 |**
0.6 |**
0.4 |*
0.2 |
    +----------------
     1 2 3 4 5 Epochs

Epoch	Loss ↓	Accuracy ↑	Observation
1	1.20	0.45	Model starts learning with high loss and low accuracy
2	0.85	0.65	Loss decreases and accuracy improves
3	0.60	0.78	Model continues to improve
4	0.45	0.85	Good convergence with lower loss and higher accuracy
5	0.35	0.90	Training stabilizes with strong performance

Prediction Trace - 5 Layers

Layer 1: Receive JSON input

Layer 2: Preprocessing

Layer 3: Model forward pass

Layer 4: Softmax activation

Layer 5: Format JSON response

Model Quiz - 3 Questions

What shape does the data have after preprocessing before model input?

A 1 tensor of shape (1, 4)

B 1 JSON object with 4 features

C 1 tensor of shape (4, 1)

D 1 tensor of shape (1, 3)

Key Insight

This visualization shows how a trained PyTorch model can be wrapped in a REST API to receive input data, preprocess it, run inference, and return predictions as probabilities. The training trace confirms the model learns well, and the prediction trace explains how raw input transforms step-by-step to output probabilities.