File size: 4,602 Bytes
30f011f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """
bert/trt_session.py
TensorRT ORT session with optimization profiles.
Profiles define min/opt/max tensor shapes so TRT compiles the most efficient
CUDA kernels for the expected (batch_size, sequence_length) distribution.
First load: ~5 min to compile the .plan engine.
Subsequent loads: instant from cache.
"""
from __future__ import annotations
import argparse
import numpy as np
from pathlib import Path
from typing import Dict, List, Tuple
import numpy as np
def build_trt_session(
model_path: str | Path,
cache_dir: str | Path = "./trt_cache",
device_id: int = 0,
# Optimization profile bounds: (min, opt, max) for (batch, seq_len)
min_shape: Tuple[int, int] = (1, 16),
opt_shape: Tuple[int, int] = (32, 128),
max_shape: Tuple[int, int] = (128, 512),
) -> "ort.InferenceSession":
"""
Build an ORT InferenceSession backed by TensorrtExecutionProvider.
The engine is cached to `cache_dir` after first compilation so
subsequent cold-starts load in milliseconds.
"""
try:
import onnxruntime as ort
except ImportError:
raise RuntimeError("pip install onnxruntime-gpu")
Path(cache_dir).mkdir(parents=True, exist_ok=True)
trt_providers = [
(
"TensorrtExecutionProvider",
{
"device_id": device_id,
"trt_fp16_enable": True,
"trt_engine_cache_enable": True,
"trt_engine_cache_path": str(cache_dir),
# Profile: input_ids and attention_mask share the same shape bounds
"trt_profile_min_shapes": f"input_ids:{min_shape[0]}x{min_shape[1]},attention_mask:{min_shape[0]}x{min_shape[1]},token_type_ids:{min_shape[0]}x{min_shape[1]}",
"trt_profile_opt_shapes": f"input_ids:{opt_shape[0]}x{opt_shape[1]},attention_mask:{opt_shape[0]}x{opt_shape[1]},token_type_ids:{opt_shape[0]}x{opt_shape[1]}",
"trt_profile_max_shapes": f"input_ids:{max_shape[0]}x{max_shape[1]},attention_mask:{max_shape[0]}x{max_shape[1]},token_type_ids:{max_shape[0]}x{max_shape[1]}",
"trt_int8_enable": False, # FP16 already halves memory
},
),
"CUDAExecutionProvider", # fallback if TRT fails a subgraph
"CPUExecutionProvider",
]
opts = ort.SessionOptions()
opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
opts.intra_op_num_threads = 4
print(f"[trt_session] loading {model_path} ...")
print(f"[trt_session] cache_dir={cache_dir} (first load may take ~5 min)")
session = ort.InferenceSession(
str(model_path),
sess_options=opts,
providers=trt_providers,
)
actual_providers = session.get_providers()
print(f"[trt_session] active providers: {actual_providers}")
return session
def run_batch(
session: "ort.InferenceSession",
input_ids: np.ndarray, # (B, L) int64
attention_mask: np.ndarray, # (B, L) int64
token_type_ids: np.ndarray | None = None,
) -> np.ndarray:
"""
Run one batch through the TRT session.
Returns logits: (B, 3) float32.
"""
if token_type_ids is None:
token_type_ids = np.zeros_like(input_ids)
feeds = {
"input_ids": input_ids.astype(np.int64),
"attention_mask": attention_mask.astype(np.int64),
"token_type_ids": token_type_ids.astype(np.int64),
}
outputs = session.run(["logits"], feeds)
return outputs[0] # (B, 3)
def softmax(x: np.ndarray) -> np.ndarray:
e = np.exp(x - x.max(axis=-1, keepdims=True))
return e / e.sum(axis=-1, keepdims=True)
def entailment_scores(logits: np.ndarray) -> np.ndarray:
"""Softmax probability of Entailment class (index 2), shape (B,)."""
return softmax(logits)[:, 2]
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, required=True, help="path to FP16 ONNX model")
parser.add_argument("--cache", type=str, default="./trt_cache")
args = parser.parse_args()
session = build_trt_session(args.model, args.cache)
# Smoke test: random batch of 4 sequences length 128
ids = np.random.randint(0, 30000, (4, 128), dtype=np.int64)
mask = np.ones((4, 128), dtype=np.int64)
logits = run_batch(session, ids, mask)
scores = entailment_scores(logits)
print(f"[smoke test] entailment scores: {scores}")
|