fix: gpu memory leaks

This commit is contained in:
Siwat Sirichai 2025-11-10 22:10:46 +07:00
parent 3a47920186
commit 593611cdb7
13 changed files with 420 additions and 166 deletions

View file

@ -24,7 +24,6 @@ from services import (
# Load environment variables
load_dotenv()
def main_single_stream():
"""Single stream example with event-driven architecture."""
print("=" * 80)
@ -33,7 +32,7 @@ def main_single_stream():
# Configuration
GPU_ID = 0
MODEL_PATH = "bangchak/models/frontal_detection_v5.pt" # PT file will be auto-converted
MODEL_PATH = "bangchak/models/frontal_detection_v5.pt" # Transparent loading: .pt, .engine, or .trt
STREAM_URL = os.getenv('CAMERA_URL_1', 'rtsp://localhost:8554/test')
BATCH_SIZE = 4
FORCE_TIMEOUT = 0.05
@ -59,10 +58,10 @@ def main_single_stream():
)
print("✓ Manager created")
# Initialize with PT model (auto-conversion)
print("\n[2/3] Initializing with PT model...")
print("Note: First load will convert PT to TensorRT (3-5 minutes)")
print("Subsequent loads will use cached TensorRT engine\n")
# Initialize with model (transparent loading - no manual parameters needed)
print("\n[2/3] Initializing model...")
print("Note: YOLO models auto-convert to native TensorRT .engine (first time only)")
print("Metadata is auto-detected from model - no manual input_shapes needed!\n")
try:
manager.initialize(
@ -70,11 +69,10 @@ def main_single_stream():
model_id="detector",
preprocess_fn=YOLOv8Utils.preprocess,
postprocess_fn=YOLOv8Utils.postprocess,
num_contexts=4,
pt_input_shapes={"images": (1, 3, 640, 640)},
pt_precision=torch.float16
num_contexts=4
# Note: No pt_input_shapes or pt_precision needed for YOLO models!
)
print("✓ Manager initialized (PT converted to TensorRT)")
print("✓ Manager initialized")
except Exception as e:
print(f"✗ Failed to initialize: {e}")
import traceback
@ -176,6 +174,7 @@ def main_single_stream():
class_counts[obj.class_name] = class_counts.get(obj.class_name, 0) + 1
print(f" Classes: {class_counts}")
except KeyboardInterrupt:
print(f"\n✓ Interrupted by user")
@ -206,7 +205,7 @@ def main_multi_stream():
# Configuration
GPU_ID = 0
MODEL_PATH = "models/yolov8n.pt" # PT file will be auto-converted
MODEL_PATH = "bangchak/models/frontal_detection_v5.pt" # Transparent loading: .pt, .engine, or .trt
BATCH_SIZE = 16
FORCE_TIMEOUT = 0.05
@ -241,17 +240,16 @@ def main_multi_stream():
)
print("✓ Manager created")
# Initialize with PT model
print("\n[2/3] Initializing with PT model...")
# Initialize model (transparent loading)
print("\n[2/3] Initializing model...")
try:
manager.initialize(
model_path=MODEL_PATH,
model_id="detector",
preprocess_fn=YOLOv8Utils.preprocess,
postprocess_fn=YOLOv8Utils.postprocess,
num_contexts=8,
pt_input_shapes={"images": (1, 3, 640, 640)},
pt_precision=torch.float16
num_contexts=8
# Note: No pt_input_shapes or pt_precision needed for YOLO models!
)
print("✓ Manager initialized")
except Exception as e:
@ -312,6 +310,7 @@ def main_multi_stream():
s_fps = stats['count'] / s_elapsed if s_elapsed > 0 else 0
print(f" {sid}: {stats['count']} ({s_fps:.1f} FPS)")
except KeyboardInterrupt:
print(f"\n✓ Interrupted")