How We Cut ERNIE Model Resource Use by 75% with Pruning, Structured Slimming, and ONNX Runtime

Abstract

With AI models booming, cost‑efficiency has become a core challenge for AI engineers. This article fully recounts a production‑level optimization of an ERNIE text‑classification model that originally occupied 128 CPU cores and 96 GB memory. By combining model selection, structured pruning, and inference‑engine innovation, we reduced resource consumption by 75% while preserving business‑critical accuracy.

Main Content

1. Introduction: When ERNIE Becomes a “Gold‑Eating Beast”

Our real‑time text anomaly detection service relied on ERNIE‑3.0 Base, which, despite its strong performance, incurred massive CPU and memory costs, leading to three major pain points: high cost, low deployment efficiency, and sluggish elasticity.

Key Points

Problem Diagnosis – Precisely locate the heavyweight NLP model bottleneck.

Three‑Step Optimization – Model slimming, structured pruning, and engine acceleration.

Pruning Decisions – Use data‑driven trade‑offs to keep recall ≥ 99.5%.

Engineering Practice – Rebuild the pipeline from Paddle Taskflow to ONNX Runtime, exposing core code.

Future Outlook – From FP32 to INT8 quantization and TensorRT.

Optimization Steps

Step One: Model Selection

Switch from ERNIE‑Base to the lighter ernie‑3.0‑medium‑zh (6 layers, 12 heads) – a 50% reduction in layers while maintaining near‑baseline recall (99.7%).

Step Two: Structured Pruning

Using PaddleSlim , we performed surgical pruning on attention heads. Multiple compression ratios were generated in a single run:

python train.py \
  --do_compress \
  --device gpu \
  --data_dir data \
  --model_name_or_path checkpoint \
  --output_dir checkpoint/prune \
  --learning_rate 3e-5 \
  --per_device_train_batch_size 32 \
  --per_device_eval_batch_size 32 \
  --num_train_epochs 1 \
  --max_length 128 \
  --logging_steps 5 \
  --save_steps 100 \
  --width_mult_list '3/4''2/3''1/2'  # produce three compression rates

Data‑driven evaluation kept recall above 99.5% while reducing attention heads from 12 to 9, achieving a 21% speed boost.

Step Three: Engine Innovation

We replaced the black‑box Paddle Taskflow inference with a transparent ONNX Runtime (ORT) pipeline, manually handling tokenization, inference, and post‑processing, and switched the service framework to FastAPI for better performance.

import onnxruntime as ort
from paddlenlp.transformers import ErnieTokenizer
from scipy.special import softmax
import numpy as np
from fastapi import FastAPI
import uvicorn

# Initialize components
tokenizer = ErnieTokenizer.from_pretrained('ernie-3.0-tiny-medium-v2-zh')
session = ort.InferenceSession('/opt/prune_75_model.onnx')
labels = ['异常标签1','异常标签2','异常标签3','异常标签4','异常标签5','正常','异常标签6']
app = FastAPI()

@app.get('/check')
def do_check(content: str):
    inputs = tokenizer(content, return_tensors='np', max_length=128, padding=True, truncation=True)
    logits = session.run(None, {
        'input_ids': inputs['input_ids'],
        'token_type_ids': inputs['token_type_ids']
    })[0]
    probabilities = softmax(logits, axis=-1)
    pred_idx = np.argmax(probabilities, axis=-1)[0]
    return {
        'content': content,
        'result': labels[pred_idx],
        'probability': float(probabilities[0][pred_idx])
    }

Results

After optimization, the service runs on a 2‑core CPU with 64 GB memory, a 75% reduction in total resource usage, while maintaining recall of 99.6% (baseline 99.7%). Inference latency dropped from 28 ms to 22 ms, and overall CPU usage fell from 128 cores to 2 cores, freeing 96 cores for other workloads.

Key takeaways:

Reject “technology inertia” – measure cost and performance, not just feasibility.

Leverage ecosystem tools (model compression, specialized runtimes) as productivity multipliers.

Let data drive trade‑offs between accuracy and efficiency.

Conclusion & Outlook

The journey from a heavyweight “black‑box” to a lean “white‑box” demonstrates that engineering excellence is as decisive as algorithmic innovation. Future work will explore INT8 quantization and TensorRT deployment for even greater speed gains.